Vectorizing matrix operation instead of for loop in circular matrices

Question

I have a vector \$X_k\$ and a matrix \$Y_{k,j}\$, where \$k \in \{1,\dots,K\}\$ and \$j \in \{1, \dots, J\}\$. They are circular, meaning that $$\eqalign{X_{k+K} &= X_k \\ Y_{k+K,j} &= Y_{k,j} \\ Y_{k,j+J} &= Y_{k,j}}$$

I want to compute a \$ZX\$ vector and \$ZY\$ matrix according to the equations $$\eqalign{ZX_k &= -X_{k-1}(X_{K-2}-X_{K+1})-X_k \\ ZY_{k,j} &= -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 \$ZX\$, \$k = 1, 2, K\$; for \$ZY\$, \$j = 1, J, J-1\$). And for the others, I use the equations above.

I'm wondering if this calculation 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).

Any suggestion of how to improve the performance?

Answer 1

This is easy using numpy.roll, for example:

zx = np.roll(x, 1) * (np.roll(x, 2) + np.roll(x, -1)) - x