Split a multidimensional numpy array using a condition

Question

I have a multidimensional numpy array. The first array indicates the quality of the data. 0 is good, 1 is not so good. For a first check I only want to use good data. How do I split the array into two new ones? My own idea does not work:

good_data = [x for x in data[0,:] if x = 1.0]
bad_data = [x for x in data[0,:] if x = 0.0]

Here is a small example indicating my problem:

import numpy as np

flag = np.array([0., 0., 0., 1., 1., 1.])
temp = np.array([300., 310., 320., 300., 222., 333.])
pressure = np.array([1013., 1013., 1013., 900., 900., 900.])
data = np.array([flag, temp, pressure])

good_data = data[0,:][data[0,:] == 1.0]
bad_data  = data[0,:][data[0,:] == 0.0]

print good_data

The print statement gives me [1., 1., 1.].

But I am looking for [[1., 1., 1.], [300., 222., 333.], [900., 900., 900.]].

Answer 1

Is this what you are looking for?

good_data = data[0,:][data[0,:] == 1.0]
bad_data  = data[0,:][data[0,:] == 0.0]

This returns a numpy.array.

Alternatively, you can do as you suggested, but convert the resulting list to numpy.array:

good_data = np.array([x for x in data[0,:] if x == 1.0])

Notice the comparison operator == in place of the assignment operator =.

For your particular example, subset data using flag == 1 while iterating over the first index:

good_data = [data[n,:][flag == 1] for n in range(data.shape[0])]

If you really want the elements of good_data to be lists, convert inside the comprehension:

good_data = [data[n,:][flag == 1].tolist() for n in range(data.shape[0])]

Thanks to Jaime who pointed out that the easy way to do this is:

good_data = data[:, data[0] == 1]