Ars3nous
Ars3nous

Reputation: 136

python summing over 2d array to get a 1d array

Say I have an 2d array (x,y) and the values of the function z=F(x,y)

   x = y = array([ 1.,  2.,  3.,  4.,  5.])

   z= array([[  1.,   2.,   3.,   4.,   5.],
       [  2.,   4.,   6.,   8.,  10.],
       [  3.,   6.,   9.,  12.,  15.],
       [  4.,   8.,  12.,  16.,  20.],
       [  5.,  10.,  15.,  20.,  25.]])

Now what i want to find is the integration $P(w)=\int F(x,y) \delta(x-y=w) dx dy $ for this I construct w matrix by doing w = x-y which gives me something like

   w= array([[ 0.,  1.,  2.,  3.,  4.],
       [-1.,  0.,  1.,  2.,  3.],
       [-2., -1.,  0.,  1.,  2.],
       [-3., -2., -1.,  0.,  1.],
       [-4., -3., -2., -1.,  0.]])

Now i must add up all values of z corresponding to w ,say w = 3 I should get 4+10=14.

Question is what is the best way to do the last part?

P.S. This is an example, the arrays are in general not equal and doesnt have the symmetry as in this example. Iterating would be a bad option, I guess, as these arrays are quite large.

Upvotes: 1

Views: 311

Answers (2)

user2357112
user2357112

Reputation: 280251

z[w == 3].sum()

w == 3 builds a boolean array representing which locations of w have a 3. z[w == 3] gives an array of the elements of z corresponding to those locations, and sum() adds them up. You'll learn a lot of this kind of stuff in the NumPy tutorial, and you'll learn even more in the NumPy reference.

Upvotes: 2

Daniel
Daniel

Reputation: 19547

This looks like a good place to use np.unique and numpy's new (v1.8.2) np.add.at function:

uvals, uidx = np.unique(w, return_inverse=True)

output = np.zeros_like(uvals)
np.add.at(output, uidx, z.ravel())

print uvals
# [-4. -3. -2. -1.  0.  1.  2.  3.  4.]
print output
# [  5.  14.  26.  40.  55.  40.  26.  14.   5.]

Upvotes: 1

Related Questions