tkowt - 1 year ago 143

Python Question

I want bellow calculation.Values used in this example are variable in real case.

`A = [[1,1,1],[2,1,1]]`

B =[[[1,2,3],[2,2,2]],[[1,1,2],[1,1,1]]]

result = apply_some_function(A,B)

>> result = [[[6,7],[6,8]],

[[4,5],[3,4]]]

Is their any efficient way satisfy above calculation with numpy.

In fact, I don't understand clearly how to applicate numpy for multi dimensional array. So if you know helpful document to understand manipulation multi dimensional array rules, I'm very glad if you let me know.

Answer Source

You could use `np.tensordot`

or `np.einsum`

.

```
In [19]: np.tensordot(B, A, (-1, -1))
Out[19]:
array([[[6, 7],
[6, 8]],
[[4, 5],
[3, 4]]])
In [20]: np.einsum('ij,klj->kli', A, B)
Out[20]:
array([[[6, 7],
[6, 8]],
[[4, 5],
[3, 4]]])
```

Both of these functions compute sums of products. Thus, they are generalizations of matrix multiplication.

It is often helpful to pay attention to the shape of the arrays. If we make `A`

, `B`

, and `result`

NumPy arrays:

```
A = np.array([[1,1,1],[2,1,1]])
B = np.array([[[1,2,3],[2,2,2]],[[1,1,2],[1,1,1]]])
result = np.array([[[6,7],[6,8]], [[4,5],[3,4]]])
```

then

```
In [6]: A.shape
Out[6]: (2, 3)
In [7]: B.shape
Out[7]: (2, 2, 3)
In [9]: result.shape
Out[9]: (2, 2, 2)
```

Notice that the axis of length 3 in `A`

and `B`

disappears in `result`

. That suggests the sum (of products) is being taken over the last axis of `A`

and `B`

.

The `(-1, -1)`

in `np.tensordot(B, A, (-1, -1))`

tells `np.tensordot`

to sum over the last axes of `A`

and `B`

.

Similarly, the `'ij,klj->kli'`

in `np.einsum('ij,klj->kli', A, B)`

says that if `A`

has indices `i`

and `j`

and if `B`

has indices `k`

, `l`

and `j`

, then the result should have indices `k`

,`l`

,`i`

. Notice the `j`

index disappears. The `j`

index is the last index in `A`

and `B`

. Thus, `'ij,klj->kli'`

tells `np.einsum`

to sum over the last index of `A`

and `B`

.

The only thing left to do is figure out the correct order of the `k`

, `l`

and `i`

indices. Since each axis in `result`

has the same length, the shape of the `result`

gives no clue. I found the correct order by trial and error.