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

Question:

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.

Asked By: Ars3nous

||

Answers:

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.

Answered By: user2357112

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.]
Answered By: Daniel