Calculate Sum Of Nth Column Of Numpy Array Entry Grouped By The Indices In First Two Columns?

I would like to loop over following check_matrix in such a way that code recognize whether the first and second element is 1 and 1 or 1 and 2 etc? Then for each separate class of p

Solution 1:

I am pretty sure you can solve this problem in an easier way and I am not 100% sure that I understood you correctly, but here is some code that does what I think you want. If you have a possibility to check if the results are valid, I would suggest you do so.

import numpy as np

n = 20
q = np.zeros((20, 3))
q[:, -1] = np.linspace(0, 10, n)

check_matrix = np.array([[1, 1, 0, 0, 0, 0,    0,       -0.7977, -0.243293],
                         [1, 1, 0, 0, 0, 0,    0,       1.5954,  0.004567],
                         [1, 2, 0, 0, 0, -1,   0,       0,       1.126557],
                         [2, 1, 0, 0, 0, 0.5,  0.86603, 1.5954,  0.038934],
                         [2, 1, 0, 0, 0, 2,    0,       -0.7977, -0.015192],
                         [2, 2, 0, 0, 0, -0.5, 0.86603, 1.5954,  0.21394]])
check_matrix[:, :2] -= 1# python indexing is zero based

matrices = np.zeros((n, 2, 2), dtype=np.complex_)

for i inrange(2):
    for j inrange(2):
        k_list = []
        for k inrange(len(check_matrix)):
            if check_matrix[k][0] == i and check_matrix[k][1] == j:
                k_list.append(check_matrix[k][8] *
                              np.exp(-1J * np.dot(q, check_matrix[k][2:5] 
                                                     - check_matrix[k][5:8])))

        matrices[:, i, j] = np.sum(k_list, axis=0)

NOTE: I changed your indices to have consistent zero-based indexing.

---

Here is another approach where I replaced the k-loop with a vectored version:

for i in range(2):
    for j in range(2):
        k = np.logical_and(check_matrix[:, 0] == i, check_matrix[:, 1] == j)
        temp = np.dot(check_matrix[k, 2:5] - check_matrix[k, 5:8], q[:, :, np.newaxis])[..., 0]
        temp = check_matrix[k, 8:] * np.exp(-1J * temp)
        matrices[:, i, j] = np.sum(temp, axis=0)

Solution 2:

3 line solution

You asked for efficient solution in your original title so how about this solution that avoids nested loops and if statements in a 3 liner, which is thus hopefully faster?

fac=2*(check_matrix[:,0]-1)+(check_matrix[:,1]-1)
grp=np.split(check_matrix[:,8], np.cumsum(np.unique(fac,return_counts=True)[1])[:-1])
[np.sum(x) for x in grp]

output:

[-0.23872600000000002, 1.126557, 0.023742000000000003, 0.21394]

How does it work?

I combine the first two columns into a single index, treating each as "bits" (i.e. base 2)

fac=2*(check_matrix[:,0]-1)+(check_matrix[:,1]-1)

( If you have indexes that exceed 2, you can still use this technique but you will need to use a different base to combine the columns. i.e. if your indices go from 1 to 18, you would need to multiply column 0 by a number equal to or larger than 18 instead of 2. )

So the result of the first line is

array([0., 0., 1., 2., 2., 3.])

Note as well it assumes the data is ordered, that one column changes fastest, if this is not the case you will need an extra step to sort the index and the original check matrix. In your example the data is ordered.

The next step groups the data according to the index, and uses the solution posted here.

np.split(check_matrix[:,8], np.cumsum(np.unique(fac,return_counts=True)[1])[:-1])

[array([-0.243293,  0.004567]), array([1.126557]), array([ 0.038934, -0.015192]), array([0.21394])]

i.e. it outputs the 8th column of check_matrix according to the grouping of fac

then the last line simply sums those... knowing how the first two columns were combined to give the single index allows you to map the result back. Or you could simply add it to check matrix as a 9th column if you wanted.