Using numpy where on multidimensional array but need to control indexing

Question

I need to modify elements of an 3D array if they exceed some threshold value. The modification is based upon related elements of another array. More concretely:

A_ijk = A_ijk if A_ijk < threshold value

= (B_(i-1)jk + B_ijk) / 2, otherwise

Numpy.where provides most of the functionality I need, but I don't know how to iterate over the first index without an explicit loop. The follow code does what I want, but uses a loop. Is there a better way? Assume A and B are same shape.

for i in xrange(A.shape[0]):
    A[i] = numpy.where(A[i] <= threshold, A[i], (B[i - 1] + B[i]) / 2)

To address the comments below: The first few rows of A are guaranteed to be below threshold. This keeps the i index from looping over to the last entry of A.

Does this wrap around? So for example if i == 0, then B[i-1] should reference the last row of the array? — user3483203
– user3483203, Commented Oct 16, 2018 at 21:59
A and B can't be of the same shape if A[i,...] references both B[i,...] and B[i-1,...]. I mean it can, but are you sure the loopy code does what you want it to do? — Andras Deak -- Слава Україні
– Andras Deak -- Слава Україні, Commented Oct 16, 2018 at 21:59

Andras Deak -- Слава Україні · Accepted Answer · 2018-10-17 08:57:16Z

1

You can vectorize your operation by using boolean indexing to replace the elements of A that are above the threshold. A little care has to be taken, since the auxiliary array corresponding to (B[i-1] + B[i])/2 has one less size along the first dimension than A, so we have to explicitly ignore the first row of A (knowing that they are all below the threshold, as explained in the question):

import numpy as np

# some dummy data
A = np.random.rand(3,4,5)
B = np.random.rand(3,4,5)
threshold = 0.5
A[0,:] *= threshold # put the first dummy row below threshhold
mask = A[1:] > threshold # to be overwritten, shape (2,4,5)

replace = (B[:-1] + B[1:])/2 # to overwrite elements in A from, shape (2,4,5)

# replace corresponding elements where `mask` is True
A[1:][mask] = replace[mask]

edited Oct 17, 2018 at 8:57

answered Oct 16, 2018 at 22:07

Andras Deak -- Слава Україні

35.4k13 gold badges94 silver badges118 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Jim Parker Over a year ago

Thanks! This will work. As mentioned above, the first row of A is known to be below threshold, so the mask values from the B matrix never "loop around"

Andras Deak -- Слава Україні Over a year ago

@JimParker thanks, I missed that detail (I had already begun writing up the answer by then). I updated the answer, it's slightly simpler with this additional assumption (no auxiliary array is needed).

cvanelteren · Accepted Answer · 2018-10-16 22:09:43Z

0

You can use where to directly index into ndarray:

a   = np.random.rand(4,3,10)
b   = np.zeros(a.shape)
idx = np.where(a < .1)
print(a)
a[idx] = b[idx]
print(a)

If a for-loop is needed however, just convert the ravel the indices and update.

a   = np.random.rand(4,3,10)
b   = np.zeros(a.shape)
idx = [np.ravel_multi_index(i, a.shape) for i in zip(*np.where(a < .1))]
print(a, idx)
for i in idx:
    a.ravel()[i] = b.ravel()[i]
print(a)

answered Oct 16, 2018 at 22:09

cvanelteren

1,70311 silver badges18 bronze badges

Collectives™ on Stack Overflow

Using numpy where on multidimensional array but need to control indexing

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related