def solve_shallow_water(t1, num_multisteps=10):
    """Iterate the model forward in time."""
    # initial conditions
    h, u, v = get_initial_conditions()
    du, dv, dh = (jnp.zeros((ny_local, nx_local)) for _ in range(3))

    state = ModelState(h, u, v, dh, du, dv)
    sol = [state]

    state = shallow_water_step(state, True)
    sol.append(state)
    t = dt

    if HAS_TQDM:
        pbar = tqdm.tqdm(
            total=math.ceil(t1 / dt),
            disable=mpi_rank != 0,
            unit="model day",
            initial=t / DAY_IN_SECONDS,
            unit_scale=dt / DAY_IN_SECONDS,
            bar_format=(
                "{l_bar}{bar}| {n:.2f}/{total:.2f} [{elapsed}<{remaining}, "
                "{rate_fmt}{postfix}]"
            ),
        )

    # pre-compile JAX kernel
    do_multistep(state, num_multisteps)

    start = time.perf_counter()
    with ExitStack() as es:
        if HAS_TQDM:
            es.enter_context(pbar)

        while t < t1:
            state = do_multistep(state, num_multisteps)
            state[0].block_until_ready()
            sol.append(state)

            t += dt * num_multisteps

            if t < t1 and HAS_TQDM:
                pbar.update(num_multisteps)

    end = time.perf_counter()

    if mpi_rank == 0:
        print(f"\nSolution took {end - start:.2f}s")

    return sol
Iterate the model forward in time.
examples/shallow_water.py