CODEWITHSUNDEEP
pythonnumpy
Steve
Steve

Reputation: 43

Efficient calculation of element-wise matrix product with index dependent function in Python

I have this interesting problem where I want to calculate the sum over the element-wise product of three matrices

enter image description here

While calculating \mathbf{p}_ {ijk} and c_{ijk} can be done apriori, I have my problem with f_{ijk}(x,y,z). Elements of this matrix are multivariate polynomials which depend upon the matrix indices, thus numpy.vectorize cannot be trivially applied. My best bet at tackling the issue would be to treat the (i,j,k) as additional variables such that numpy.vectorize is then subsequently applied to a 6-dimensional instead of 3-dimensional input. However, I am not sure if more efficient or alternative ways exist.

Upvotes: 1

Views: 52

Answers (1)

javidcf
javidcf

Reputation: 59701

This is a simple way to implement that formula efficiently:

import numpy as np

np.random.seed(0)
l, m, n = 4, 5, 6
x, y, z = np.random.rand(3)
p = np.random.rand(l, m, n)
c = np.random.rand(l, m, n)
i, j, k = map(np.arange, (l, m, n))
xi = (x ** (l - i)) * (x ** l)
yj = (y ** (m - j)) * (y ** m)
zk = (z ** (n - k)) * (z ** n)
res = np.einsum('ijk,ijk,i,j,k->', p, c, xi, yj, zk)
print(res)
# 0.0007208482648476157

Or even slightly more compact:

import numpy as np

np.random.seed(0)
l, m, n = 4, 5, 6
x, y, z = np.random.rand(3)
p = np.random.rand(l, m, n)
c = np.random.rand(l, m, n)
t = map(lambda v, s: (v ** (s - np.arange(s))) * (v ** s), (x, y, z), (l, m, n))
res = np.einsum('ijk,ijk,i,j,k->', p, c, *t)
print(res)
# 0.0007208482648476157

Using np.einsum you minimize the need for intermediate arrays, so it should be faster that making f first (which you could get e.g. as f = np.einsum('i,j,k->ijk', xi, yj, zk)), multiplying p, c and f and then summing the result.

Upvotes: 2

Related Questions