NumPy apply function to groups of rows corresponding to another numpy array

user_12323228 0
Viewed 583

I have a NumPy array with each row representing some (x, y, z) coordinate like so:

a = array([[0, 0, 1],
           [1, 1, 2],
           [4, 5, 1],
           [4, 5, 2]])

I also have another NumPy array with unique values of the z-coordinates of that array like so:

b = array([1, 2])

How can I apply a function, let's call it "f", to each of the groups of rows in a which correspond to the values in b? For example, the first value of b is 1 so I would get all rows of a which have a 1 in the z-coordinate. Then, I apply a function to all those values.

In the end, the output would be an array the same shape as b.

I'm trying to vectorize this to make it as fast as possible. Thanks!

Example of an expected output (assuming that f is count()):

c = array([2, 2])

because there are 2 rows in array a which have a z value of 1 in array b and also 2 rows in array a which have a z value of 2 in array b.

A trivial solution would be to iterate over array b like so:

for val in b:
    apply function to a based on val
    append to an array c

My attempt:

I tried doing something like this, but it just returns an empty array.

func(a[a[:, 2]==b])
numpy python
3 Answers

The problem is that the groups of rows with the same Z can have different sizes so you cannot stack them into one 3D numpy array which would allow to easily apply a function along the third dimension. One solution is to use a for-loop, another is to use np.split:

a = np.array([[0, 0, 1],
              [1, 1, 2],
              [4, 5, 1],
              [4, 5, 2],
              [4, 3, 1]])


a_sorted = a[a[:,2].argsort()]

inds = np.unique(a_sorted[:,2], return_index=True)[1]

a_split = np.split(a_sorted, inds)[1:]

# [array([[0, 0, 1],
#         [4, 5, 1],
#         [4, 3, 1]]),

#  array([[1, 1, 2],
#         [4, 5, 2]])]

f = np.sum  # example of a function

result = list(map(f, a_split))
# [19, 15]

But imho the best solution is to use pandas and groupby as suggested by FBruzzesi. You can then convert the result to a numpy array.

EDIT: For completeness, here are the other two solutions

List comprehension:

b = np.unique(a[:,2])
result = [f(a[a[:,2] == z]) for z in b]

Pandas:

df = pd.DataFrame(a, columns=list('XYZ'))
result = df.groupby(['Z']).apply(lambda x: f(x.values)).tolist()

This is the performance plot I got for a = np.random.randint(0, 100, (n, 3)):

enter image description here

As you can see, approximately up to n = 10^5 the "split solution" is the fastest, but after that the pandas solution performs better.

user_29453570

If you are allowed to use pandas:

import pandas as pd
df=pd.DataFrame(a, columns=['x','y','z'])

df.groupby('z').agg(f)

Here f can be any custom function working on grouped data.

Numeric example:

a = np.array([[0, 0, 1],
              [1, 1, 2],
              [4, 5, 1],
              [4, 5, 2]])
df=pd.DataFrame(a, columns=['x','y','z'])
df.groupby('z').size()

z
1    2
2    2
dtype: int64

Remark that .size is the way to count number of rows per group.

To keep it into pure numpy, maybe this can suit your case:

tmp = np.array([a[a[:,2]==i] for i in b])
tmp 
array([[[0, 0, 1],
        [4, 5, 1]],

       [[1, 1, 2],
        [4, 5, 2]]])

which is an array with each group of arrays.

user_124115360 
c = np.array([])
for x in np.nditer(b):
    c = np.append(c, np.where((a[:,2] == x))[0].shape[0])

Output:

[2. 2.]
user_115579740
Related