Question

现在我有一个按预期工作的二维功能：

def wavestep2d(U,dx, dz, dt):
    U = (dt/dx)**2 * fourth_order(U, dx, axis=0) +
        (dt/dz)**2 * fourth_order(U, dz, axis=1)   
    return U

当然我可以扩展到3d只添加一行：

def wavestep3d(U,dv, dt):
    """dv here is a N sized array with dx_1, dx_2, ..., dx_N"""
    U = (dt/dv[0])**2 * fourth_order(U, dv[0], axis=0) +
        (dt/dv[1])**2 * fourth_order(U, dv[1], axis=1) +
        (dt/dv[2])**2 * fourth_order(U, dv[2], axis=2)   
    return U

并且可以对4d（一条线），5d（另外两条线）做同样的事情并继续进行。但无限不是我的朋友。因此，我想为N维案例扩展它，如下所示：

def wavestepNd(U, dv, dt):
    """dv here is a N sized array with dx_1, dx_2, ..., dx_N"""
    for i in range(len(U.shape)):
        U += (dt/dv[i])**2 * fourth_order(U, dv[i], axis=i)   
    return U

这里的问题是在for循环的每次迭代中使用相同的U来改变U.

当然，我可以这样做：

def wavestepNd_with_copy(U, dv, dt):
    """dv here is a N sized array with dx_1, dx_2, ..., dx_N"""
    U_temp = np.zeros_like(U)
    for i in range(len(U.shape)):
        U_temp += (dt/dv[i])**2 * fourth_order(U, dv[i], axis=i)  
    return U_temp

但是这种情况会导致内存翻倍，而U会变得非常大。

对此有好的方法吗？

编辑1：

按要求添加引用的函数。

def fourth_order_term(U, axis=-1):
    U = U.swapaxes(0, axis)
    fm2 = U[ :-4]
    fm1 = U[1:-3]
    fc0 = U[2:-2]
    fp1 = U[3:-1]
    fp2 = U[4:  ]
    return (fm2 - 16*(fm1+fp1) + 30*fc0 + fp2).swapaxes(0, axis)

def fourth_order(U, dv, axis=-1):
    mul = 1/(12 * dv[axis]**2)
    return mul * fourth_order_term(U, axis=axis)

没有数据复制的N维求和

编辑1：

0 个答案: