Pfthb

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.