Set numpy array values to zero where x,y,z are not all multiples of n-CodePudding

Here is a naive implementation of a function which sets every value of a numpy array with index x,y,z to zero if any of x,y, or z is not a multiple of one_per_n:

import numpy as np

X = np.random.rand(10,10,10)

one_per_n = 4

for x in range(X.shape[0]):
    for y in range(X.shape[1]):
        for z in range(X.shape[2]):
            if x % one_per_n != 0 or y % one_per_n != 0 or z % one_per_n != 0:
                X[x,y,z] = 0

I am looking for a more efficient method.

CodePudding user response：

In each dimension, the elements that retain their values are the ones where the modulo is zero. They are spaced one_per_n apart, so utilizing a slice with one_per_n as step size is useful.

You can construct a mask of elements to set to 0 for each dimension, and then zero out X at all these positions.

The inverse of

# elements to set to zero
(x % one_per_n != 0) or (y % one_per_n != 0) or (z % one_per_n != 0)

# elements to keep
(x % one_per_n == 0) and (y % one_per_n == 0) and (z % one_per_n == 0)

As the inverse selection masks are connected by logical ands, they can be sliced at the same time:

# init with zeroing out everything
zero_out_mask = np.ones_like(X, dtype=bool)

# except for the modulo == 0 elements
zero_out_mask[::one_per_n, ::one_per_n, ::one_per_n] = False

X[zero_out_mask] = 0

CodePudding user response：

A quite literal numpy vectorization of the code would be:

x, y, z = np.indices(X.shape, sparse=True)
X[(x % one_per_n != 0) | (y % one_per_n != 0) | (z % one_per_n != 0)] = 0

This is still fairly efficient because modulus is only computed for 1D arrays thanks to sparse=True.