Xtykutl

I have written a computational simulation that accepts some optional parameters (as boundary conditions). As shown in the code below

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

While the code works fine, it for sure triggers CPU branching since I check if I added boundary condition in main and secondary grid (u and v) at left and right (_l and _r), for example, u_l means left boundary condition in main grid.

u_l, u_r, v_l and v_r are numpy arrays like np.array([0]). I need to add dimension, otherwise (if a scaler) it has no len and also that makes things more uniform.

In other hands sometimes I do not want to add boundary conditions, in these cases I use the if statement to catch it.

My question is, is there any way to avoid the branching (and all ifs)?

Of course, I could create 'specialized' functions with and without boundary condition (in main and secondary grids), but I'm looking a more generic approach.

Also, any other hints to improve the code are very welcome.