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.]].

OK, see the update to my answer.

milancurcic
– milancurcic

2013-10-24 21:04:22 +00:00
Commented Oct 24, 2013 at 21:04 — milancurcic
– milancurcic, Commented Oct 24, 2013 at 21:04

milancurcic · Accepted Answer · 2013-10-25 15:28:39Z

2

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]

edited Oct 25, 2013 at 15:28

answered Oct 24, 2013 at 20:29

milancurcic

6,2413 gold badges39 silver badges49 bronze badges

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

2 Comments

paulchen Over a year ago

Thanks for your help. Unfortunately it is not really what I was looking for. I will try to specify my question, please see above.

Jaime Over a year ago

That's the idea, but you are overcomplicating things, it's actually data[:, data[0] == 1] and data[:, data[0] == 0] that the OP needs.

Collectives™ on Stack Overflow

Split a multidimensional numpy array using a condition

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related