从 3D numpy 数组中删除切片

Question

I have a 3D numpy array (I call it a tensor) shaped (5,8, 15000) .

Because of some calculations which filled it, there are some NaN 's inside the tensor. The last axis shows the simulation index. I had a process which repeated 15,000 times on the computer with slightly changed dynamics.

I want to go through all the 15,000 simulations and look at the 2D array (shaped (5,8) ) corresponding to one such simulation. If there's a NaN across the 2nd row, across all the columns in the 2D array (ie [1,:] ), I want to delete the simulation from the 3D tensor. This shall repeat for all 15,000 simulations.

That is, the new tensor shall be of shape (5, 8, 15000 - number_of_bad_sims) , with all the 2D arrays corresponding to that simulations which contained at least 1 NaN deleted. The rest of the tensor shall remain the same. I don't want to touch any of the 2D arrays which don't have a NaN across their 2nd row.

I have tried:

# threeDimTensor has shape (5,8, 15000)

for idx in range(threeDimTensor.shape[2]): # for all simulations
    boolean_array = np.isnan(threeDimTensor[1, :, idx]) # check if any of the results is NaN
    condition = False
    for element in boolean_array:
        if element == True and condition == False:
            condition = True
            np.delete(threeDimTensor, idx, axis=2) # along last axis ???

My question is how do I store the np.delete() result? I don't know how many NaN's I have, thus I cannot initially create a np.zeros(input_shape) 3D array to be populated with not-NaN's values.

Also, is my np.delete() instruction correct?

What would be a correct, efficient way to store the ''clean'' new three dimensional tensor? That is, how do I finish this piece of code to return the new three dimensional tensor and to use it further?

Thank you!

Answer 1

You could just have used isnan for your entire array, generate an array that masks the bad sims from that, and use it to select the good sims

My advice for using Numpy is to vectorize your code, try to think of a way in which you do not need explicit for-loops:

is_element_nan = np.isnan(threeDimTensor[1, :, :]) # Across your 2nd row
any_nan = np.any(is_element_nan, axis=0) #Flattens the 2d matrix
clean_tensor = threeDimTensor[:, :, ~any_nan] # ~is bitwise not