当步长大于1时，通过数组切片和numpy.diff替换python中的for循环

Posted 2023-03-12

【中文标题】当步长大于1时，通过数组切片和numpy.diff替换python中的for循环

【英文标题】：Replacing a for loop in python by array slicing and numpy.diff when step size is larger than 1

【发布时间】：2022-01-09 22:37:03

【问题描述】：

我想通过有限差分计算导数，但请注意，我这里的步长为 2 个网格点。

这可以在for 循环中轻松实现，但速度很慢。由于我正在处理numpy 数组，因此我用数组切片替换了for 循环以显着提高速度。然后我想测试/比较numpy.diff，但使用它得到了不同的结果。

import matplotlib.pyplot as plt
import numpy as np

# create some data
n   = 100
x   = np.linspace(0,2*np.pi,n)
y   = np.linspace(0,2*np.pi,n)
X,Y = np.meshgrid(x,y)
Z   = np.sin(X) * np.sin(Y)

# calculate centered finite difference using for loop
Z_diff1 = Z*.0
for ii in range(1,Z.shape[0]-1,2):
    for jj in range(1,Z.shape[1]-1,2):
        Z_diff1[ii,jj]  = Z[ii-1, jj] - Z[ii+1,jj]

# calculate centered finite difference using array slicing
Z_diff2 = Z[:-2:2,::2] - Z[2::2,::2]

# calculate centered finite difference using numpy.diff
Z_diff3 = np.diff(Z[::2,::2], axis=0)

fig = plt.figure( figsize=(8,6) )

ax1 = fig.add_subplot( 1,4,1, aspect='equal' )
ax1.pcolormesh(Z)
ax1.set_title('original data')

ax2 = fig.add_subplot( 1,4,2, aspect='equal' )
ax2.pcolormesh(Z_diff1[1::2,1::2])
ax2.set_title('for loops')

ax3 = fig.add_subplot( 1,4,3, aspect='equal' )
ax3.pcolormesh(Z_diff2)
ax3.set_title('slicing')

ax4 = fig.add_subplot( 1,4,4, aspect='equal' )
ax4.pcolormesh(Z_diff3)
ax4.set_title('numpy.diff')

plt.show()

从图中可以看出，numpy.diff 的结果看起来不一样——我在这里遗漏了什么？

【参考方案1】：

根据docs，第一个差异由out[i] = a[i+1] - a[i]沿给定轴给出。

>>> a = np.array([1, 2, 3, 4, 5, 4, 3, 2, 1])
>>> np.diff(a[::2])
array([ 2,  2, -2, -2])

你正在做相反的事情：

>>> a[:-2:2] - a[2::2]
array([-2, -2,  2,  2])

【讨论】：

哦，确实，感谢您指出显而易见的解决方案:)