Update 2D NumPy array indexed with cross product with repeated indices

user_8260213 0
Viewed 99

I'm interested in the version of Increment Numpy multi-d array with repeated indices indexed with a cross-product.

In particular, I want to perform the operation done by the following code using matrix operations to accelerate it:

def get_s(image, grid_size):
    W, H = image.shape
    s = np.zeros((W, H))
    for w in range(W):
        for h in range(H):
            i, j = int(w / grid_size), int(h / grid_size)
            s[i, j] += image[w, h]
    return s

My idea was to compute all the (i, j) indices at once and use NumPy's ix_ method to index the matrix s:

def get_s(image, grid_size):
    W, H = image.shape
    s = np.zeros((W, H))
    w_idx, h_idx = np.arange(W), np.arange(H)
    x_idx, y_idx = np.trunc(w_idx / grid_size).astype(int), np.trunc(h_idx / grid_size).astype(int)
    s[np.ix_(x_idx, y_idx)] += image
    return s

It is easier to understand the code above with NumPy's example:

Using ix_ one can quickly construct index arrays that will index the cross product. a[np.ix_([1,3],[2,5])] returns the array [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

In my case, it's likely that some indices will be repeated (as for example with grid_size=2, int(0 / grid_size) = int(1 / grid_size)). And that's where the Increment Numpy multi-d array with repeated indices question comes.

In case the indices are repeated, I would like to update the matrix with the image value by the same number of times. I cannot get any solution to this problem without any additional loops (e.g., zipping the indices; but you essentially have to perform the actual cross product of the indices for s and the image).

numpy arrays python
2 Answers

Does skimage.measure.block_reduce do what you want?

from skimage.measure import block_reduce
s = block_reduce(image, block_size=(grid_size, grid_size), func=np.sum)
user_5656350

I don't think this is the best way to do it but here's one way.

import numpy as np

image = np.arange(9).reshape(3, 3)
s = np.zeros((5, 5))
x_idx, y_idx = np.meshgrid([0, 0, 2], [1, 1, 2])
# find unique destinations
idxs = np.stack((x_idx.flatten(), y_idx.flatten())).T
idxs_unique, counts = np.unique(idxs, axis = 0, return_counts = True)
# create mask for the source and sumthe source pixels headed to the same destination
idxs_repeated = idxs[None, :, :].repeat(len(idxs_unique), axis = 0)
image_mask = (idxs_repeated == idxs_unique[:, None, :]).all(-1)
pixel_sum = (image.flatten()[None, :]*image_mask).sum(-1)
# assign summed sources to destination
s[tuple(idxs_unique.T)] += pixel_sum

EDIT 1:

If you run into problems caused by memory constraints you can do the image masking and summation in batches as done in the following implementation. I set the batch size to 10 but that parameter can be set to whatever works on your machine.

import numpy as np

image = np.arange(12).reshape(3, 4)
s = np.zeros((5, 5))
x_idx, y_idx = np.meshgrid([0, 0, 2], [1, 1, 2, 1])
idxs = np.stack((x_idx.flatten(), y_idx.flatten())).T
idxs_unique, counts = np.unique(idxs, axis = 0, return_counts = True)
batch_size = 10
pixel_sum = []
for i in range(len(unique_idxs)//batch_size + ((len(unique_idxs)%batch_size)!=0)):
    batch = idxs_unique[i*batch_size:(i+1)*batch_size, None, :]
    idxs_repeated = idxs[None, :, :].repeat(len(batch), axis = 0)
    image_mask = (idxs_repeated == idxs_unique[i*batch_size:(i+1)*batch_size, None, :]).all(-1)
    pixel_sum.append((image.flatten()[None, :]*image_mask).sum(-1))
pixel_sum = np.concatenate(pixel_sum)
s[tuple(idxs_unique.T)] += pixel_sum

EDIT 2: OP's method seems to be faster by far if you use numba.

import numpy as np
from numba import jit

@jit(nopython=True)
def get_s(image, grid_size):
    W, H = image.shape
    s = np.zeros((W, H))
    for w in range(W):
        for h in range(H):
            i, j = int(w / grid_size), int(h / grid_size)
            s[i, j] += image[w, h]
    return s

def get_s_vec(image, grid_size, batch_size = 10):
    W, H = image.shape
    s = np.zeros((W, H))
    w_idx, h_idx = np.arange(W), np.arange(H)
    x_idx, y_idx = np.trunc(w_idx / grid_size).astype(int), np.trunc(h_idx / grid_size).astype(int)
    y_idx, x_idx = np.meshgrid(y_idx, x_idx)
    idxs = np.stack((x_idx.flatten(), y_idx.flatten())).T
    idxs_unique, counts = np.unique(idxs, axis = 0, return_counts = True)
    pixel_sum = []
    for i in range(len(unique_idxs)//batch_size + ((len(unique_idxs)%batch_size)!=0)):
        batch = idxs_unique[i*batch_size:(i+1)*batch_size, None, :]
        idxs_repeated = idxs[None, :, :].repeat(len(batch), axis = 0)
        image_mask = (idxs_repeated == idxs_unique[i*batch_size:(i+1)*batch_size, None, :]).all(-1)
        pixel_sum.append((image.flatten()[None, :]*image_mask).sum(-1))
    pixel_sum = np.concatenate(pixel_sum)
    s[tuple(idxs_unique.T)] += pixel_sum
    return s


print(f'loop result = {get_s(image, 2)}')
print(f'vector result = {get_s_vec(image, 2)}')

%timeit get_s(image, 2)
%timeit get_s_vec(image, 2)

output:

loop result = [[10. 18.  0.  0.]
 [17. 21.  0.  0.]
 [ 0.  0.  0.  0.]]
vector result = [[10. 18.  0.  0.]
 [17. 21.  0.  0.]
 [ 0.  0.  0.  0.]]
The slowest run took 15.00 times longer than the fastest. This could mean that an intermediate result is being cached.
1000000 loops, best of 5: 751 ns per loop
1000 loops, best of 5: 195 µs per loop
user_153066060
Related