numpy.ufunc 操作等效

Question

I have 4D numpy array a and 3D array b . Also I have 2D arrays of indices i0 , j0 , k0 . Suppose I want to use the following construction:

np.add.at(a, (slice(None), k0, i0, j0), b)

As fas as I understood, a[slice(None), k0, i0, j0] += b is not equivalent to this np.add.at . The question is how can this np.add.at line can be replaced with a simple numpy adding a[...] += b[...] properly?

Answer 1

You wrote that i0 , j0 and k0 are 2-D arrays.

This is why I assumed that each row in these arrays defines a 3-D slice in b array (each array in the respective dimension).

For the test I defined these arrays as:

k0 = np.array([[0,1],[1,2]])
i0 = np.array([[1,3],[2,4]])
j0 = np.array([[0,2],[2,4]])

so that the first 3-D slice from b is: b[0:1, 1:3, 0:2] .

Because k0 , i0 and j0 (taken "by row") define consecutive slices, you can not join them in a single instruction. My proposition is to perform your addition in the following loop:

for sl in np.dstack([k0, i0, j0]):
    idx = tuple(np.insert(np.r_[np.apply_along_axis(
        lambda row: slice(*row), 0, sl)], 0, slice(None)))
    a[tuple(idx)] += b[tuple(idx[1:])]

idx is a 4-tuple to index array " a ". The first element is slice(None) - operate on all elements in the first dimension. Following elements are slices for following dimensions.

To index b array the same idx is used, but without the first element.

To test the above code, I defined both arrays as:

aDim = (2, 2, 4, 4)
bDim = aDim[1:]
a = np.arange(1, np.array(aDim).prod() + 1).reshape(aDim)
b = np.arange(501, np.array(bDim).prod() + 501).reshape(bDim)

(print them to see the initial content).

After my code was executed, a contains:

array([[[[  1,   2,   3,   4],
         [510, 512,   7,   8],
         [518, 520,  11,  12],
         [ 13,  14,  15,  16]],

        [[ 17,  18,  19,  20],
         [ 21,  22,  23,  24],
         [ 25,  26, 554, 556],
         [ 29,  30, 562, 564]]],


       [[[ 33,  34,  35,  36],
         [542, 544,  39,  40],
         [550, 552,  43,  44],
         [ 45,  46,  47,  48]],

        [[ 49,  50,  51,  52],
         [ 53,  54,  55,  56],
         [ 57,  58, 586, 588],
         [ 61,  62, 594, 596]]]])