numpy 有没有可能我的python代码计算数值偏导数可以进一步简化?

8oomwypt  于 2023-10-19  发布在  Python
关注(0)|答案(1)|浏览(98)

我想计算具有形状(Nmesh,Nmesh,Nmesh)的3D笛卡尔网格的一阶偏导数,例如。Nmesh=512通过四阶精度方案,即
f ′(n)= 2*(f(n+1)-f(n-1))/(3dh)-(f(n+2)-f(n-2))/(12dh)。
为此,我写了一个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
8zzbczxx

8zzbczxx1#

正如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)

相关问题