Elegant Numpy ways to implement this matrix operation-CodePudding

I'm looking to vectorize the following equation:

Where x represents an N x T list of observations, and gamma is an N x K matrix. The intended output, theta, is a K x T matrix. Of course, I have a naive implementation that works by just looping through the elements and setting the array manually, but I just know there's a much cleaner solution lurking out here somewhere.

Below is my naive implementation, however it is much too slow for my purposes.

    theta = np.zeros((K, T))
    for k in range(K):
        numerator = np.zeros((T,))
        denominator = 0
        for i in range(N):
            gamma_ik = r[i,k]
            numerator = numerator   gamma_ik * x[i]
            denominator = denominator   gamma_ik
        theta[k] = numerator / denominator

CodePudding user response：

Slightly more explicit answer that may be easier to follow. Again, the key function call is numpy.matmul.

import numpy as np

# dimensions N, T, K
N = 2
T = 3
K = 4

x = np.cumsum(np.ones([N,T]), axis = 0)
#[[1. 1. 1.]
#[2. 2. 2.]]

gamma = np.cumsum(np.ones([K,N]), axis = 0)
#[[1. 1.]
# [2. 2.]
# [3. 3.]
# [4. 4.]]

y = np.matmul(gamma, x) # numerator
#[[ 3.  3.  3.]
# [ 6.  6.  6.]
# [ 9.  9.  9.]
# [12. 12. 12.]]

# denominator
sumgamma = np.matmul(gamma, np.ones(N))
#[2. 4. 6. 8.]

theta = y / sumgamma[:,None]
#[[1.5 1.5 1.5]
# [1.5 1.5 1.5]
# [1.5 1.5 1.5]
# [1.5 1.5 1.5]]

CodePudding user response：

Both the numerator and denominator can straightforwardly be expressed as matrix products.

import numpy as np

n = 3
G = np.arange(n**2).reshape((n,n)) # the matrix Gamma
x = np.array([1,2,4])              # the vector x
y = np.ones(3)

theta = (G.T@x)/(G.T@y)
print('theta =',theta)
# theta = [3.33333333 3.08333333 2.93333333]