CODEWITHSUNDEEP
pythonperformanceloopsmatrixoptimization
Xbel
Xbel

Reputation: 777

Optimization of loop for circular matrices operations

I have a vector X_k and a matrix Y_{k,j}, where k = (k,...,K) and j = (1, .. J). They are circular, that means that 

X_{k+K} = X_k, Y_{k+K,j} = Y_{k,j} and Y_{k,j+J} = Y_{k,j}

I want to compute a Zx vector and Zy matrix according to the following expressions:

Zx = -X_{k-1}*(X_{K-2}-X_{K+1})-X_k

Zy = -Y_{k,j+1}*(Y_{k,j+2}-Y_{k,j-1})-Y_{k,j}+X_k

Currently, I'm doing this through a loop, where first I compute the edge cases (for X: k = 1, 2, K. For Y: j = 1, J, J-1). And for the others, I use the formula.

I'm wondering if this calculus can be vectorized. Here is the example code. 

import numpy as np

np.random.seed(10)
K = 20
J = 10

# initial state (equilibrium)
x = np.random.uniform(size=K) 
y = np.random.uniform(size=(K*J)) 
y = y.reshape(K,J)

# zy
zy = np.zeros((K*J))
zy = zy.reshape(K,J)

# Edge case of Y
for k in range(K):
    zy[k,0] = -y[k,1]*(y[k,2]-y[k,J-1])-y[k,0]+ x[k]
    zy[k,J-1] = -y[k,0]*(y[k,1]-y[k,J-2])-y[k,J-1]+ x[k]
    zy[k,J-2] = -y[k,J-1]*(y[k,0]-y[k,J-3])-y[k,J-2]+ x[k]

    # General Case of Y
    for j in range(1,J-2):
        zy[k,j] = -y[k,j+1]*(y[k,j+2]-y[k,j-1])-y[k,j]+ x[k]

# zx
zx = np.zeros(K)

# first the 3 edge cases: k = 1, 2, K
zx[0] = -x[K-1]*(-x[1] + x[K-2])  - x[0]
zx[1] = - x[0]*(-x[2] + x[K-1])- x[1] 
zx[K-1] = -x[K-2]*(-x[0] + x[K-3]) - x[K-1]

# then the general case for X
for k in range(2, K-1)
    zx[k] = -x[k-1]*(-x[k+1] + x[k-2]) - x[k]

print(zx)
print(zy)

I suspect that is possible to optimize with matrix operations but not sure if it is possible without the loops (at least for the edge cases).

Upvotes: 0

Views: 50

Answers (1)

Xbel
Xbel

Reputation: 777

From Code Review I get the correct answer. The way to vectorize this kind of operation is through np.roll().

For example, the Y variable would be:

zy = np.roll(y,-1,axis=1)*(np.roll(y,-2,axis=1)-np.roll(y,1,axis=1))- y +x[:,None]

Upvotes: 1

Related Questions