def get_initial_conditions():
    """For the initial conditions, we use a horizontal jet in geostrophic balance."""
    # global initial conditions
    u0_global = 10 * jnp.exp(
        -((yy_global - 0.5 * length_y) ** 2) / (0.02 * length_x) ** 2
    )
    v0_global = jnp.zeros_like(u0_global)

    # approximate balance h_y = -(f/g)u
    coriolis_global = CORIOLIS_F + yy_global * CORIOLIS_BETA
    h_geostrophy = np.cumsum(-dy * u0_global * coriolis_global / GRAVITY, axis=0)
    h0_global = (
        DEPTH
        + h_geostrophy
        # make sure h0 is centered around depth
        - h_geostrophy.mean()
        # small perturbation to break symmetry
        + 0.2
        * np.sin(xx_global / length_x * 10 * np.pi)
        * np.cos(yy_global / length_y * 8 * np.pi)
    )

    h0_local = h0_global[local_slice]
    u0_local = u0_global[local_slice]
    v0_local = v0_global[local_slice]

    token = jax.lax.create_token()
    h0_local, token = enforce_boundaries(h0_local, "h", token)
    u0_local, token = enforce_boundaries(u0_local, "u", token)
    v0_local, token = enforce_boundaries(v0_local, "v", token)

    return h0_local, u0_local, v0_local
For the initial conditions, we use a horizontal jet in geostrophic balance.
examples/shallow_water.py