Home > Back-end >  Replacing a for loop in python by array slicing and numpy.diff when step size is larger than 1
Replacing a for loop in python by array slicing and numpy.diff when step size is larger than 1

Time:12-05

I want to calculate the derivative via finite differences, note though that I have a stepsize of 2 grid points here.

This can be easily realized in a for loop, which is however quite slow. Since I am working on numpy arrays, I replaced the for loop by array slicing to get a significant speed up. Then I wanted to test/compare numpy.diff but I get a different result using it.

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()

As can be seen from the figure, the result from numpy.diff looks different - what am I missing here?

enter image description here

CodePudding user response:

According to the docs, the first difference is given by out[i] = a[i 1] - a[i] along the given axis.

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

You are doing the opposite:

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