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

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]


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?

  • 2
    \$\begingroup\$ Can you explain the purpose of this code? What do \$X\$, \$Y\$, \$ZX\$, \$ZY\$ represent? Can you check the mathematics in the post — in particular, the post says \$ZX_k = -X_{k-1}(X_{K-2}-X_{K+1})-X_k\$ but the code computes \$ZX_k = -X_{k-1}(X_{k-2}-X_{k+1})-X_k\$? \$\endgroup\$ Commented Oct 12, 2018 at 6:29
  • \$\begingroup\$ It sounds like you're after a re-write rather than a review. \$\endgroup\$ Commented Oct 12, 2018 at 7:35
  • \$\begingroup\$ I was after numpy.roll() function. And now the code works vectorized. I'm not sure if that's going after a re-writte or a review. \$\endgroup\$
    – Xbel
    Commented Oct 12, 2018 at 7:55

1 Answer 1


This is easy using numpy.roll, for example:

zx = np.roll(x, 1) * (np.roll(x, 2) + np.roll(x, -1)) - x
  • \$\begingroup\$ Thanks, that is what I was looking for. Thanks as well for tidying the post. \$\endgroup\$
    – Xbel
    Commented Oct 12, 2018 at 6:59

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.