I want to compute the 1st-order partial derivatives of a 3D Cartesian mesh with shape (Nmesh, Nmesh, Nmesh), e.g. Nmesh=512 by the fourth-order accuracy scheme, i.e.
f'(n) = 2*(f(n+1)-f(n-1))/(3*dh) - (f(n+2)-f(n-2))/(12*dh) .
For this end, I wrote a Python function as shown below. However, I see that some rows are very similar between different conditional branches. So can I further simplify this function?
import numpy as np
def partial( arr, Nmesh, dh, axis=0 ):
'''
The partial derivative of the 3D Cartesian mesh with periodic
boundary conditions, computed by the fourth-order accuracy scheme.
'''
dh1 = 2/(3*dh)
dh2 = -1/(12*dh)
if (axis==0):
arr_ = np.pad(arr, [(2,2), (0,0), (0,0)], mode='wrap')
diff1 = arr_[3:Nmesh+3,:,:] - arr_[1:Nmesh+1,:,:]
diff2 = arr_[4:Nmesh+4,:,:] - arr_[0:Nmesh,:,:]
return dh1*diff1 + dh2*diff2
elif (axis==1):
arr_ = np.pad(arr, [(0,0), (2,2), (0,0)], mode='wrap')
diff1 = arr_[:,3:Nmesh+3,:] - arr_[:,1:Nmesh+1,:]
diff2 = arr_[:,4:Nmesh+4,:] - arr_[:,0:Nmesh,:]
return dh1*diff1 + dh2*diff2
elif (axis==2):
arr_ = np.pad(arr, [(0,0), (0,0), (2,2)], mode='wrap')
diff1 = arr_[:,:,3:Nmesh+3] - arr_[:,:,1:Nmesh+1]
diff2 = arr_[:,:,4:Nmesh+4] - arr_[:,:,0:Nmesh]
return dh1*diff1 + dh2*diff2
As in mozway's comment, you could use the swapaxes function to do something like:
import numpy as np
def partial(arr, Nmesh, dh, axis=0):
'''
The partial derivative of the 3D Cartesian mesh with periodic
boundary conditions, computed by the fourth-order accuracy scheme.
'''
dh1 = 2/(3*dh)
dh2 = -1/(12*dh)
# use swapaxes to put required axis first
arr_ = np.pad(
np.swapaxes(arr, axis, 0), [(2, 2), (0, 0), (0, 0)], mode="wrap"
)
diff1 = arr_[3:Nmesh + 3,:,:] - arr_[1:Nmesh + 1,:,:]
diff2 = arr_[4:Nmesh + 4,:,:] - arr_[0:Nmesh,:,:]
# swap axes back on return
return np.swapaxes(dh1*diff1 + dh2*diff2, 0, axis)