def navier_stokes_rk(tableau: ButcherTableau, equation: ExplicitNavierStokesODE, time_step: float) -> TimeStepFn:
    'Create a forward Runge-Kutta time-stepper for incompressible Navier-Stokes.\n\n  This function implements the reference method (equations 16-21), rather than\n  the fast projection method, from:\n    "Fast-Projection Methods for the Incompressible Navier–Stokes Equations"\n    Fluids 2020, 5, 222; doi:10.3390/fluids5040222\n\n  Args:\n    tableau: Butcher tableau.\n    equation: equation to use.\n    time_step: overall time-step size.\n\n  Returns:\n    Function that advances one time-step forward.\n  '
    dt = time_step
    F = tree_math.pytree_to_vector_fun(equation.explicit_terms)
    P = tree_math.pytree_to_vector_fun(equation.pressure_projection)
    a = tableau.a
    b = tableau.b
    num_steps = len(b)

    @tree_math.vector_to_pytree_fun
    def step_fn(u0):
        u = ([None] * num_steps)
        k = ([None] * num_steps)
        u[0] = u0
        k[0] = F(u0)
        for i in range(1, num_steps):
            u_star = (u0 + (dt * sum(((a[(i - 1)][j] * k[j]) for j in range(i) if a[(i - 1)][j]))))
            u[i] = P(u_star)
            k[i] = F(u[i])
        u_star = (u0 + (dt * sum(((b[j] * k[j]) for j in range(num_steps) if b[j]))))
        u_final = P(u_star)
        return u_final
    return step_fn
Create a forward Runge-Kutta time-stepper for incompressible Navier-Stokes.

This function implements the reference method (equations 16-21), rather than
the fast projection method, from:
  "Fast-Projection Methods for the Incompressible Navier–Stokes Equations"
  Fluids 2020, 5, 222; doi:10.3390/fluids5040222

Args:
  tableau: Butcher tableau.
  equation: equation to use.
  time_step: overall time-step size.

Returns:
  Function that advances one time-step forward.
def navier_stokes_rk(tableau: ButcherTableau, equation: ExplicitNavierStokesODE, time_step: float) -> TimeStepFn:
    'Create a forward Runge-Kutta time-stepper for incompressible Navier-Stokes.\n\n  This function implements the reference method (equations 16-21), rather than\n  the fast projection method, from:\n    "Fast-Projection Methods for the Incompressible Navier–Stokes Equations"\n    Fluids 2020, 5, 222; doi:10.3390/fluids5040222\n\n  Args:\n    tableau: Butcher tableau.\n    equation: equation to use.\n    time_step: overall time-step size.\n\n  Returns:\n    Function that advances one time-step forward.\n  '
    dt = time_step
    F = tree_math.pytree_to_vector_fun(equation.explicit_terms)
    P = tree_math.pytree_to_vector_fun(equation.pressure_projection)
    a = tableau.a
    b = tableau.b
    num_steps = len(b)

    @tree_math.vector_to_pytree_fun
    def step_fn(u0):
        u = ([None] * num_steps)
        k = ([None] * num_steps)
        u[0] = u0
        k[0] = F(u0)
        for i in range(1, num_steps):
            u_star = (u0 + (dt * sum(((a[(i - 1)][j] * k[j]) for j in range(i) if a[(i - 1)][j]))))
            u[i] = P(u_star)
            k[i] = F(u[i])
        u_star = (u0 + (dt * sum(((b[j] * k[j]) for j in range(num_steps) if b[j]))))
        u_final = P(u_star)
        return u_final
    return step_fn