Reputation: 11
I'm trying to write the following expression using the einsum function in NumPy:
for j in range(100):
p[j] = 0
for i in range(100):
if i!=j:
p[j] += S[i,j]*B[T[i,j], i]
p.shape = (1,100)
S.shape = (100,100)
B.shape = (N,100)
T.shape = (100,100)
N is chosen such that it exceeds the maximum value in the matrix T.
Is this possible to implement using the einsum function? If not, is there another way to optimise this?
I've searched the documentation, and it looks like it'll need some kind of broadcasting operation before being inserted into the einsum expression. But I can't get the expression right.
Upvotes: 1
Views: 54
Reputation: 6583
I don't think einsum
has any nested indexing. But you can use vectorization:
import numpy as np
# Preliminaries:
S = np.random.rand(100, 100)
B = np.random.rand(10, 100)
T = np.random.randint(10, size=(100, 100))
# Solution:
idx = np.arange(100)[:, None]
arr = S * B[T, idx]
p = arr.sum(axis=0) - arr.diagonal()
I subtract the diagonal to account for the fact that you don't want to sum when i == j
. You may choose to use keepdims=True
in the sum to get p
as shape (1, 100)
instead.
Upvotes: 1