Оптимизация цикла для операций с круговыми матрицами

Question

У меня есть вектор X_k и матрица Y_ {k, j}, где k = (k, ..., K) и j = (1, .. J).Они круглые, это означает, что

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

Я хочу вычислить вектор Zx и матрицу Zy в соответствии со следующими выражениями:

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

В настоящее время я делаю это черезцикл, где сначала я вычисляю крайние случаи (для X: k = 1, 2, K. Для Y: j = 1, J, J-1).А для других я использую формулу.

Мне интересно, можно ли векторизовать это исчисление.Вот пример кода.

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)

Я подозреваю, что можно оптимизировать с помощью матричных операций, но не уверен, что это возможно без циклов (по крайней мере, для краевых случаев).

Answer 1

С Проверка кода Я получаю правильный ответ.Способ векторизации такого рода операции - через np.roll().

Например, переменная Y будет:

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