CODEWITHSUNDEEP
pythonarraysnumpy
iron2man
iron2man

Reputation: 1837

A more efficient method of passing multiple conditions for each element in a 3D numpy array

I have a 3 dimensional numpy array that I am checking multiple conditions for. I am checking each element to see if they are less than a certain number. If each 3d element is indexed by i, where i=[0,1,2] in a what I call array3, and if one of the elements is greater than the number I have set, maybe giving a boolean array [False, True, True] or [False, False, True], then this index is eliminated from array3.

I have a dumb method for each element less than 20:

import numpy as np

wx = np.where( np.abs(array3[:,0]) <= 20.0 ) # x values less than 20
xarray3x = array3[:,0][wx]
yarray3x = array3[:,1][wx]
zarray3x = array3[:,2][wx]

wy = np.where( np.abs(yarray3x) <= 20.0 ) # y values less than 20
xarray3xy = xarray3x[wy]
yarray3xy = yarray3x[wy]
zarray3xy = zarray3x[wy]

wz = np.where( np.abs(zarray3xy) <= 20.0 ) # z values less than 20
xarray3xyz = xarray3xy[wz]
yarray3xyz = yarray3xy[wz]
zarray3xyz = zarray3xy[wz]

Which works, but it can be annoying to keep up with the variables I have named. So now I am trying to write something that takes up less lines (and hopefully less compiling time).

I was thinking of constructing a for loop for each indices like so:

for i in range(3):
    w = np.where( abs(array3[:,i]).all() <= 20.  )
n_array = array3[w]

But I will only construct one value instead of many values.

Upvotes: 3

Views: 120

Answers (1)

Chiel
Chiel

Reputation: 6194

I think that using a 4D array is the easiest in this example. In that case, you can check whether the element with the lowest value in the vector is below the threshold. Then, with np.newaxis, you can apply the mask to the entire vector and create a masked array.

import numpy as np

n = 4 # Size of the first three dimensions.   
array3 = 100.*np.random.rand(n, n, n, 3) # Random numbers between 0 and 100.

thres = 20.
m = np.empty(array3.shape, dtype=bool)
m[:,:,:,:] = (np.min(array3, axis=-1) < thres)[:,:,:,np.newaxis]

array3_masked = np.ma.masked_array(array3, mask=m)

Upvotes: 1

Related Questions