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\$ – Gareth Rees Oct 12 '18 at 6:29
  • \$\begingroup\$ It sounds like you're after a re-write rather than a review. \$\endgroup\$ – Toby Speight Oct 12 '18 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 Oct 12 '18 at 7:55

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 Oct 12 '18 at 6:59

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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