CODEWITHSUNDEEP
pythonnumpyvectorization
ToughMind
ToughMind

Reputation: 1009

How to vectorize this for loops in python?

Code is as below

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)

I have a numpy array "data" with shape (h, w). Its elements is int number ∈[0, 10).
I create a array dataMask with shape (h, w, 10). dataMask[i, j, k] indicats the number of points whose value is k within an area in data. This area in data has center (i,j) and r = 2, and is a square.

How to vectorize those for loops in the code? Thank you!

Upvotes: 2

Views: 170

Answers (2)

javidcf
javidcf

Reputation: 59731

This is a "partially vectorized" solution that just iterates over the window size.

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

You can actually make this fully vectorized by just tiling the data more (which uses more memory):

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

Upvotes: 1

Paul Panzer
Paul Panzer

Reputation: 53119

Here is one method using 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))

Example run with h,w,r = 30,40,20

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

Upvotes: 2

Related Questions