Suppose I have a collection of square matrices Sigma1, Sigma2, ...
. These are stored in a numpy array Sigma
such that Sigma[0, :, :] = Sigma1, Sigma[1, :, :] = Sigma2
, etc.
I have a corresponding set of vectors v1, v2, ...
, stored in a matrix v
such that v[0, :] = v1, v[1, :] = v2
, etc.
I am trying to produce a vector a such that a[0] is v1.T @ Sigma1 @ v1
(a scalar value), a[1] is v2.T @ Sigma2 @ v2
, etc.
I'm a bit bewildered by the documentation for einsum
tensordot
and the like, which seem promising but which I have been unable to coerce into giving me the desired answer.
For example:
Sigma = np.empty(4, 3, 3)
for i in range(4):
Z = np.random.randn(10, 3)
Sigma[i, :, :] = Z.T @ Z
y = np.random.randn(4, 3)
I want a function foo(Sigma y)
that returns a vector equivalent to
bar = np.empty(4)
for i in range 4:
bar[i] = y[i, :].T @ Sigma[i, :, :] @ y[i, :]
return bar
CodePudding user response:
np.einsum('ij,ijk,ik -> i', y, Sigma, y)
is what you are looking for.
In [1]: Sigma = np.random.rand(4, 3, 3)
In [2]: y = np.random.rand(4, 3)
In [3]: bar = np.empty(4)
...: for i in range(4):
...: bar[i] = y[i, :].T @ Sigma[i, :, :] @ y[i, :]
...:
In [4]: bar_einsum = np.einsum('ij,ijk,ik -> i', y, Sigma, y)
In [5]: np.allclose(bar, bar_einsum)
Out[5]: True
CodePudding user response:
Looks like you want (not tested)
y[i, :].T @ Sigma[i, :, :] @ y[i, :]
np.einsum('ij,ijk,ik',y,Sigma,y)
A matmul
version should be possible, but einsum
is easier to write