matrix multiplication with ndarray

Question

How to perform a vector and matrix multiplication where v is a vector (v1, v2, v3) and A is a 3x3 matrix? Python complains about shapes not aligned, maybe because v is an ndarray. Any ideas on how to do this operation? the final result should be a scalar at each coordinate point (v1, v2, v3). The basic code below breaks when trying to carry the multiplication.

import numpy as np
a = np.linspace(0, 10, 21)
b = np.linspace(0, 20, 41)
a, b = np.meshgrid(a,b)
v = np.array([a*b, a+b, a])
A = np.ones([3,3])
s = v.T @ A @ v     # doesn't work

Error

----> 1 s = v.T @ A @ v    
ValueError: shapes (21,41,3) and (3,41,21) not aligned: 3 (dim 2) != 41 (dim 1)

Edit: the matrix operation should be done at each point v, where v is usually a large array (of vectors). For example take a 1m cube with its center at the origin, and evaluate the matrix operation at each grid point, say every 10cm in each coordinate axis.

edit 2 example for a single point at (x,y,z)

A = np.zeros([3,3])
A[0][0] = 1
A[1][1] = 2
A[2][2] = 3
x,y,z = 1, 1, 0
v = np.array([x, y, z])
s = v.T @ A @ v   # should give s=3

The next step is to make the code work for a large array of vectors v. Except it's a little more involved because the vector coordinates (x,y,z) need to be parameterized in terms of coordinates (a,b). The original code above tries to do that, but doesn't work and may not be the best approach. Any other ideas?

Answer 1

When you multiply two N-dimensional matrices with numpy, I suppose it automatically multiplies the two last dimensions together and keeps the first ones. The N-dimensional matrix multiplication is more or less like the 2D multiplication. Your matrices must have the same shape except for 2 dimensions. In these 2 dimensions, you should respect the rules of the 2D multiplication. As an exemple, if you have a matrix A with shape (a,b,c, ...,d,e,f) and you want to multiply it with a matrix B, the shape of B should be (a,b,c,...,d,f,g) and the result's shape will be (a,b,c, ...,d,e,g).

Let's forget we are in a 4D space. If you had only one point, v^T*A*v should be shaped as (1,3)x(3,3)x(3,1). We simply want to apply that for each point in a (41,21) grid. This gives us the last dimensions of each of the component we need to multiply. To be consistent, v^T*A*v should have a shape of (41,21,1,3)x(3,3)x(41,21,3,1).

 import numpy as np
 a = np.linspace(0, 10, 21)
 b = np.linspace(0, 20, 41)
 a, b = np.meshgrid(a,b)
 a = np.expand_dims(a, axis=0)
 b = np.expand_dims(b, axis=0)
 print("Shape a = {}, Shape b = {}".format(a.shape, b.shape))
 v = np.array([a*b, a+b, a])
 print("Shape v = {}".format(v.shape))
 u1 = v.transpose((2,3,1,0))
 print("Shape u1 = {}".format(u1.shape))
 s = u1 @ A
 u2 = v.transpose((2,3,0,1))
 print("Shape u2 = {}".format(u2.shape))
 s = s @ u2
 print("{} x {} x {} = {} x {} = {}".format(u1.shape, A.shape, u2.shape, (u1 @ A).shape, u2.shape, s.shape))

returns :

Shape a = (1, 41, 21), Shape b = (1, 41, 21)
Shape v = (3, 1, 41, 21)
Shape u1 = (41, 21, 1, 3)
Shape u2 = (41, 21, 3, 1)
(41, 21, 1, 3) x (3, 3) x (41, 21, 3, 1) = (41, 21, 1, 3) x (41, 21, 3, 1) = (41, 21, 1, 1)

I propose you this solution. You start by adding a dimension of size 1 to your vectors a and b. Instead of having a shape of (41,21), they will have a shape of (1,41,21). Now, when you construct v, you obtain a shape of (3,1,41,21). Now, if you use the usual transpose, you just reverse all the dimensions, and that is not what you want. You want v^T to be multipliable by A, of shape (3,3). So you define by hand how to reverse the dimensions of your vector to go from (3,1,41,21) to (41,21,1,3) and to (41,21,3,1). At the end, you can finally multiply it, and it is consistent.

NOTE 1 In the theory, you can multiply wrt other dimensions than the last ones, as long as you respect the rule of the 2D multiplication for these dimensions, with same other dimensions. But this is the way to do it in Python.

NOTE 2 You may wonder why we can multiply a matrix of shape (41,21,1,3) by a matrix of shape (3,3). This is exactly the same mechanics that when you multiply a 2D matrix by a scalar. When you do that, you increase the dimension of the scalar to 2 dimensions (basically a matrix with the scalar everywhere), and you perform the element-wise multiplication. Similarly, you create a matrix of shape (41,21,3,3) and you multiply element-wise, or "block-wise" (2D matrix multiplication). The elements give multiplication (1,3)x(3,3).

Answer 2

It seems by the mention of the vector v of three elements, you meant an ndarray with three elements along its first axis, with each elements holding a n-dim array data. For the listed sample, you have the same as a 3D array. It also seems that the output is to be reduced to a scalar for each of the vector of three elements, i.e. the output would be 2D. So, to solve your case, we need to use tensor multiplication for the first : V.T @ A sum-reducing the first axes, giving us a 3D array. Then, use einsum to keep the first two axes aligned and sum-reduce the last ones, like so -

p1 = np.tensordot(v,A,axes=((0,0)))
out = np.einsum('jkl,ljk->jk',p1,v)

Alternatively, using einsum we can do all in one step, like so -

out = np.einsum('ijk,il,ljk->jk',v,A,v)

We can possibly make it faster with einsum's optional arg : optimize set as True :

np.einsum(..., optimize=True)