Question

I have two 2D numpy arrays – real r, which contains points in space, given by their Cartesian coordinates, and v, a complex vector defined at each of these points. I would like to split both of these arrays, based on some condition on r.

e.g., r1 contains all points with the first cartesian coordinate is positive, and v1 gives the corresponding values of v. All other points and their corresponding vectors go into .

Based on this question, and the fact that zip is essentially it’s own inverse, I currently have the following solution:

r1, v1 = zip(*[rv for rv in zip(r, v) if rv[0][0] > 0.0])
r2, v2 = zip(*[rv for rv in zip(r, v) if rv[0][0] <= 0.0])
r1 = np.array(r1)
r2 = np.array(r2)
v1 = np.array(v1)
v2 = np.array(v2)

This works well enough for my purposes, however it involves conversion to large lists of arrays, which is surely quite inefficient.

Is there an alternative solution, which is fast, concise and avoids the creation of intermediate lists?

Answer 1

you can use bool array as index to filter out values:

create some random test data first:

import numpy as np
np.random.seed(0)
r = np.random.rand(10,2)-0.5
v = np.random.rand(10) + np.random.rand(10)*1j

then:

idx = r[:,0] > 0 # idx is a bool array 
r1 = r[idx]
v1 = v[idx]

r2 = r[~idx] # ~idx compute bit-wise NOT, element-wise
v2 = v[~idx]