我计算数值偏导数的Python代码是否可以进一步简化？

Question

我想计算形状为 (Nmesh, Nmesh, Nmesh) 的 3D 笛卡尔网格的一阶偏导数，例如按四阶精度方案Nmesh=512，即

f'(n) = 2*(f(n+1)-f(n-1))/(3*dh) - (f(n+2)-f(n-2))/(12* DH）.

为此，我编写了一个Python函数，如下所示。但是，我发现不同条件分支之间的某些行非常相似。那么我可以进一步简化这个功能吗？

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

Answer 1

如 mozway 的评论 所示，您可以使用

swapaxes

函数执行以下操作：

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)