Question

我编写了一个计算仿真，该仿真接受一些可选参数（作为边界条件）。如下面的代码所示

def leapfrog(dx, dt, u_0, v_0, U, V, u_l=None, u_r=None, v_l=None, v_r=None, 
        n_max=None, **kwargs):
    """Every call returns a timestep of leapfrog iteration given:
        dx: Spatial step
        dt: Time step
        u0: Initial condition of regular grid
        v0: Initial condition of staggered grid
        U: Operator to apply over the regular grid
        V: Operator to apply over staggered grid
        u_l: left boundary condition to regular grid
        u_r: right boundary condition to regular grid
        v_l: left boundary condition to staggered grid
        v_r: right boundary conditiion to staggered grid
        n_max: If nmax is given, stops after n_max iterations (> 2)

        For every call returns another timestep

    """
    # Initial condition
    yield u_0, v_0

    # First iteration
    u = u_0 - (dt/dx) * (V @ v_0)
    if u_l is not None:
        u[:len(u_l)] = u_l
    if u_r is not None:
        u[-len(u_r):] = u_r

    v = v_0 - (dt/dx) * (U @ u_0)
    if v_l is not None:
        v[:len(v_l)] = v_l
    if v_r is not None:
        v[-len(v_r):] = v_r
    yield u, v

    # All other iterations
    for n in itertools.count(start=2, step=1):
        u = u - (dt/dx) * (V @ v)
        if u_l is not None:
            u[:u_l.size] = u_l
        if u_r is not None:
            u[-u_r.size:] = u_r

        v = v - (dt/dx) * (U @ u)
        if v_l is not None:
            v[:v_l.size] = v_l
        if v_r is not None:
            v[-v_r.size:] = v_r

        if n_max is not None and n >= n_max:
            break

        yield u,v

虽然代码可以正常工作，但是由于我检查是否在左右两边的u和v中添加了边界条件，因此肯定会触发CPU分支}和_l），例如_r表示主网格中的左边界条件。

u_l，u_l，u_r和v_l是像v_r这样的numpy数组。我需要添加尺寸，否则（如果是缩放器）它没有len，这也使事情变得统一。

在另一方面，有时我不想添加边界条件，在这种情况下，我使用np.array([0])语句来捕获它。

我的问题是，有什么办法可以避免分支（以及所有if）？

当然，我可以在有和没有边界条件的情况下（在主网格和辅助网格中）创建“专用”函数，但是我正在寻找一种更通用的方法。

此外，我们非常欢迎其他任何改进代码的提示。

在跨越式仿真中避免CPU分支

0 个答案: