Как векторизовать это для циклов в Python?

Question

код, как показано ниже

import numpy as np


data = np.random.randint(0, 10, 12).reshape(3, 4)
print(data)

h, w = data.shape[:2]
dataMask = np.zeros((h, w, 10), np.int)

r = 2

for i in range(h):
    for j in range(w):
        for ir in range(i - r, i + r):
            for jr in range(j - r, j + r):
                if ir >= 0 and ir < h and jr >= 0 and jr < w:
                    dataMask[i, j, data[ir, jr]] += 1

print(dataMask)

У меня есть массив данных с формой (h, w). Его элементами является целое число ∈ [0, 10).
Я создаю массив dataMask с формой (h, w, 10). dataMask [i, j, k] указывает количество точек, значение которых равно k в области данных. Эта область в данных имеет центр (i, j) и r = 2 и является квадратом.

Как векторизовать эти циклы в коде? Спасибо!

Answer 1

Вот один метод, использующий cumsum:

import numpy as np


data = np.random.randint(0, 10, 1200).reshape(30, 40)
print(data)

h, w = data.shape[:2]
dataMask = np.zeros((h, w, 10), np.int)

r = 20

from time import time
T = []

T.append(time())

for i in range(h):
    for j in range(w):
        for ir in range(i - r, i + r):
            for jr in range(j - r, j + r):
                if ir >= 0 and ir < h and jr >= 0 and jr < w:
                    dataMask[i, j, data[ir, jr]] += 1

T.append(time())

m1 = np.zeros((h, w, 10), np.int)
np.put_along_axis(m1, data[...,None], 1, 2)
m2 = np.empty_like(m1)
m1 = m1.cumsum(1)
m2[: ,:-r+1] = m1[:, r-1:]
m2[:, -r+1:] = m1[:, -1, None]
m2[:, r+1:] -= m1[:, :-r-1]
m2 = m2.cumsum(0)
m1[:-r+1] = m2[r-1:]
m1[-r+1:] = m2[-1, None]
m1[r+1:] -= m2[:-r-1]

T.append(time())


assert (dataMask==m1).all()

print(np.diff(T))

Пример выполнения с h,w,r = 30,40,20

# time [seconds] used by
# OP            cumsum
[9.23162699e-01 3.41892242e-04]

Answer 2

Это «частично векторизованное» решение, которое просто перебирает размер окна.

import numpy as np
from itertools import product

# Input data
np.random.seed(0)
data = np.random.randint(0, 10, 12).reshape(3, 4)
h, w = data.shape[:2]
dataMask = np.zeros((h, w, 10), np.int)
r = 2

# Original solution
for i in range(h):
    for j in range(w):
        for ir in range(i - r, i + r):
            for jr in range(j - r, j + r):
                if ir >= 0 and ir < h and jr >= 0 and jr < w:
                    dataMask[i, j, data[ir, jr]] += 1

# Partially vectorized solution
idx_i, idx_j = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
idx_i = idx_i.ravel()
idx_j = idx_j.ravel()
idx_k = data.ravel()
dataMask2 = np.zeros((h, w, 10), np.int)
for i, j in product(range(-r + 1, r + 1), repeat=2):
    ii = idx_i + i
    jj = idx_j + j
    m = (ii >= 0) & (ii < h) & (jj >= 0) & (jj < w)
    ii = ii[m]
    jj = jj[m]
    kk = idx_k[m]
    np.add.at(dataMask2, (ii, jj, kk), 1)

print(np.all(dataMask == dataMask2))
# True

На самом деле вы можете сделать это полностью векторизованным, просто разбив данные на листы (что использует больше памяти):

import numpy as np

# Fully vectorized
idx_i, idx_j = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
w_i, w_j = np.meshgrid(np.arange(-r + 1, r + 1), np.arange(-r + 1, r + 1), indexing='ij')
ii = (idx_i[:, :, np.newaxis, np.newaxis] + w_i).ravel()
jj = (idx_j[:, :, np.newaxis, np.newaxis] + w_j).ravel()
kk = np.tile(data[:, :, np.newaxis, np.newaxis], (1, 1, 2 * r, 2 * r)).ravel()
m = (ii >= 0) & (ii < h) & (jj >= 0) & (jj < w)
ii = ii[m]
jj = jj[m]
kk = kk[m]
dataMask3 = np.zeros((h, w, 10), np.int)
np.add.at(dataMask3, (ii, jj, kk), 1)
print(np.all(dataMask == dataMask3))
# True