Time Integration

The generic time steppers live in integrix (integrix.march); finitevolX keeps the semi-Lagrangian scheme and the solve_ocean_pde wrapper.

Deprecated finitevolX names

finitevolx.RK3SSP, finitevolx.rk3_ssp_step and the other generic steppers below still work for one release and emit a DeprecationWarning. The diffrax solvers and the pure step functions are the integrix objects themselves, so import them from integrix.march. finitevolx.AB2Solver, finitevolx.LeapfrogRAFSolver and finitevolx.SplitExplicitRKSolver are adapters that keep the old init / step call shapes; their integrix replacements are integrix.march.AbstractSteppers run by integrix.march.run(stepper, rhs, t0, y0, dt, n_steps) with rhs(t, y, args).

Explicit Runge-Kutta (Pure Functional)

integrix.march.euler_step(state, rhs_fn, dt)

One forward-Euler step: yₙ₊₁ = yₙ + h f(yₙ).

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required

Returns:

Type	Description
PyTree	The state yₙ₊₁.

Examples:

>>> import jax.numpy as jnp
>>> bool(euler_step(jnp.array(1.0), lambda y: -y, 0.1) == 0.9)
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

def euler_step(
    state: PyTree, rhs_fn: Callable[[PyTree], PyTree], dt: ScalarLike
) -> PyTree:
    """One forward-Euler step: yₙ₊₁ = yₙ + h f(yₙ).

    Args:
        state: The state yₙ, any pytree.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.

    Returns:
        The state yₙ₊₁.

    Examples:
        >>> import jax.numpy as jnp
        >>> bool(euler_step(jnp.array(1.0), lambda y: -y, 0.1) == 0.9)
        True
    """
    k1 = rhs_fn(state)
    return jax.tree.map(lambda y, f: y + dt * f, state, k1)

integrix.march.heun_step(state, rhs_fn, dt)

One Heun (RK2) step, order 2, SSP with C = 1.

k₁ = f(yₙ), k₂ = f(yₙ + h k₁), yₙ₊₁ = yₙ + (h/2)(k₁ + k₂).

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required

Returns:

Type	Description
PyTree	The state yₙ₊₁.

Examples:

>>> import jax.numpy as jnp
>>> bool(jnp.allclose(heun_step(jnp.array(1.0), lambda y: -y, 0.1), 0.905))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

def heun_step(
    state: PyTree, rhs_fn: Callable[[PyTree], PyTree], dt: ScalarLike
) -> PyTree:
    """One Heun (RK2) step, order 2, SSP with C = 1.

    k₁ = f(yₙ), k₂ = f(yₙ + h k₁), yₙ₊₁ = yₙ + (h/2)(k₁ + k₂).

    Args:
        state: The state yₙ, any pytree.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.

    Returns:
        The state yₙ₊₁.

    Examples:
        >>> import jax.numpy as jnp
        >>> bool(jnp.allclose(heun_step(jnp.array(1.0), lambda y: -y, 0.1), 0.905))
        True
    """
    k1 = rhs_fn(state)
    state_star = jax.tree.map(lambda y, f: y + dt * f, state, k1)
    k2 = rhs_fn(state_star)
    return jax.tree.map(lambda y, f1, f2: y + 0.5 * dt * (f1 + f2), state, k1, k2)

integrix.march.rk3_ssp_step(state, rhs_fn, dt)

One third-order SSP Runge–Kutta step in Shu–Osher form, C = 1.

y⁽¹⁾ = yₙ + h f(yₙ); y⁽²⁾ = ¾ yₙ + ¼ y⁽¹⁾ + ¼ h f(y⁽¹⁾); yₙ₊₁ = ⅓ yₙ + ⅔ y⁽²⁾ + ⅔ h f(y⁽²⁾).

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required

Returns:

Type	Description
PyTree	The state yₙ₊₁.

Examples:

>>> import jax.numpy as jnp
>>> y = rk3_ssp_step(jnp.array(1.0), lambda y: -y, 0.1)
>>> bool(jnp.allclose(y, 1.0 - 0.1 + 0.1**2 / 2 - 0.1**3 / 6))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

def rk3_ssp_step(
    state: PyTree, rhs_fn: Callable[[PyTree], PyTree], dt: ScalarLike
) -> PyTree:
    """One third-order SSP Runge–Kutta step in Shu–Osher form, C = 1.

    y⁽¹⁾ = yₙ + h f(yₙ); y⁽²⁾ = ¾ yₙ + ¼ y⁽¹⁾ + ¼ h f(y⁽¹⁾);
    yₙ₊₁ = ⅓ yₙ + ⅔ y⁽²⁾ + ⅔ h f(y⁽²⁾).

    Args:
        state: The state yₙ, any pytree.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.

    Returns:
        The state yₙ₊₁.

    Examples:
        >>> import jax.numpy as jnp
        >>> y = rk3_ssp_step(jnp.array(1.0), lambda y: -y, 0.1)
        >>> bool(jnp.allclose(y, 1.0 - 0.1 + 0.1**2 / 2 - 0.1**3 / 6))
        True
    """
    k1 = rhs_fn(state)
    y1 = jax.tree.map(lambda y, f: y + dt * f, state, k1)

    k2 = rhs_fn(y1)
    y2 = jax.tree.map(
        lambda y, y1_, f: 0.75 * y + 0.25 * y1_ + 0.25 * dt * f,
        state,
        y1,
        k2,
    )

    k3 = rhs_fn(y2)
    return jax.tree.map(
        lambda y, y2_, f: (1.0 / 3.0) * y + (2.0 / 3.0) * y2_ + (2.0 / 3.0) * dt * f,
        state,
        y2,
        k3,
    )

integrix.march.rk4_step(state, rhs_fn, dt)

One classical fourth-order Runge–Kutta step.

k₁ = f(yₙ), k₂ = f(yₙ + (h/2) k₁), k₃ = f(yₙ + (h/2) k₂), k₄ = f(yₙ + h k₃), yₙ₊₁ = yₙ + (h/6)(k₁ + 2k₂ + 2k₃ + k₄).

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required

Returns:

Type	Description
PyTree	The state yₙ₊₁.

Examples:

>>> import jax.numpy as jnp
>>> y = rk4_step({"a": jnp.array(1.0)}, lambda y: {"a": -y["a"]}, 0.1)
>>> bool(jnp.allclose(y["a"], jnp.exp(-0.1), rtol=1e-6))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

def rk4_step(
    state: PyTree, rhs_fn: Callable[[PyTree], PyTree], dt: ScalarLike
) -> PyTree:
    """One classical fourth-order Runge–Kutta step.

    k₁ = f(yₙ), k₂ = f(yₙ + (h/2) k₁), k₃ = f(yₙ + (h/2) k₂), k₄ = f(yₙ + h k₃),
    yₙ₊₁ = yₙ + (h/6)(k₁ + 2k₂ + 2k₃ + k₄).

    Args:
        state: The state yₙ, any pytree.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.

    Returns:
        The state yₙ₊₁.

    Examples:
        >>> import jax.numpy as jnp
        >>> y = rk4_step({"a": jnp.array(1.0)}, lambda y: {"a": -y["a"]}, 0.1)
        >>> bool(jnp.allclose(y["a"], jnp.exp(-0.1), rtol=1e-6))
        True
    """
    k1 = rhs_fn(state)
    y1 = jax.tree.map(lambda y, f: y + 0.5 * dt * f, state, k1)

    k2 = rhs_fn(y1)
    y2 = jax.tree.map(lambda y, f: y + 0.5 * dt * f, state, k2)

    k3 = rhs_fn(y2)
    y3 = jax.tree.map(lambda y, f: y + dt * f, state, k3)

    k4 = rhs_fn(y3)
    return jax.tree.map(
        lambda y, f1, f2, f3, f4: y + (dt / 6.0) * (f1 + 2.0 * f2 + 2.0 * f3 + f4),
        state,
        k1,
        k2,
        k3,
        k4,
    )

Multistep Methods (Pure Functional)

integrix.march.ab2_step(state, rhs_fn, dt, rhs_nm1)

One second-order Adams–Bashforth step.

yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁), one evaluation of f per step.

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required
rhs_nm1	PyTree	fₙ₋₁, the right-hand side one step back.	required

Returns:

Type	Description
tuple[PyTree, PyTree, PyTree]	(yₙ₊₁, fₙ, fₙ₋₁): thread fₙ as the next rhs_nm1.

Examples:

>>> import jax.numpy as jnp
>>> y, f_n, _ = ab2_step(jnp.array(1.0), lambda y: -y, 0.1, -1.0)
>>> bool(jnp.allclose(y, 0.9)), bool(f_n == -1.0)
(True, True)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def ab2_step(
    state: PyTree,
    rhs_fn: Callable[[PyTree], PyTree],
    dt: ScalarLike,
    rhs_nm1: PyTree,
) -> tuple[PyTree, PyTree, PyTree]:
    """One second-order Adams–Bashforth step.

    yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁), one evaluation of f per step.

    Args:
        state: The state yₙ, any pytree.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.
        rhs_nm1: fₙ₋₁, the right-hand side one step back.

    Returns:
        ``(yₙ₊₁, fₙ, fₙ₋₁)``: thread fₙ as the next ``rhs_nm1``.

    Examples:
        >>> import jax.numpy as jnp
        >>> y, f_n, _ = ab2_step(jnp.array(1.0), lambda y: -y, 0.1, -1.0)
        >>> bool(jnp.allclose(y, 0.9)), bool(f_n == -1.0)
        (True, True)
    """
    rhs_n = rhs_fn(state)
    # yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁)
    new_state = jax.tree.map(
        lambda y, fn, fnm1: y + (dt / 2.0) * (3.0 * fn - fnm1),
        state,
        rhs_n,
        rhs_nm1,
    )
    return new_state, rhs_n, rhs_nm1

integrix.march.ab3_step(state, rhs_fn, dt, rhs_nm1, rhs_nm2)

One third-order Adams–Bashforth step.

yₙ₊₁ = yₙ + (h/12)(23 fₙ − 16 fₙ₋₁ + 5 fₙ₋₂), one evaluation of f per step. The start must be at least second-order accurate to keep the global order 3: bootstrap y₁ and y₂ with rk3_ssp_step, as AB3Solver does; the crude start fₙ₋₂ = fₙ₋₁ = f(y₀) degrades AB3 to second order.

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required
rhs_nm1	PyTree	fₙ₋₁, the right-hand side one step back.	required
rhs_nm2	PyTree	fₙ₋₂, two steps back.	required

Returns:

Type	Description
PyTree	(yₙ₊₁, fₙ, fₙ₋₁): thread fₙ as the next rhs_nm1 and fₙ₋₁ as
PyTree	the next rhs_nm2.

Examples:

>>> import jax.numpy as jnp
>>> f = lambda y: -y
>>> y, f_n, f_nm1 = ab3_step(jnp.array(1.0), f, 0.1, f(1.0), f(1.0))
>>> bool(jnp.allclose(y, 0.9)), bool(f_n == -1.0)
(True, True)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def ab3_step(
    state: PyTree,
    rhs_fn: Callable[[PyTree], PyTree],
    dt: ScalarLike,
    rhs_nm1: PyTree,
    rhs_nm2: PyTree,
) -> tuple[PyTree, PyTree, PyTree]:
    """One third-order Adams–Bashforth step.

    yₙ₊₁ = yₙ + (h/12)(23 fₙ − 16 fₙ₋₁ + 5 fₙ₋₂), one evaluation of f per
    step. The start must be at least second-order accurate to keep the
    global order 3: bootstrap y₁ and y₂ with `rk3_ssp_step`, as `AB3Solver`
    does; the crude start fₙ₋₂ = fₙ₋₁ = f(y₀) degrades AB3 to second order.

    Args:
        state: The state yₙ, any pytree.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.
        rhs_nm1: fₙ₋₁, the right-hand side one step back.
        rhs_nm2: fₙ₋₂, two steps back.

    Returns:
        ``(yₙ₊₁, fₙ, fₙ₋₁)``: thread fₙ as the next ``rhs_nm1`` and fₙ₋₁ as
        the next ``rhs_nm2``.

    Examples:
        >>> import jax.numpy as jnp
        >>> f = lambda y: -y
        >>> y, f_n, f_nm1 = ab3_step(jnp.array(1.0), f, 0.1, f(1.0), f(1.0))
        >>> bool(jnp.allclose(y, 0.9)), bool(f_n == -1.0)
        (True, True)
    """
    rhs_n = rhs_fn(state)
    # yₙ₊₁ = yₙ + (h/12)(23 fₙ − 16 fₙ₋₁ + 5 fₙ₋₂)
    new_state = jax.tree.map(
        lambda y, fn, fnm1, fnm2: (
            y + (dt / 12.0) * (23.0 * fn - 16.0 * fnm1 + 5.0 * fnm2)
        ),
        state,
        rhs_n,
        rhs_nm1,
        rhs_nm2,
    )
    return new_state, rhs_n, rhs_nm1

integrix.march.leapfrog_raf_step(state, state_nm1, rhs_fn, dt, alpha=0.05)

One leapfrog step with a Robert–Asselin filter on the middle level.

ỹₙ₊₁ = yₙ₋₁ + 2h f(yₙ) and ȳₙ = yₙ + α (yₙ₋₁ − 2yₙ + ỹₙ₊₁). The filtered ȳₙ damps the computational mode of the three-level scheme. Second order for α = 0 and first order for α > 0 (the filter's O(αh) damping; the measured rate is about 1.1 at the default 0.05).

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
state_nm1	PyTree	The (filtered) state yₙ₋₁.	required
rhs_fn	Callable[[PyTree], PyTree]	The autonomous right-hand side f(y).	required
dt	ScalarLike	The step h.	required
alpha	float	The filter coefficient α, typically 0.01 to 0.1.	0.05

Returns:

Type	Description
tuple[PyTree, PyTree]	(yₙ₊₁, ȳₙ): yₙ₊₁ is the next state and ȳₙ the next state_nm1.

Examples:

>>> import jax.numpy as jnp
>>> y_next, filtered = leapfrog_raf_step(1.0, 1.0, lambda y: 0.0 * y, 0.1)
>>> bool(y_next == 1.0), bool(filtered == 1.0)
(True, True)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def leapfrog_raf_step(
    state: PyTree,
    state_nm1: PyTree,
    rhs_fn: Callable[[PyTree], PyTree],
    dt: ScalarLike,
    alpha: float = 0.05,
) -> tuple[PyTree, PyTree]:
    """One leapfrog step with a Robert–Asselin filter on the middle level.

    ỹₙ₊₁ = yₙ₋₁ + 2h f(yₙ) and ȳₙ = yₙ + α (yₙ₋₁ − 2yₙ + ỹₙ₊₁). The filtered
    ȳₙ damps the computational mode of the three-level scheme. Second order
    for α = 0 and first order for α > 0 (the filter's O(αh) damping; the
    measured rate is about 1.1 at the default 0.05).

    Args:
        state: The state yₙ, any pytree.
        state_nm1: The (filtered) state yₙ₋₁.
        rhs_fn: The autonomous right-hand side f(y).
        dt: The step h.
        alpha: The filter coefficient α, typically 0.01 to 0.1.

    Returns:
        ``(yₙ₊₁, ȳₙ)``: yₙ₊₁ is the next state and ȳₙ the next ``state_nm1``.

    Examples:
        >>> import jax.numpy as jnp
        >>> y_next, filtered = leapfrog_raf_step(1.0, 1.0, lambda y: 0.0 * y, 0.1)
        >>> bool(y_next == 1.0), bool(filtered == 1.0)
        (True, True)
    """
    rhs_n = rhs_fn(state)

    # ỹₙ₊₁ = yₙ₋₁ + 2h f(yₙ)
    y_next = jax.tree.map(
        lambda ynm1, fn: ynm1 + 2.0 * dt * fn,
        state_nm1,
        rhs_n,
    )

    # ȳₙ = yₙ + α (yₙ₋₁ − 2yₙ + ỹₙ₊₁)
    state_filtered = jax.tree.map(
        lambda ynm1, yn, ynp1: yn + alpha * (ynm1 - 2.0 * yn + ynp1),
        state_nm1,
        state,
        y_next,
    )

    return y_next, state_filtered

IMEX (Pure Functional)

integrix.march.imex_ssp2_step(state, rhs_explicit, rhs_implicit, implicit_solve, dt)

One IMEX-SSP2(2,2,2) step, with the caller's implicit solve.

Second order in the explicit part, the implicit part and the coupled split (the implicit row (1 − 2γ, γ) gives ∑ᵢ bᵢ cᵢ = ½).

Algorithm

1. Y₁ = yₙ + γh f_I(Y₁): implicit_solve(yₙ, γh).
2. Y₂ = yₙ + h f_E(Y₁) + (1 − 2γ) h f_I(Y₁), then Y₂ = Y₂ + γh f_I(Y₂): implicit_solve(Y₂*, γh).
3. yₙ₊₁ = yₙ + (h/2)(f_E(Y₁) + f_E(Y₂)) + (h/2)(f_I(Y₁) + f_I(Y₂)).

Parameters:

Name	Type	Description	Default
state	PyTree	The state yₙ, any pytree.	required
rhs_explicit	Callable[[PyTree], PyTree]	The non-stiff right-hand side f_E(y).	required
rhs_implicit	Callable[[PyTree], PyTree]	The stiff right-hand side f_I(y).	required
implicit_solve	Callable[[PyTree, ScalarLike], PyTree]	Solves Y − γh f_I(Y) = r for Y, given (r, γh); for vertical diffusion, a tridiagonal solve along columns.	required
dt	ScalarLike	The step h.	required

Returns:

Type	Description
PyTree	The state yₙ₊₁.

Examples:

>>> import jax.numpy as jnp
>>> y = imex_ssp2_step(
...     jnp.array(1.0),
...     lambda y: -y,
...     lambda y: -10.0 * y,
...     lambda r, gamma_dt: r / (1.0 + 10.0 * gamma_dt),
...     0.01,
... )
>>> bool(jnp.allclose(y, jnp.exp(-0.11), rtol=1e-4))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_imex.py

def imex_ssp2_step(
    state: PyTree,
    rhs_explicit: Callable[[PyTree], PyTree],
    rhs_implicit: Callable[[PyTree], PyTree],
    implicit_solve: Callable[[PyTree, ScalarLike], PyTree],
    dt: ScalarLike,
) -> PyTree:
    """One IMEX-SSP2(2,2,2) step, with the caller's implicit solve.

    Second order in the explicit part, the implicit part and the coupled
    split (the implicit row (1 − 2γ, γ) gives ∑ᵢ bᵢ cᵢ = ½).

    Algorithm:
        1. Y₁ = yₙ + γh f_I(Y₁): ``implicit_solve(yₙ, γh)``.
        2. Y₂* = yₙ + h f_E(Y₁) + (1 − 2γ) h f_I(Y₁), then
           Y₂ = Y₂* + γh f_I(Y₂): ``implicit_solve(Y₂*, γh)``.
        3. yₙ₊₁ = yₙ + (h/2)(f_E(Y₁) + f_E(Y₂)) + (h/2)(f_I(Y₁) + f_I(Y₂)).

    Args:
        state: The state yₙ, any pytree.
        rhs_explicit: The non-stiff right-hand side f_E(y).
        rhs_implicit: The stiff right-hand side f_I(y).
        implicit_solve: Solves Y − γh f_I(Y) = r for Y, given (r, γh); for
            vertical diffusion, a tridiagonal solve along columns.
        dt: The step h.

    Returns:
        The state yₙ₊₁.

    Examples:
        >>> import jax.numpy as jnp
        >>> y = imex_ssp2_step(
        ...     jnp.array(1.0),
        ...     lambda y: -y,
        ...     lambda y: -10.0 * y,
        ...     lambda r, gamma_dt: r / (1.0 + 10.0 * gamma_dt),
        ...     0.01,
        ... )
        >>> bool(jnp.allclose(y, jnp.exp(-0.11), rtol=1e-4))
        True
    """
    gamma_dt = _GAMMA * dt

    # Stage 1: Y₁ = yₙ + γh f_I(Y₁)
    y1 = implicit_solve(state, gamma_dt)
    fe_1 = rhs_explicit(y1)
    fi_1 = rhs_implicit(y1)

    # Stage 2: Y₂* = yₙ + h f_E(Y₁) + (1 − 2γ) h f_I(Y₁); Y₂ = Y₂* + γh f_I(Y₂)
    y2_star = jax.tree.map(
        lambda y, fe, fi: y + dt * fe + (1.0 - 2.0 * _GAMMA) * dt * fi,
        state,
        fe_1,
        fi_1,
    )
    y2 = implicit_solve(y2_star, gamma_dt)
    fe_2 = rhs_explicit(y2)
    fi_2 = rhs_implicit(y2)

    # yₙ₊₁ = yₙ + (h/2)(f_E(Y₁) + f_E(Y₂)) + (h/2)(f_I(Y₁) + f_I(Y₂))
    return jax.tree.map(
        lambda y, fe1, fe2, fi1, fi2: (
            y + 0.5 * dt * (fe1 + fe2) + 0.5 * dt * (fi1 + fi2)
        ),
        state,
        fe_1,
        fe_2,
        fi_1,
        fi_2,
    )

Split-Explicit (Pure Functional)

integrix.march.split_explicit_step(state_3d, state_2d, rhs_3d, rhs_2d, couple_2d_to_3d, dt_slow, n_substeps, slow_forcing=None, t0=0.0)

One split-explicit barotropic/baroclinic step.

The fast subcycle is neutrally stable for the oscillatory gravity-wave mode (RK4 amplification 1 − (ωh)⁶/144 per substep); the slow step is first order in h. The returned 2-D state is the end-of-cycle y_M, which stays on time; the slow step used the mean ⟨y⟩, and making the 3-D state consistent with it (replacing the depth mean of the 3-D velocity by the averaged transport, say) is couple_2d_to_3d's job.

Algorithm

1. F = slow_forcing(state_3d), if given, held fixed over the cycle.
2. Subcycle the 2-D state with M RK4 steps of h/M on y′ = rhs_2d(t, y, state_3d) + F, and form its trapezoidal mean ⟨y⟩ over the cycle.
3. One forward-Euler step of the 3-D state: y₃ + h rhs_3d(state_3d, ⟨y⟩).
4. couple_2d_to_3d(new state_3d, ⟨y⟩).

Parameters:

Name	Type	Description	Default
state_3d	PyTree	The slow (baroclinic) state.	required
state_2d	PyTree	The fast (barotropic) state.	required
rhs_3d	Callable[[PyTree, PyTree], PyTree]	The slow tendency rhs_3d(state_3d, mean_2d).	required
rhs_2d	Callable[[ScalarLike, PyTree, PyTree], PyTree]	The fast tendency rhs_2d(t, state_2d, state_3d), at the absolute RK4 stage time t.	required
couple_2d_to_3d	Callable[[PyTree, PyTree], PyTree]	couple_2d_to_3d(state_3d, mean_2d), applied after the slow step.	required
dt_slow	ScalarLike	The slow step h.	required
n_substeps	int	The number of fast substeps M ≥ 1 per slow step.	required
slow_forcing	Callable[[PyTree], PyTree] | None	slow_forcing(state_3d), the baroclinic forcing of the barotropic mode, evaluated once per slow step and added to every fast tendency.	None
t0	ScalarLike	The time at the start of the slow step.	0.0

Returns:

Type	Description
tuple[PyTree, PyTree]	(new_state_3d, new_state_2d).

Raises:

Type	Description
ValidationError	If n_substeps is below 1.

Examples:

>>> import jax.numpy as jnp
>>> slow, fast = split_explicit_step(
...     jnp.array(1.0),
...     jnp.array(1.0),
...     lambda y3, mean_2d: -mean_2d,
...     lambda t, y2, y3: 0.0 * y2,
...     lambda y3, mean_2d: y3,
...     0.1,
...     5,
... )
>>> bool(jnp.allclose(slow, 0.9)), bool(fast == 1.0)
(True, True)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_split_explicit.py

def split_explicit_step(
    state_3d: PyTree,
    state_2d: PyTree,
    rhs_3d: Callable[[PyTree, PyTree], PyTree],
    rhs_2d: Callable[[ScalarLike, PyTree, PyTree], PyTree],
    couple_2d_to_3d: Callable[[PyTree, PyTree], PyTree],
    dt_slow: ScalarLike,
    n_substeps: int,
    slow_forcing: Callable[[PyTree], PyTree] | None = None,
    t0: ScalarLike = 0.0,
) -> tuple[PyTree, PyTree]:
    """One split-explicit barotropic/baroclinic step.

    The fast subcycle is neutrally stable for the oscillatory gravity-wave
    mode (RK4 amplification 1 − (ωh)⁶/144 per substep); the slow step is first
    order in h. The returned 2-D state is the end-of-cycle y_M, which stays on
    time; the slow step used the mean ⟨y⟩, and making the 3-D state consistent
    with it (replacing the depth mean of the 3-D velocity by the averaged
    transport, say) is ``couple_2d_to_3d``'s job.

    Algorithm:
        1. F = ``slow_forcing(state_3d)``, if given, held fixed over the cycle.
        2. Subcycle the 2-D state with M RK4 steps of h/M on
           y′ = ``rhs_2d(t, y, state_3d)`` + F, and form its trapezoidal
           mean ⟨y⟩ over the cycle.
        3. One forward-Euler step of the 3-D state:
           y₃ + h ``rhs_3d(state_3d, ⟨y⟩)``.
        4. ``couple_2d_to_3d(new state_3d, ⟨y⟩)``.

    Args:
        state_3d: The slow (baroclinic) state.
        state_2d: The fast (barotropic) state.
        rhs_3d: The slow tendency ``rhs_3d(state_3d, mean_2d)``.
        rhs_2d: The fast tendency ``rhs_2d(t, state_2d, state_3d)``, at the
            absolute RK4 stage time t.
        couple_2d_to_3d: ``couple_2d_to_3d(state_3d, mean_2d)``, applied after
            the slow step.
        dt_slow: The slow step h.
        n_substeps: The number of fast substeps M ≥ 1 per slow step.
        slow_forcing: ``slow_forcing(state_3d)``, the baroclinic forcing of the
            barotropic mode, evaluated once per slow step and added to every
            fast tendency.
        t0: The time at the start of the slow step.

    Returns:
        ``(new_state_3d, new_state_2d)``.

    Raises:
        ValidationError: If ``n_substeps`` is below 1.

    Examples:
        >>> import jax.numpy as jnp
        >>> slow, fast = split_explicit_step(
        ...     jnp.array(1.0),
        ...     jnp.array(1.0),
        ...     lambda y3, mean_2d: -mean_2d,
        ...     lambda t, y2, y3: 0.0 * y2,
        ...     lambda y3, mean_2d: y3,
        ...     0.1,
        ...     5,
        ... )
        >>> bool(jnp.allclose(slow, 0.9)), bool(fast == 1.0)
        (True, True)
    """
    if n_substeps < 1:
        raise ValidationError(f"n_substeps must be >= 1, got {n_substeps}")

    dt_fast = dt_slow / n_substeps
    forcing = None if slow_forcing is None else slow_forcing(state_3d)

    # --- fast (barotropic) subcycle ---
    y_2d_end, y_2d_avg = _barotropic_cycle(
        state_2d,
        lambda t_sub, y_2d: rhs_2d(t_sub, y_2d, state_3d),
        t0,
        dt_fast,
        n_substeps,
        forcing,
    )

    # --- slow (baroclinic) step ---
    f_slow = rhs_3d(state_3d, y_2d_avg)
    y_3d_new = jax.tree.map(lambda y, f: y + dt_slow * f, state_3d, f_slow)

    # --- coupling ---
    y_3d_new = couple_2d_to_3d(y_3d_new, y_2d_avg)

    return y_3d_new, y_2d_end

Semi-Lagrangian (Pure Functional)

finitevolx.semi_lagrangian_step(field, u, v, dx, dy, dt, interp_order=1, bc='periodic')

Advect a 2D scalar field using semi-Lagrangian backtracking.

Algorithm::

1. Compute departure points: x_dep = x_i - u * dt,
                              y_dep = y_j - v * dt
2. Interpolate ``field`` at the departure points.
3. Return interpolated values as the new field.

Parameters:

Name	Type	Description	Default
field	Array[Ny, Nx]	Scalar field to advect.	required
u	Array[Ny, Nx]	Velocity components at the same grid points as field, in physical units (m/s).	required
v	Array[Ny, Nx]	Velocity components at the same grid points as field, in physical units (m/s).	required
dx	float	Grid spacing in x and y (m).	required
dy	float	Grid spacing in x and y (m).	required
dt	float	Timestep (s).	required
interp_order	int	Interpolation order passed to :func:jax.scipy.ndimage.map_coordinates. Currently JAX only supports order <= 1. Default 1.	1
bc	str	Boundary handling: "periodic" (wrap) or "edge" (Neumann-like clamp). Default "periodic".	'periodic'

Returns:

Type	Description
Array[Ny, Nx]	Advected field.

Source code in finitevolx/_src/timestepping/semi_lagrangian.py

def semi_lagrangian_step(
    field: jax.Array,
    u: jax.Array,
    v: jax.Array,
    dx: float,
    dy: float,
    dt: float,
    interp_order: int = 1,
    bc: str = "periodic",
) -> jax.Array:
    """Advect a 2D scalar field using semi-Lagrangian backtracking.

    Algorithm::

        1. Compute departure points: x_dep = x_i - u * dt,
                                      y_dep = y_j - v * dt
        2. Interpolate ``field`` at the departure points.
        3. Return interpolated values as the new field.

    Parameters
    ----------
    field : Array[Ny, Nx]
        Scalar field to advect.
    u, v : Array[Ny, Nx]
        Velocity components at the same grid points as ``field``, in
        **physical units** (m/s).
    dx, dy : float
        Grid spacing in x and y (m).
    dt : float
        Timestep (s).
    interp_order : int, optional
        Interpolation order passed to :func:`jax.scipy.ndimage.map_coordinates`.
        Currently JAX only supports ``order <= 1``.  Default 1.
    bc : str, optional
        Boundary handling: ``"periodic"`` (wrap) or ``"edge"``
        (Neumann-like clamp).  Default ``"periodic"``.

    Returns
    -------
    Array[Ny, Nx]
        Advected field.
    """
    if bc not in {"periodic", "edge"}:
        raise ValueError(f"bc must be 'periodic' or 'edge', got {bc!r}")
    if interp_order not in {0, 1}:
        raise ValueError(
            f"interp_order must be 0 or 1 (JAX limitation), got {interp_order}"
        )

    ny, nx = field.shape

    # Target grid coordinates (in index space)
    y_coords, x_coords = jnp.meshgrid(
        jnp.arange(ny, dtype=field.dtype),
        jnp.arange(nx, dtype=field.dtype),
        indexing="ij",
    )

    # Departure points in index space
    x_dep = x_coords - u * dt / dx
    y_dep = y_coords - v * dt / dy

    mode = "wrap" if bc == "periodic" else "nearest"

    return jax.scipy.ndimage.map_coordinates(
        field,
        [y_dep, x_dep],
        order=interp_order,
        mode=mode,
    )

Diffrax Solvers

integrix.march.ForwardEulerDfx

Bases: AbstractSolver

Forward Euler as a diffrax solver: yₙ₊₁ = yₙ + h F(tₙ, yₙ).

Order 1, one stage, no error estimate (fixed steps only). Included for completeness and parity with finitevolX; diffrax.Euler is the same scheme.

Examples:

>>> import diffrax as dfx
>>> import jax.numpy as jnp
>>> solution = dfx.diffeqsolve(
...     dfx.ODETerm(lambda t, y, args: -y),
...     ForwardEulerDfx(),
...     t0=0.0,
...     t1=1.0,
...     dt0=0.1,
...     y0=jnp.array(1.0),
... )
>>> bool(jnp.allclose(solution.ys[0], 0.9**10))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

class ForwardEulerDfx(dfx.AbstractSolver):
    """Forward Euler as a diffrax solver: yₙ₊₁ = yₙ + h F(tₙ, yₙ).

    Order 1, one stage, no error estimate (fixed steps only). Included for
    completeness and parity with finitevolX; `diffrax.Euler` is the same
    scheme.

    Examples:
        >>> import diffrax as dfx
        >>> import jax.numpy as jnp
        >>> solution = dfx.diffeqsolve(
        ...     dfx.ODETerm(lambda t, y, args: -y),
        ...     ForwardEulerDfx(),
        ...     t0=0.0,
        ...     t1=1.0,
        ...     dt0=0.1,
        ...     y0=jnp.array(1.0),
        ... )
        >>> bool(jnp.allclose(solution.ys[0], 0.9**10))
        True
    """

    term_structure: ClassVar[Any] = dfx.AbstractTerm
    interpolation_cls: ClassVar[Any] = dfx.LocalLinearInterpolation

    def order(self, terms):
        return 1

    def init(self, terms, t0, t1, y0, args):
        return None

    def step(self, terms, t0, t1, y0, args, solver_state, made_jump):
        del solver_state, made_jump
        dt = t1 - t0
        f0 = terms.vf(t0, y0, args)
        # yₙ₊₁ = yₙ + h F(tₙ, yₙ)
        y1 = jax.tree.map(lambda y, f: y + dt * f, y0, f0)
        dense_info = dict(y0=y0, y1=y1)
        return y1, None, dense_info, None, RESULTS.successful

    def func(self, terms, t0, y0, args):
        return terms.vf(t0, y0, args)

integrix.march.RK2Heun

Bases: AbstractERK

Heun's method (RK2), with an embedded forward-Euler error estimate.

Order 2, two stages, SSP with coefficient C = 1. The error weights are b̂ − b = [½, −½] (as diffrax.Heun), so adaptive controllers work.

Butcher tableau: c = (0, 1), a₂₁ = 1, b = (½, ½).

Examples:

>>> RK2Heun().order(None)
2

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

class RK2Heun(dfx.AbstractERK):
    """Heun's method (RK2), with an embedded forward-Euler error estimate.

    Order 2, two stages, SSP with coefficient C = 1. The error weights are
    b̂ − b = [½, −½] (as `diffrax.Heun`), so adaptive controllers work.

    Butcher tableau: c = (0, 1), a₂₁ = 1, b = (½, ½).

    Examples:
        >>> RK2Heun().order(None)
        2
    """

    tableau: ClassVar[dfx.ButcherTableau] = dfx.ButcherTableau(
        c=np.array([1.0]),
        b_sol=np.array([0.5, 0.5]),
        # b_error = b_euler − b_sol, b_euler = [1, 0]
        b_error=np.array([0.5, -0.5]),
        a_lower=(np.array([1.0]),),
    )
    interpolation_cls: ClassVar[Any] = (
        dfx.ThirdOrderHermitePolynomialInterpolation.from_k
    )

    def order(self, terms):
        return 2

integrix.march.RK3SSP

Bases: AbstractERK

Third-order strong-stability-preserving Runge–Kutta (Shu–Osher).

Order 3, three stages, SSP coefficient C = 1 (optimal): it keeps the monotonicity, positivity and TVD properties of forward Euler within the same CFL limit. The embedded second-order solution is Heun on the first two stages (b̂ = [½, ½, 0]), so adaptive controllers work.

Butcher tableau: c = (0, 1, ½), a₂₁ = 1, a₃₁ = a₃₂ = ¼, b = (⅙, ⅙, ⅔).

Examples:

>>> RK3SSP().order(None)
3

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

class RK3SSP(dfx.AbstractERK):
    """Third-order strong-stability-preserving Runge–Kutta (Shu–Osher).

    Order 3, three stages, SSP coefficient C = 1 (optimal): it keeps the
    monotonicity, positivity and TVD properties of forward Euler within the
    same CFL limit. The embedded second-order solution is Heun on the first
    two stages (b̂ = [½, ½, 0]), so adaptive controllers work.

    Butcher tableau: c = (0, 1, ½), a₂₁ = 1, a₃₁ = a₃₂ = ¼, b = (⅙, ⅙, ⅔).

    Examples:
        >>> RK3SSP().order(None)
        3
    """

    tableau: ClassVar[dfx.ButcherTableau] = dfx.ButcherTableau(
        c=np.array([1.0, 0.5]),
        b_sol=np.array([1.0 / 6, 1.0 / 6, 2.0 / 3]),
        # b_error = b_hat − b_sol, b_hat = [1/2, 1/2, 0] (Heun, 2nd order)
        b_error=np.array([1.0 / 3, 1.0 / 3, -2.0 / 3]),
        a_lower=(
            np.array([1.0]),
            np.array([0.25, 0.25]),
        ),
    )
    interpolation_cls: ClassVar[Any] = (
        dfx.ThirdOrderHermitePolynomialInterpolation.from_k
    )

    def order(self, terms):
        return 3

integrix.march.RK4Classic

Bases: _FixedStepERK

The classical fourth-order Runge–Kutta method, fixed steps only.

Order 4, four stages, not SSP. The tableau has no embedded lower-order solution, so an adaptive controller raises; diffrax.Tsit5 or diffrax.Dopri5 are the adaptive fourth- and fifth-order choices.

Butcher tableau: c = (0, ½, ½, 1), a₂₁ = a₃₂ = ½, a₄₃ = 1, b = (⅙, ⅓, ⅓, ⅙).

Examples:

>>> RK4Classic().order(None)
4

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

class RK4Classic(_FixedStepERK):
    """The classical fourth-order Runge–Kutta method, fixed steps only.

    Order 4, four stages, not SSP. The tableau has no embedded lower-order
    solution, so an adaptive controller raises; `diffrax.Tsit5` or
    `diffrax.Dopri5` are the adaptive fourth- and fifth-order choices.

    Butcher tableau: c = (0, ½, ½, 1), a₂₁ = a₃₂ = ½, a₄₃ = 1,
    b = (⅙, ⅓, ⅓, ⅙).

    Examples:
        >>> RK4Classic().order(None)
        4
    """

    tableau: ClassVar[dfx.ButcherTableau] = dfx.ButcherTableau(
        c=np.array([0.5, 0.5, 1.0]),
        b_sol=np.array([1.0 / 6, 1.0 / 3, 1.0 / 3, 1.0 / 6]),
        b_error=np.zeros(4),
        a_lower=(
            np.array([0.5]),
            np.array([0.0, 0.5]),
            np.array([0.0, 0.0, 1.0]),
        ),
    )
    interpolation_cls: ClassVar[Any] = (
        dfx.ThirdOrderHermitePolynomialInterpolation.from_k
    )

    def order(self, terms):
        return 4

integrix.march.SSP_RK2

Bases: AbstractERK

Second-order SSP Runge–Kutta, the same scheme as RK2Heun.

Order 2, two stages, SSP coefficient C = 1, with an embedded forward-Euler error estimate.

Examples:

>>> SSP_RK2().order(None)
2

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

class SSP_RK2(dfx.AbstractERK):
    """Second-order SSP Runge–Kutta, the same scheme as `RK2Heun`.

    Order 2, two stages, SSP coefficient C = 1, with an embedded forward-Euler
    error estimate.

    Examples:
        >>> SSP_RK2().order(None)
        2
    """

    tableau: ClassVar[dfx.ButcherTableau] = dfx.ButcherTableau(
        c=np.array([1.0]),
        b_sol=np.array([0.5, 0.5]),
        # b_error = b_euler − b_sol, b_euler = [1, 0]
        b_error=np.array([0.5, -0.5]),
        a_lower=(np.array([1.0]),),
    )
    interpolation_cls: ClassVar[Any] = (
        dfx.ThirdOrderHermitePolynomialInterpolation.from_k
    )

    def order(self, terms):
        return 2

integrix.march.SSP_RK104

Bases: _FixedStepERK

Ketcheson's ten-stage, fourth-order SSP Runge–Kutta, fixed steps only.

Order 4, ten stages, SSP coefficient C = 6, the largest achievable at fourth order. There is no embedded error estimate, so an adaptive controller raises.

Examples:

>>> SSP_RK104().order(None)
4

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_erk.py

class SSP_RK104(_FixedStepERK):
    """Ketcheson's ten-stage, fourth-order SSP Runge–Kutta, fixed steps only.

    Order 4, ten stages, SSP coefficient C = 6, the largest achievable at
    fourth order. There is no embedded error estimate, so an adaptive
    controller raises.

    Examples:
        >>> SSP_RK104().order(None)
        4
    """

    tableau: ClassVar[dfx.ButcherTableau] = dfx.ButcherTableau(
        # c[i] = sum(a_lower[i]) for i = 0..8 (excludes the first stage, c = 0)
        c=np.array(
            [
                1.0 / 6,
                1.0 / 3,
                1.0 / 2,
                2.0 / 3,
                1.0 / 3,
                1.0 / 2,
                2.0 / 3,
                5.0 / 6,
                1.0,
            ]
        ),
        b_sol=np.array([1.0 / 10] * 10),
        b_error=np.zeros(10),
        a_lower=(
            # Rows 1-4: all entries 1/6
            np.array([1.0 / 6]),
            np.array([1.0 / 6, 1.0 / 6]),
            np.array([1.0 / 6, 1.0 / 6, 1.0 / 6]),
            np.array([1.0 / 6, 1.0 / 6, 1.0 / 6, 1.0 / 6]),
            # Rows 5-9: 1/15 for cols 0-4, then 1/6 for later cols
            # (convex combination at stage 6: 3/5*y0 + 2/5*y5 maps 1/6 -> 1/15)
            np.array([1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 15]),
            np.array([1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 6]),
            np.array(
                [1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 15, 1.0 / 6, 1.0 / 6]
            ),
            np.array([*[1.0 / 15] * 5, *[1.0 / 6] * 3]),
            np.array([*[1.0 / 15] * 5, *[1.0 / 6] * 4]),
        ),
    )
    interpolation_cls: ClassVar[Any] = (
        dfx.ThirdOrderHermitePolynomialInterpolation.from_k
    )

    def order(self, terms):
        return 4

integrix.march.IMEX_SSP2

Bases: AbstractRungeKutta, AbstractImplicitSolver

IMEX-SSP2(2,2,2) for a diffrax.MultiTerm of an explicit and an implicit term.

The explicit part is SSP with C = 1; the implicit part is A-stable SDIRK with γ = 1 − 1/√2. Order 2. The embedded first-order solution y + h (F_E(Y₁) + F_I(Y₁)) (b̂ = [1, 0] in both tableaux) gives the error estimate for adaptive step-size control. Dense output is linear between steps.

Explicit tableau: c = (0, 1), a₂₁ = 1, b = (½, ½). Implicit tableau: c = (γ, 1 − γ), a₁₁ = a₂₂ = γ, a₂₁ = 1 − 2γ, b = (½, ½).

With fixed steps, give the implicit root finder its own tolerances (IMEX_SSP2(root_finder=dfx.VeryChord(rtol=1e-8, atol=1e-8))); with a diffrax.PIDController, the default root finder takes the controller's.

Examples:

>>> import diffrax as dfx
>>> import jax.numpy as jnp
>>> terms = dfx.MultiTerm(
...     dfx.ODETerm(lambda t, y, args: -y),
...     dfx.ODETerm(lambda t, y, args: -10.0 * y),
... )
>>> solver = IMEX_SSP2(root_finder=dfx.VeryChord(rtol=1e-10, atol=1e-10))
>>> solution = dfx.diffeqsolve(
...     terms, solver, 0.0, 1.0, 0.01, jnp.array(1.0)
... )
>>> bool(jnp.allclose(solution.ys[0], jnp.exp(-11.0), rtol=1e-2))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_imex.py

class IMEX_SSP2(dfx.AbstractRungeKutta, dfx.AbstractImplicitSolver):
    """IMEX-SSP2(2,2,2) for a `diffrax.MultiTerm` of an explicit and an implicit term.

    The explicit part is SSP with C = 1; the implicit part is A-stable SDIRK
    with γ = 1 − 1/√2. Order 2. The embedded first-order solution
    y + h (F_E(Y₁) + F_I(Y₁)) (b̂ = [1, 0] in both tableaux) gives the error
    estimate for adaptive step-size control. Dense output is linear between
    steps.

    Explicit tableau: c = (0, 1), a₂₁ = 1, b = (½, ½). Implicit tableau:
    c = (γ, 1 − γ), a₁₁ = a₂₂ = γ, a₂₁ = 1 − 2γ, b = (½, ½).

    With fixed steps, give the implicit root finder its own tolerances
    (``IMEX_SSP2(root_finder=dfx.VeryChord(rtol=1e-8, atol=1e-8))``); with a
    `diffrax.PIDController`, the default root finder takes the controller's.

    Examples:
        >>> import diffrax as dfx
        >>> import jax.numpy as jnp
        >>> terms = dfx.MultiTerm(
        ...     dfx.ODETerm(lambda t, y, args: -y),
        ...     dfx.ODETerm(lambda t, y, args: -10.0 * y),
        ... )
        >>> solver = IMEX_SSP2(root_finder=dfx.VeryChord(rtol=1e-10, atol=1e-10))
        >>> solution = dfx.diffeqsolve(
        ...     terms, solver, 0.0, 1.0, 0.01, jnp.array(1.0)
        ... )
        >>> bool(jnp.allclose(solution.ys[0], jnp.exp(-11.0), rtol=1e-2))
        True
    """

    tableau: ClassVar[dfx.MultiButcherTableau] = dfx.MultiButcherTableau(
        # Explicit part
        dfx.ButcherTableau(
            c=np.array([1.0]),
            b_sol=np.array([0.5, 0.5]),
            # b_error = b_hat − b_sol, b_hat = [1, 0]
            b_error=np.array([0.5, -0.5]),
            a_lower=(np.array([1.0]),),
        ),
        # Implicit part (SDIRK): c₁ = γ, c₂ = (1 − 2γ) + γ = 1 − γ
        dfx.ButcherTableau(
            c=np.array([1.0 - _GAMMA]),
            b_sol=np.array([0.5, 0.5]),
            # b_error = b_hat − b_sol, b_hat = [1, 0] (same embedded weights)
            b_error=np.array([0.5, -0.5]),
            a_lower=(np.array([1.0 - 2.0 * _GAMMA]),),
            a_diagonal=np.array([_GAMMA, _GAMMA]),
            a_predictor=(np.array([1.0]),),
            c1=_GAMMA,
        ),
    )

    @staticmethod
    def interpolation_cls(t0, t1, y0, y1, k):
        # linear dense output; k (one stage tuple per term) is not needed
        del k
        return dfx.LocalLinearInterpolation(t0=t0, t1=t1, y0=y0, y1=y1)

    term_structure: ClassVar[Any] = dfx.MultiTerm[
        tuple[dfx.AbstractTerm, dfx.AbstractTerm]
    ]
    calculate_jacobian: ClassVar[Any] = dfx.CalculateJacobian.first_stage

    # diffrax declares these on AbstractImplicitSolver for subclasses to fill,
    # as its own KenCarp3 does; ty reads the declaration as a class variable.
    root_finder: Any = dfx.with_stepsize_controller_tols(dfx.VeryChord)()  # ty: ignore[invalid-attribute-override]
    root_find_max_steps: int = eqx.field(static=True, default=10)  # ty: ignore[invalid-attribute-override]

    def order(self, terms):
        return 2

integrix.march.AB3Solver

Bases: AbstractSolver

Third-order Adams–Bashforth as a diffrax solver, with an SSP-RK3 start.

yₙ₊₁ = yₙ + (h/12)(23 fₙ − 16 fₙ₋₁ + 5 fₙ₋₂), one evaluation of F per step. The two-level history lives in the diffrax solver state, and the first two steps are SSP-RK3 steps (a third-order bootstrap).

Fixed, uniform steps only (dt0=, or StepTo with equal spacing): the weights assume a constant step and there is no error estimate, so an adaptive controller raises. Order 3; dense output is linear.

Examples:

>>> import diffrax as dfx
>>> import jax.numpy as jnp
>>> solution = dfx.diffeqsolve(
...     dfx.ODETerm(lambda t, y, args: -y),
...     AB3Solver(),
...     t0=0.0,
...     t1=1.0,
...     dt0=0.01,
...     y0=jnp.array(1.0),
... )
>>> bool(jnp.allclose(solution.ys[0], jnp.exp(-1.0), rtol=1e-5))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

class AB3Solver(dfx.AbstractSolver):
    """Third-order Adams–Bashforth as a diffrax solver, with an SSP-RK3 start.

    yₙ₊₁ = yₙ + (h/12)(23 fₙ − 16 fₙ₋₁ + 5 fₙ₋₂), one evaluation of F per
    step. The two-level history lives in the diffrax solver state, and the
    first two steps are SSP-RK3 steps (a third-order bootstrap).

    Fixed, uniform steps only (``dt0=``, or ``StepTo`` with equal spacing):
    the weights assume a constant step and there is no error estimate, so an
    adaptive controller raises. Order 3; dense output is linear.

    Examples:
        >>> import diffrax as dfx
        >>> import jax.numpy as jnp
        >>> solution = dfx.diffeqsolve(
        ...     dfx.ODETerm(lambda t, y, args: -y),
        ...     AB3Solver(),
        ...     t0=0.0,
        ...     t1=1.0,
        ...     dt0=0.01,
        ...     y0=jnp.array(1.0),
        ... )
        >>> bool(jnp.allclose(solution.ys[0], jnp.exp(-1.0), rtol=1e-5))
        True
    """

    term_structure: ClassVar[Any] = dfx.AbstractTerm
    interpolation_cls: ClassVar[Any] = dfx.LocalLinearInterpolation

    def order(self, terms):
        return 3

    def init(self, terms, t0, t1, y0, args):
        # (fₙ₋₁, fₙ₋₂, number of steps taken)
        f0 = terms.vf(t0, y0, args)
        zeros = jax.tree.map(jnp.zeros_like, f0)
        return zeros, zeros, jnp.asarray(0, dtype=jnp.int32)

    def step(self, terms, t0, t1, y0, args, solver_state, made_jump):
        del made_jump
        f_nm1, f_nm2, n_taken = solver_state
        dt = t1 - t0
        f_n = terms.vf(t0, y0, args)

        def _rk3(_):
            # SSP-RK3 bootstrap (reuses fₙ as its first stage)
            y1 = jax.tree.map(lambda y, f: y + dt * f, y0, f_n)
            k2 = terms.vf(t1, y1, args)
            y2 = jax.tree.map(
                lambda y, a, b: 0.75 * y + 0.25 * (a + dt * b), y0, y1, k2
            )
            k3 = terms.vf(t0 + 0.5 * dt, y2, args)
            return jax.tree.map(
                lambda y, a, b: y / 3.0 + (2.0 / 3.0) * (a + dt * b), y0, y2, k3
            )

        def _ab3(_):
            # fₙ is already evaluated, so hand ab3_step a constant right-hand side
            y_next, _, _ = ab3_step(y0, lambda _y: f_n, dt, f_nm1, f_nm2)
            return y_next

        y1 = jax.lax.cond(n_taken < 2, _rk3, _ab3, None)
        dense_info = dict(y0=y0, y1=y1)
        new_state = (f_n, f_nm1, n_taken + 1)
        return y1, None, dense_info, new_state, RESULTS.successful

    def func(self, terms, t0, y0, args):
        return terms.vf(t0, y0, args)

Steppers Run by integrix.march.run

integrix.march.run(stepper, rhs, t0, y0, dt, n_steps, *, args=None, saveat=None, adjoint=None)

March y₀ through n_steps constant steps of h from t₀.

yₙ₊₁ = Φ(yₙ, tₙ, h) for n = 0, …, N − 1, with tₙ = t₀ + n h.

Algorithm

1. A diffrax solver goes to diffrax.diffeqsolve on the step times t₀ + n h (diffrax.StepTo), so it takes exactly N steps.
2. An AbstractStepper takes its first step with init and the other N − 1 with step inside jax.lax.scan, carrying the state and its history.

Parameters:

Name	Type	Description	Default
stepper	AbstractStepper | AbstractSolver	An AbstractStepper, or a diffrax solver (the steppers of integrix.march, or diffrax's own).	required
rhs	Callable[..., PyTree] | Any	The vector field rhs(t, y, args). A diffrax solver may also take a diffrax term (a diffrax.MultiTerm for IMEX, say); a split-explicit stepper takes its SplitExplicitTerms.	required
t0	ScalarLike	The start time t₀.	required
y0	PyTree	The initial state y₀, any pytree.	required
dt	ScalarLike	The step h.	required
n_steps	int	The number of steps N ≥ 1, a Python integer.	required
args	Any	Passed to the vector field.	None
saveat	SaveAt | None	None for the final state y_N, or SaveAt(steps=True) for every state y₁, …, y_N, stacked on a new leading axis.	None
adjoint	Any	How reverse mode runs back through the march. None is plain reverse mode (every state stored for a stepper; diffrax's default for a solver). integrix.adjoint.Checkpointed stores every k-th state and recomputes between (for a diffrax solver, RecursiveCheckpointAdjoint(checkpoints=⌈N / k⌉)). integrix.adjoint.Truncated cuts the gradient through the state before the last k steps, for a stepper only. A diffrax solver also takes a diffrax adjoint or a pipekit-cycle spec (integrix.adjoint.to_diffrax_adjoint). Forward values never depend on the adjoint.	None

Returns:

Type	Description
PyTree	y_N, or the stacked y₁, …, y_N.

Raises:

Type	Description
ValidationError	If n_steps is not a positive integer, saveat asks for anything but the final state or every step, adjoint does not apply to the stepper (Truncated for a diffrax solve, a diffrax adjoint for an AbstractStepper), or stepper is neither kind.

Examples:

>>> import jax.numpy as jnp
>>> from integrix.march import RK4Classic
>>> y = run(
...     RK4Classic(), lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.1, 10
... )
>>> bool(jnp.allclose(y, jnp.exp(-1.0), rtol=1e-6))
True
>>> from integrix.march import AB2Solver
>>> every = run(
...     AB2Solver(),
...     lambda t, y, args: -y,
...     0.0,
...     jnp.ones(3),
...     0.1,
...     5,
...     saveat=dfx.SaveAt(steps=True),
... )
>>> every.shape
(5, 3)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_loop.py

def run(
    stepper: AbstractStepper | dfx.AbstractSolver,
    rhs: Callable[..., PyTree] | Any,
    t0: ScalarLike,
    y0: PyTree,
    dt: ScalarLike,
    n_steps: int,
    *,
    args: Any = None,
    saveat: dfx.SaveAt | None = None,
    adjoint: Any = None,
) -> PyTree:
    """March y₀ through ``n_steps`` constant steps of h from t₀.

    yₙ₊₁ = Φ(yₙ, tₙ, h) for n = 0, …, N − 1, with tₙ = t₀ + n h.

    Algorithm:
        1. A diffrax solver goes to `diffrax.diffeqsolve` on the step times
           t₀ + n h (`diffrax.StepTo`), so it takes exactly N steps.
        2. An `AbstractStepper` takes its first step with ``init`` and the
           other N − 1 with ``step`` inside `jax.lax.scan`, carrying the
           state and its history.

    Args:
        stepper: An `AbstractStepper`, or a diffrax solver (the steppers of
            `integrix.march`, or diffrax's own).
        rhs: The vector field ``rhs(t, y, args)``. A diffrax solver may also
            take a diffrax term (a `diffrax.MultiTerm` for IMEX, say); a
            split-explicit stepper takes its `SplitExplicitTerms`.
        t0: The start time t₀.
        y0: The initial state y₀, any pytree.
        dt: The step h.
        n_steps: The number of steps N ≥ 1, a Python integer.
        args: Passed to the vector field.
        saveat: ``None`` for the final state y_N, or ``SaveAt(steps=True)``
            for every state y₁, …, y_N, stacked on a new leading axis.
        adjoint: How reverse mode runs back through the march. ``None`` is
            plain reverse mode (every state stored for a stepper; diffrax's
            default for a solver). `integrix.adjoint.Checkpointed` stores
            every k-th state and recomputes between (for a diffrax solver,
            ``RecursiveCheckpointAdjoint(checkpoints=⌈N / k⌉)``).
            `integrix.adjoint.Truncated` cuts the gradient through the state
            before the last k steps, for a stepper only. A diffrax solver
            also takes a diffrax adjoint or a pipekit-cycle spec
            (`integrix.adjoint.to_diffrax_adjoint`). Forward values never
            depend on the adjoint.

    Returns:
        y_N, or the stacked y₁, …, y_N.

    Raises:
        ValidationError: If ``n_steps`` is not a positive integer, ``saveat``
            asks for anything but the final state or every step, ``adjoint``
            does not apply to the stepper (`Truncated` for a diffrax solve, a
            diffrax adjoint for an `AbstractStepper`), or ``stepper`` is
            neither kind.

    Examples:
        >>> import jax.numpy as jnp
        >>> from integrix.march import RK4Classic
        >>> y = run(
        ...     RK4Classic(), lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.1, 10
        ... )
        >>> bool(jnp.allclose(y, jnp.exp(-1.0), rtol=1e-6))
        True
        >>> from integrix.march import AB2Solver
        >>> every = run(
        ...     AB2Solver(),
        ...     lambda t, y, args: -y,
        ...     0.0,
        ...     jnp.ones(3),
        ...     0.1,
        ...     5,
        ...     saveat=dfx.SaveAt(steps=True),
        ... )
        >>> every.shape
        (5, 3)
    """
    try:
        n_steps = operator.index(n_steps)
    except TypeError as err:
        raise ValidationError(
            f"'n_steps' must be a Python integer, got {type(n_steps).__name__}"
        ) from err
    if n_steps < 1:
        raise ValidationError(f"'n_steps' must be at least 1, got {n_steps}")
    save_steps = _save_steps(saveat)

    if isinstance(stepper, dfx.AbstractSolver):
        return _run_diffrax(
            stepper, rhs, t0, y0, dt, n_steps, args, save_steps, adjoint
        )
    if not isinstance(stepper, AbstractStepper):
        raise ValidationError(
            f"'stepper' must be an AbstractStepper or a diffrax solver, got "
            f"{type(stepper).__name__}"
        )
    scan = _loop_scan(adjoint, n_steps)

    time_dtype = state_time_dtype(y0)
    start = jnp.asarray(t0, dtype=time_dtype)
    step_size = jnp.asarray(dt, dtype=time_dtype)

    state, y1 = stepper.init(rhs, start, y0, step_size, args)
    # y₁ as saved: its own step's gradient survives any truncation
    y1_saved = y1
    if isinstance(adjoint, Truncated) and adjoint.k < n_steps:
        # the first step lies before the last k, so the carry it hands on is cut
        y1, state = jax.lax.stop_gradient((y1, state))

    def body(carry, index):
        y, history = carry
        # tₙ = t₀ + n h
        t = start + index.astype(time_dtype) * step_size
        y_next, history = stepper.step(rhs, t, y, step_size, history, args)
        return (y_next, history), (y_next if save_steps else None)

    indices = jnp.arange(1, n_steps)
    (y_final, _), later = scan(body, (y1, state), indices)
    if not save_steps:
        return y_final
    return jax.tree.map(
        lambda first, rest: jnp.concatenate([first[None], rest]), y1_saved, later
    )

integrix.march.AbstractStepper

Bases: Module

A fixed-step scheme whose history lives in the loop carry.

A stepper takes the diffrax vector field rhs(t, y, args), so a model writes its tendencies once for both kinds of stepper. init takes the first step, bootstrapping whatever history the scheme needs, and step takes every later one; neither mutates the module.

Examples:

>>> import jax.numpy as jnp
>>> class Euler(AbstractStepper):
...     @property
...     def order(self):
...         return 1
...
...     def init(self, rhs, t0, y0, dt, args):
...         return None, y0 + dt * rhs(t0, y0, args)
...
...     def step(self, rhs, t, y, dt, state, args):
...         return y + dt * rhs(t, y, args), state
>>> y = run(Euler(), lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.1, 10)
>>> bool(jnp.allclose(y, 0.9**10))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_loop.py

class AbstractStepper(eqx.Module):
    """A fixed-step scheme whose history lives in the loop carry.

    A stepper takes the diffrax vector field ``rhs(t, y, args)``, so a model
    writes its tendencies once for both kinds of stepper. ``init`` takes the
    first step, bootstrapping whatever history the scheme needs, and ``step``
    takes every later one; neither mutates the module.

    Examples:
        >>> import jax.numpy as jnp
        >>> class Euler(AbstractStepper):
        ...     @property
        ...     def order(self):
        ...         return 1
        ...
        ...     def init(self, rhs, t0, y0, dt, args):
        ...         return None, y0 + dt * rhs(t0, y0, args)
        ...
        ...     def step(self, rhs, t, y, dt, state, args):
        ...         return y + dt * rhs(t, y, args), state
        >>> y = run(Euler(), lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.1, 10)
        >>> bool(jnp.allclose(y, 0.9**10))
        True
    """

    @property
    @abc.abstractmethod
    def order(self) -> int:
        """The global order of accuracy p, error O(hᵖ)."""

    @abc.abstractmethod
    def init(
        self,
        rhs: Any,
        t0: ScalarLike,
        y0: PyTree,
        dt: ScalarLike,
        args: Any,
    ) -> tuple[PyTree, PyTree]:
        """Take the first step from (t₀, y₀): returns ``(state, y₁)``."""

    @abc.abstractmethod
    def step(
        self,
        rhs: Any,
        t: ScalarLike,
        y: PyTree,
        dt: ScalarLike,
        state: PyTree,
        args: Any,
    ) -> tuple[PyTree, PyTree]:
        """Take one step from (tₙ, yₙ) with the history: returns ``(yₙ₊₁, state)``."""

order abstractmethod property

The global order of accuracy p, error O(hᵖ).

init(rhs, t0, y0, dt, args) abstractmethod

Take the first step from (t₀, y₀): returns (state, y₁).

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_loop.py

@abc.abstractmethod
def init(
    self,
    rhs: Any,
    t0: ScalarLike,
    y0: PyTree,
    dt: ScalarLike,
    args: Any,
) -> tuple[PyTree, PyTree]:
    """Take the first step from (t₀, y₀): returns ``(state, y₁)``."""

step(rhs, t, y, dt, state, args) abstractmethod

Take one step from (tₙ, yₙ) with the history: returns (yₙ₊₁, state).

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_loop.py

@abc.abstractmethod
def step(
    self,
    rhs: Any,
    t: ScalarLike,
    y: PyTree,
    dt: ScalarLike,
    state: PyTree,
    args: Any,
) -> tuple[PyTree, PyTree]:
    """Take one step from (tₙ, yₙ) with the history: returns ``(yₙ₊₁, state)``."""

integrix.march.AB2Solver

Bases: AbstractStepper

Second-order Adams–Bashforth, run by integrix.march.run.

yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁), one evaluation of the vector field per step. init bootstraps with one Heun step and keeps f(t₀, y₀) as the history; fixed, uniform steps only (variable steps break the weights).

Examples:

>>> import jax.numpy as jnp
>>> from integrix.march import run
>>> y = run(
...     AB2Solver(), lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.01, 100
... )
>>> bool(jnp.allclose(y, jnp.exp(-1.0), rtol=1e-4))
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

class AB2Solver(AbstractStepper):
    """Second-order Adams–Bashforth, run by `integrix.march.run`.

    yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁), one evaluation of the vector field per
    step. ``init`` bootstraps with one Heun step and keeps f(t₀, y₀) as the
    history; fixed, uniform steps only (variable steps break the weights).

    Examples:
        >>> import jax.numpy as jnp
        >>> from integrix.march import run
        >>> y = run(
        ...     AB2Solver(), lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.01, 100
        ... )
        >>> bool(jnp.allclose(y, jnp.exp(-1.0), rtol=1e-4))
        True
    """

    @property
    def order(self) -> int:
        """Second order."""
        return 2

    def init(self, rhs, t0, y0, dt, args):
        """Heun to y₁; the history is fₙ₋₁ = f(t₀, y₀)."""
        f0, y1 = _heun_bootstrap(rhs, t0, y0, dt, args)
        return f0, y1

    def step(self, rhs, t, y, dt, state, args):
        """yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁); the history becomes fₙ."""
        f_curr = rhs(t, y, args)
        y_next = jax.tree.map(
            lambda yi, fi, fi_1: yi + (dt / 2.0) * (3.0 * fi - fi_1),
            y,
            f_curr,
            state,
        )
        return y_next, f_curr

order property

Second order.

init(rhs, t0, y0, dt, args)

Heun to y₁; the history is fₙ₋₁ = f(t₀, y₀).

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def init(self, rhs, t0, y0, dt, args):
    """Heun to y₁; the history is fₙ₋₁ = f(t₀, y₀)."""
    f0, y1 = _heun_bootstrap(rhs, t0, y0, dt, args)
    return f0, y1

step(rhs, t, y, dt, state, args)

yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁); the history becomes fₙ.

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def step(self, rhs, t, y, dt, state, args):
    """yₙ₊₁ = yₙ + (h/2)(3 fₙ − fₙ₋₁); the history becomes fₙ."""
    f_curr = rhs(t, y, args)
    y_next = jax.tree.map(
        lambda yi, fi, fi_1: yi + (dt / 2.0) * (3.0 * fi - fi_1),
        y,
        f_curr,
        state,
    )
    return y_next, f_curr

integrix.march.LeapfrogRAFSolver

Bases: AbstractStepper

Leapfrog with a Robert–Asselin filter, run by integrix.march.run.

ỹₙ₊₁ = ȳₙ₋₁ + 2h F(tₙ, yₙ) and ȳₙ = yₙ + α (ȳₙ₋₁ − 2yₙ + ỹₙ₊₁): the filter damps the computational mode, which grows without it. Second order for α = 0 and first order for α > 0 (the Robert–Asselin degradation). init bootstraps with one Heun step and keeps y₀ as the history; fixed, uniform steps only.

Attributes:

Name	Type	Description
alpha	float | Float[Array, '']	The filter coefficient α ≥ 0, typically 0.01 to 0.1. It is an array leaf, so it can be traced, vmapped over or differentiated; a Python or NumPy value is checked when the stepper is built.

Raises:

Type	Description
ValidationError	If a concrete alpha is negative or not finite.

Examples:

>>> import jax.numpy as jnp
>>> from integrix.march import run
>>> stepper = LeapfrogRAFSolver(alpha=0.05)
>>> y = run(
...     stepper, lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.001, 1000
... )
>>> bool(jnp.allclose(y, jnp.exp(-1.0), rtol=1e-3)), stepper.order
(True, 1)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

class LeapfrogRAFSolver(AbstractStepper):
    """Leapfrog with a Robert–Asselin filter, run by `integrix.march.run`.

    ỹₙ₊₁ = ȳₙ₋₁ + 2h F(tₙ, yₙ) and ȳₙ = yₙ + α (ȳₙ₋₁ − 2yₙ + ỹₙ₊₁): the filter
    damps the computational mode, which grows without it. Second order for
    α = 0 and first order for α > 0 (the Robert–Asselin degradation).
    ``init`` bootstraps with one Heun step and keeps y₀ as the history; fixed,
    uniform steps only.

    Attributes:
        alpha: The filter coefficient α ≥ 0, typically 0.01 to 0.1. It is an
            array leaf, so it can be traced, vmapped over or differentiated;
            a Python or NumPy value is checked when the stepper is built.

    Raises:
        ValidationError: If a concrete ``alpha`` is negative or not finite.

    Examples:
        >>> import jax.numpy as jnp
        >>> from integrix.march import run
        >>> stepper = LeapfrogRAFSolver(alpha=0.05)
        >>> y = run(
        ...     stepper, lambda t, y, args: -y, 0.0, jnp.array(1.0), 0.001, 1000
        ... )
        >>> bool(jnp.allclose(y, jnp.exp(-1.0), rtol=1e-3)), stepper.order
        (True, 1)
    """

    alpha: float | Float[Array, ""] = 0.05

    def __check_init__(self) -> None:
        if isinstance(self.alpha, jax.Array):
            return
        if not (math.isfinite(self.alpha) and self.alpha >= 0.0):
            raise ValidationError(
                f"'alpha' must be finite and non-negative, got {self.alpha}"
            )

    @property
    def order(self) -> int:
        """Second order without the filter, first order with it (or when traced)."""
        try:
            unfiltered = bool(self.alpha == 0.0)
        except jax.errors.ConcretizationTypeError:
            unfiltered = False
        return 2 if unfiltered else 1

    def init(self, rhs, t0, y0, dt, args):
        """Heun to y₁; the history is ȳ₀ = y₀."""
        _, y1 = _heun_bootstrap(rhs, t0, y0, dt, args)
        return y0, y1

    def step(self, rhs, t, y, dt, state, args):
        """Leapfrog from the filtered ȳₙ₋₁, then filter yₙ; the history becomes ȳₙ."""
        f_curr = rhs(t, y, args)
        # ỹₙ₊₁ = ȳₙ₋₁ + 2h F(tₙ, yₙ)
        y_next = jax.tree.map(lambda yp, fc: yp + 2.0 * dt * fc, state, f_curr)
        # ȳₙ = yₙ + α (ȳₙ₋₁ − 2yₙ + ỹₙ₊₁)
        y_filtered = jax.tree.map(
            lambda yp, yc, yn: yc + self.alpha * (yp - 2.0 * yc + yn),
            state,
            y,
            y_next,
        )
        return y_next, y_filtered

order property

Second order without the filter, first order with it (or when traced).

init(rhs, t0, y0, dt, args)

Heun to y₁; the history is ȳ₀ = y₀.

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def init(self, rhs, t0, y0, dt, args):
    """Heun to y₁; the history is ȳ₀ = y₀."""
    _, y1 = _heun_bootstrap(rhs, t0, y0, dt, args)
    return y0, y1

step(rhs, t, y, dt, state, args)

Leapfrog from the filtered ȳₙ₋₁, then filter yₙ; the history becomes ȳₙ.

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_multistep.py

def step(self, rhs, t, y, dt, state, args):
    """Leapfrog from the filtered ȳₙ₋₁, then filter yₙ; the history becomes ȳₙ."""
    f_curr = rhs(t, y, args)
    # ỹₙ₊₁ = ȳₙ₋₁ + 2h F(tₙ, yₙ)
    y_next = jax.tree.map(lambda yp, fc: yp + 2.0 * dt * fc, state, f_curr)
    # ȳₙ = yₙ + α (ȳₙ₋₁ − 2yₙ + ỹₙ₊₁)
    y_filtered = jax.tree.map(
        lambda yp, yc, yn: yc + self.alpha * (yp - 2.0 * yc + yn),
        state,
        y,
        y_next,
    )
    return y_next, y_filtered

integrix.march.SplitExplicitRKSolver

Bases: AbstractStepper

Split-explicit stepping as an AbstractStepper, for integrix.march.run.

The state is the pair y = (y₃, y₂) of the slow and fast modes, and rhs is a SplitExplicitTerms. Each step is split_explicit_step with the fast mode in M RK4 substeps; M follows the external-wave CFL limit. The scheme carries no history, so init is a plain first step. First order in the slow step.

Attributes:

Name	Type	Description
n_substeps	int	The number of fast substeps M ≥ 1 per slow step.

Raises:

Type	Description
ValidationError	If n_substeps is below 1.

Examples:

>>> import jax.numpy as jnp
>>> from integrix.march import run
>>> terms = SplitExplicitTerms(
...     slow=lambda t, y3, mean_2d, args: -mean_2d,
...     fast=lambda t, y2, y3, args: 0.0 * y2,
... )
>>> slow, fast = run(
...     SplitExplicitRKSolver(n_substeps=5),
...     terms,
...     0.0,
...     (jnp.array(1.0), jnp.array(1.0)),
...     0.1,
...     10,
... )
>>> bool(jnp.allclose(slow, 0.0, atol=1e-6)), bool(fast == 1.0)
(True, True)

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_split_explicit.py

class SplitExplicitRKSolver(AbstractStepper):
    """Split-explicit stepping as an `AbstractStepper`, for `integrix.march.run`.

    The state is the pair y = (y₃, y₂) of the slow and fast modes, and ``rhs``
    is a `SplitExplicitTerms`. Each step is `split_explicit_step` with the
    fast mode in M RK4 substeps; M follows the external-wave CFL limit. The
    scheme carries no history, so ``init`` is a plain first step. First order
    in the slow step.

    Attributes:
        n_substeps: The number of fast substeps M ≥ 1 per slow step.

    Raises:
        ValidationError: If ``n_substeps`` is below 1.

    Examples:
        >>> import jax.numpy as jnp
        >>> from integrix.march import run
        >>> terms = SplitExplicitTerms(
        ...     slow=lambda t, y3, mean_2d, args: -mean_2d,
        ...     fast=lambda t, y2, y3, args: 0.0 * y2,
        ... )
        >>> slow, fast = run(
        ...     SplitExplicitRKSolver(n_substeps=5),
        ...     terms,
        ...     0.0,
        ...     (jnp.array(1.0), jnp.array(1.0)),
        ...     0.1,
        ...     10,
        ... )
        >>> bool(jnp.allclose(slow, 0.0, atol=1e-6)), bool(fast == 1.0)
        (True, True)
    """

    n_substeps: int = eqx.field(static=True, default=50)

    def __check_init__(self) -> None:
        if isinstance(self.n_substeps, bool):
            raise ValidationError("'n_substeps' must be an integer, got a bool")
        try:
            count = operator.index(self.n_substeps)
        except TypeError as err:
            raise ValidationError(
                f"'n_substeps' must be an integer, got {self.n_substeps!r}"
            ) from err
        if count < 1:
            raise ValidationError(f"'n_substeps' must be at least 1, got {count}")

    @property
    def order(self) -> int:
        """First order in the slow step."""
        return 1

    def init(self, rhs, t0, y0, dt, args):
        """No history to bootstrap: the first step is an ordinary one."""
        y1, _ = self.step(rhs, t0, y0, dt, None, args)
        return None, y1

    def step(self, rhs, t, y, dt, state, args):
        """One `split_explicit_step` of the pair (y₃, y₂)."""
        terms = _as_terms(rhs)
        state_3d, state_2d = y
        slow_forcing = None
        if terms.slow_forcing is not None:
            forcing_fn = terms.slow_forcing

            def slow_forcing(y3: PyTree) -> PyTree:
                return forcing_fn(t, y3, args)

        couple = terms.couple

        def couple_2d_to_3d(y3: PyTree, mean_2d: PyTree) -> PyTree:
            return y3 if couple is None else couple(y3, mean_2d, args)

        new_3d, new_2d = split_explicit_step(
            state_3d,
            state_2d,
            lambda y3, mean_2d: terms.slow(t, y3, mean_2d, args),
            lambda t_sub, y2, y3: terms.fast(t_sub, y2, y3, args),
            couple_2d_to_3d,
            dt,
            self.n_substeps,
            slow_forcing=slow_forcing,
            t0=t,
        )
        return (new_3d, new_2d), state

order property

First order in the slow step.

init(rhs, t0, y0, dt, args)

No history to bootstrap: the first step is an ordinary one.

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_split_explicit.py

def init(self, rhs, t0, y0, dt, args):
    """No history to bootstrap: the first step is an ordinary one."""
    y1, _ = self.step(rhs, t0, y0, dt, None, args)
    return None, y1

step(rhs, t, y, dt, state, args)

One split_explicit_step of the pair (y₃, y₂).

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_split_explicit.py

def step(self, rhs, t, y, dt, state, args):
    """One `split_explicit_step` of the pair (y₃, y₂)."""
    terms = _as_terms(rhs)
    state_3d, state_2d = y
    slow_forcing = None
    if terms.slow_forcing is not None:
        forcing_fn = terms.slow_forcing

        def slow_forcing(y3: PyTree) -> PyTree:
            return forcing_fn(t, y3, args)

    couple = terms.couple

    def couple_2d_to_3d(y3: PyTree, mean_2d: PyTree) -> PyTree:
        return y3 if couple is None else couple(y3, mean_2d, args)

    new_3d, new_2d = split_explicit_step(
        state_3d,
        state_2d,
        lambda y3, mean_2d: terms.slow(t, y3, mean_2d, args),
        lambda t_sub, y2, y3: terms.fast(t_sub, y2, y3, args),
        couple_2d_to_3d,
        dt,
        self.n_substeps,
        slow_forcing=slow_forcing,
        t0=t,
    )
    return (new_3d, new_2d), state

integrix.march.SplitExplicitTerms

Bases: Module

The vector fields of a split-explicit model, passed to run as rhs.

Attributes:

Name	Type	Description
slow	Callable[..., PyTree]	The slow tendency slow(t, y_3d, mean_2d, args).
fast	Callable[..., PyTree]	The fast tendency fast(t, y_2d, y_3d, args), at the absolute RK4 stage time t.
slow_forcing	Callable[..., PyTree] | None	slow_forcing(t, y_3d, args), the baroclinic forcing of the barotropic mode, held fixed over the cycle; None for none.
couple	Callable[..., PyTree] | None	couple(y_3d, mean_2d, args), applied after the slow step; None leaves the slow state as it is.

Examples:

>>> terms = SplitExplicitTerms(
...     slow=lambda t, y3, mean_2d, args: -y3,
...     fast=lambda t, y2, y3, args: -y2,
... )
>>> terms.couple is None
True

Source code in .venv/lib/python3.12/site-packages/integrix/_src/march/_split_explicit.py

class SplitExplicitTerms(eqx.Module):
    """The vector fields of a split-explicit model, passed to `run` as ``rhs``.

    Attributes:
        slow: The slow tendency ``slow(t, y_3d, mean_2d, args)``.
        fast: The fast tendency ``fast(t, y_2d, y_3d, args)``, at the
            absolute RK4 stage time t.
        slow_forcing: ``slow_forcing(t, y_3d, args)``, the baroclinic forcing
            of the barotropic mode, held fixed over the cycle; ``None`` for
            none.
        couple: ``couple(y_3d, mean_2d, args)``, applied after the slow step;
            ``None`` leaves the slow state as it is.

    Examples:
        >>> terms = SplitExplicitTerms(
        ...     slow=lambda t, y3, mean_2d, args: -y3,
        ...     fast=lambda t, y2, y3, args: -y2,
        ... )
        >>> terms.couple is None
        True
    """

    slow: Callable[..., PyTree]
    fast: Callable[..., PyTree]
    slow_forcing: Callable[..., PyTree] | None = None
    couple: Callable[..., PyTree] | None = None

Semi-Lagrangian Solver

finitevolx.SemiLagrangianSolver

Bases: AbstractSolver

Semi-Lagrangian advection solver for diffrax.

Traces characteristic curves backward in time and interpolates the old field at departure points. Unconditionally stable (CFL > 1 allowed).

The term's vf(t, y, args) must return (u, v) velocity components in grid index units per second (i.e. physical velocity divided by grid spacing); wrap the velocity function in :class:SemiLagrangianTerm to use the solver with diffeqsolve. The departure point is found with one backward Euler step along a velocity frozen at t0, so the scheme is first order in time (:meth:order is 1); there is no error estimate, so it runs with fixed steps only. Dense output (SaveAt(ts=...)) is linear between steps.

Parameters:

Name	Type	Description	Default
interpolation_order	int	0 = nearest-neighbour, 1 = linear (diffusive, monotone). JAX currently only supports order <= 1.	required

Source code in finitevolx/_src/timestepping/semi_lagrangian.py

class SemiLagrangianSolver(dfx.AbstractSolver):
    """Semi-Lagrangian advection solver for diffrax.

    Traces characteristic curves backward in time and interpolates the old
    field at departure points.  Unconditionally stable (CFL > 1 allowed).

    The term's ``vf(t, y, args)`` must return ``(u, v)`` velocity components
    in **grid index units per second** (i.e. physical velocity divided by
    grid spacing); wrap the velocity function in :class:`SemiLagrangianTerm`
    to use the solver with ``diffeqsolve``.  The departure point is found
    with one backward Euler step along a velocity frozen at ``t0``, so the
    scheme is first order in time (:meth:`order` is 1); there is no error
    estimate, so it runs with fixed steps only.  Dense output (``SaveAt(ts=...)``)
    is linear between steps.

    Parameters
    ----------
    interpolation_order : int
        0 = nearest-neighbour, 1 = linear (diffusive, monotone).
        JAX currently only supports ``order <= 1``.
    """

    term_structure: ClassVar[Any] = dfx.AbstractTerm
    interpolation_cls: ClassVar[Any] = dfx.LocalLinearInterpolation
    interpolation_order: int = eqx.field(static=True, default=1)

    def order(self, terms):
        # first-order backward trajectory; the spatial interpolation order
        # (interpolation_order) is a separate, spatial accuracy
        return 1

    def init(self, terms, t0, t1, y0, args):
        return None

    def step(self, terms, t0, t1, y0, args, solver_state, made_jump):
        del solver_state, made_jump
        dt = t1 - t0

        # Velocity in grid-index units/second
        u, v = terms.vf(t0, y0, args)

        ny, nx = y0.shape
        y_coords, x_coords = jnp.meshgrid(
            jnp.arange(ny, dtype=y0.dtype),
            jnp.arange(nx, dtype=y0.dtype),
            indexing="ij",
        )

        x_dep = x_coords - u * dt
        y_dep = y_coords - v * dt

        y1 = jax.scipy.ndimage.map_coordinates(
            y0, [y_dep, x_dep], order=self.interpolation_order, mode="wrap"
        )

        dense_info = dict(y0=y0, y1=y1)
        return y1, None, dense_info, None, RESULTS.successful

    def func(self, terms, t0, y0, args):
        return terms.vf(t0, y0, args)

finitevolx.SemiLagrangianTerm

Bases: AbstractTerm

Velocity term for :class:SemiLagrangianSolver in diffeqsolve.

dfx.ODETerm requires its vector field to return a pytree shaped like y0, but semi-Lagrangian advection needs the velocity (u, v) of a scalar field y. This term carries that velocity function.

Parameters:

Name	Type	Description	Default
velocity	Callable[[float, Array, PyTree], tuple[Array, Array]]	velocity(t, y, args) -> (u, v), each shaped like y, in grid index units per unit time (physical velocity divided by the grid spacing).	required

Examples:

>>> term = SemiLagrangianTerm(
...     lambda t, y, args: (jnp.ones_like(y), jnp.zeros_like(y))
... )
>>> sol = dfx.diffeqsolve(term, SemiLagrangianSolver(), 0.0, 4.0, 1.0, y0)

Source code in finitevolx/_src/timestepping/semi_lagrangian.py

class SemiLagrangianTerm(dfx.AbstractTerm):
    """Velocity term for :class:`SemiLagrangianSolver` in ``diffeqsolve``.

    ``dfx.ODETerm`` requires its vector field to return a pytree shaped like
    ``y0``, but semi-Lagrangian advection needs the velocity ``(u, v)`` of a
    scalar field ``y``.  This term carries that velocity function.

    Parameters
    ----------
    velocity : Callable[[float, Array, PyTree], tuple[Array, Array]]
        ``velocity(t, y, args) -> (u, v)``, each shaped like ``y``, in
        **grid index units per unit time** (physical velocity divided by
        the grid spacing).

    Examples
    --------
    >>> term = SemiLagrangianTerm(
    ...     lambda t, y, args: (jnp.ones_like(y), jnp.zeros_like(y))
    ... )
    >>> sol = dfx.diffeqsolve(term, SemiLagrangianSolver(), 0.0, 4.0, 1.0, y0)
    """

    velocity: Callable

    def vf(self, t, y, args):
        return self.velocity(t, y, args)

    def contr(self, t0, t1, **kwargs):
        del kwargs
        return t1 - t0

    def prod(self, vf, control):
        return jax.tree.map(lambda v: v * control, vf)

Deprecated Adapters

The objects behind finitevolx.AB2Solver, finitevolx.LeapfrogRAFSolver and finitevolx.SplitExplicitRKSolver for one release.

finitevolx._src.timestepping._deprecated.AB2Solver

Bases: Module

Deprecated: use integrix.march.AB2Solver with integrix.march.run.

Adams-Bashforth 2nd order, y_{n+1} = y_n + (dt/2)(3 f_n - f_{n-1}), with f_prev held in the module. The arithmetic is integrix.march.AB2Solver's.

Usage::

solver = AB2Solver()
solver, y = solver.init(rhs_fn, t0, y0, dt)
for n in range(1, n_steps):
    y, solver = solver.step(rhs_fn, t0 + n * dt, y, dt)

Source code in finitevolx/_src/timestepping/_deprecated.py

class AB2Solver(eqx.Module):
    """Deprecated: use ``integrix.march.AB2Solver`` with ``integrix.march.run``.

    Adams-Bashforth 2nd order, ``y_{n+1} = y_n + (dt/2)(3 f_n - f_{n-1})``,
    with ``f_prev`` held in the module. The arithmetic is
    ``integrix.march.AB2Solver``'s.

    Usage::

        solver = AB2Solver()
        solver, y = solver.init(rhs_fn, t0, y0, dt)
        for n in range(1, n_steps):
            y, solver = solver.step(rhs_fn, t0 + n * dt, y, dt)
    """

    f_prev: PyTree | None = None

    def init(self, rhs_fn: Callable, t0: float, y0, dt: float):
        """Bootstrap with a Heun step, returning ``(updated_solver, y1)``.

        Stores ``f_prev = f(t0, y0)`` so that the first AB2 step uses the
        correct history.
        """
        f0, y1 = ixm.AB2Solver().init(_with_args(rhs_fn), t0, y0, dt, None)
        return eqx.tree_at(lambda s: s.f_prev, self, f0), y1

    def step(self, rhs_fn: Callable, t: float, y, dt: float):
        """AB2 step: ``y_{n+1} = y_n + (dt/2)(3 f_n - f_{n-1})``."""
        y_next, f_curr = ixm.AB2Solver().step(
            _with_args(rhs_fn), t, y, dt, self.f_prev, None
        )
        return y_next, eqx.tree_at(lambda s: s.f_prev, self, f_curr)

init(rhs_fn, t0, y0, dt)

Bootstrap with a Heun step, returning (updated_solver, y1).

Stores f_prev = f(t0, y0) so that the first AB2 step uses the correct history.

Source code in finitevolx/_src/timestepping/_deprecated.py

def init(self, rhs_fn: Callable, t0: float, y0, dt: float):
    """Bootstrap with a Heun step, returning ``(updated_solver, y1)``.

    Stores ``f_prev = f(t0, y0)`` so that the first AB2 step uses the
    correct history.
    """
    f0, y1 = ixm.AB2Solver().init(_with_args(rhs_fn), t0, y0, dt, None)
    return eqx.tree_at(lambda s: s.f_prev, self, f0), y1

step(rhs_fn, t, y, dt)

AB2 step: y_{n+1} = y_n + (dt/2)(3 f_n - f_{n-1}).

Source code in finitevolx/_src/timestepping/_deprecated.py

def step(self, rhs_fn: Callable, t: float, y, dt: float):
    """AB2 step: ``y_{n+1} = y_n + (dt/2)(3 f_n - f_{n-1})``."""
    y_next, f_curr = ixm.AB2Solver().step(
        _with_args(rhs_fn), t, y, dt, self.f_prev, None
    )
    return y_next, eqx.tree_at(lambda s: s.f_prev, self, f_curr)

finitevolx._src.timestepping._deprecated.LeapfrogRAFSolver

Bases: Module

Deprecated: use integrix.march.LeapfrogRAFSolver with integrix.march.run.

Leapfrog with a Robert-Asselin filter on the middle level, with the filtered y_prev held in the module. The arithmetic is integrix.march.LeapfrogRAFSolver's: first order for alpha > 0, second order for alpha = 0.

Usage::

solver = LeapfrogRAFSolver(alpha=0.05)
solver, y = solver.init(rhs_fn, t0, y0, dt)
for n in range(1, n_steps):
    y, solver = solver.step(rhs_fn, t0 + n * dt, y, dt)

Source code in finitevolx/_src/timestepping/_deprecated.py

class LeapfrogRAFSolver(eqx.Module):
    """Deprecated: use ``integrix.march.LeapfrogRAFSolver`` with ``integrix.march.run``.

    Leapfrog with a Robert-Asselin filter on the middle level, with the
    filtered ``y_prev`` held in the module. The arithmetic is
    ``integrix.march.LeapfrogRAFSolver``'s: first order for ``alpha > 0``,
    second order for ``alpha = 0``.

    Usage::

        solver = LeapfrogRAFSolver(alpha=0.05)
        solver, y = solver.init(rhs_fn, t0, y0, dt)
        for n in range(1, n_steps):
            y, solver = solver.step(rhs_fn, t0 + n * dt, y, dt)
    """

    alpha: float = 0.05
    y_prev: PyTree | None = None

    def init(self, rhs_fn: Callable, t0: float, y0, dt: float):
        """Bootstrap with a Heun step, returning ``(updated_solver, y1)``."""
        y_prev, y1 = ixm.LeapfrogRAFSolver(alpha=self.alpha).init(
            _with_args(rhs_fn), t0, y0, dt, None
        )
        return eqx.tree_at(lambda s: s.y_prev, self, y_prev), y1

    def step(self, rhs_fn: Callable, t: float, y_curr, dt: float):
        """Leapfrog + RAF step, returning ``(y_next, updated_solver)``."""
        y_next, y_filtered = ixm.LeapfrogRAFSolver(alpha=self.alpha).step(
            _with_args(rhs_fn), t, y_curr, dt, self.y_prev, None
        )
        return y_next, eqx.tree_at(lambda s: s.y_prev, self, y_filtered)

init(rhs_fn, t0, y0, dt)

Bootstrap with a Heun step, returning (updated_solver, y1).

Source code in finitevolx/_src/timestepping/_deprecated.py

def init(self, rhs_fn: Callable, t0: float, y0, dt: float):
    """Bootstrap with a Heun step, returning ``(updated_solver, y1)``."""
    y_prev, y1 = ixm.LeapfrogRAFSolver(alpha=self.alpha).init(
        _with_args(rhs_fn), t0, y0, dt, None
    )
    return eqx.tree_at(lambda s: s.y_prev, self, y_prev), y1

step(rhs_fn, t, y_curr, dt)

Leapfrog + RAF step, returning (y_next, updated_solver).

Source code in finitevolx/_src/timestepping/_deprecated.py

def step(self, rhs_fn: Callable, t: float, y_curr, dt: float):
    """Leapfrog + RAF step, returning ``(y_next, updated_solver)``."""
    y_next, y_filtered = ixm.LeapfrogRAFSolver(alpha=self.alpha).step(
        _with_args(rhs_fn), t, y_curr, dt, self.y_prev, None
    )
    return y_next, eqx.tree_at(lambda s: s.y_prev, self, y_filtered)

finitevolx._src.timestepping._deprecated.SplitExplicitRKSolver

Bases: Module

Deprecated: use integrix.march.SplitExplicitRKSolver with integrix.march.run.

Split-explicit barotropic/baroclinic step with finitevolX's old call shape. The arithmetic is integrix.march.SplitExplicitRKSolver's: RK4 substeps of the fast (2D) mode, its trapezoidal time mean over the cycle, and one Forward-Euler step of the slow (3D) mode with that mean.

Parameters:

Name	Type	Description	Default
n_substeps	int	Number of barotropic substeps per baroclinic step.	required

Source code in finitevolx/_src/timestepping/_deprecated.py

class SplitExplicitRKSolver(eqx.Module):
    """Deprecated: use ``integrix.march.SplitExplicitRKSolver`` with ``integrix.march.run``.

    Split-explicit barotropic/baroclinic step with finitevolX's old call
    shape. The arithmetic is ``integrix.march.SplitExplicitRKSolver``'s: RK4
    substeps of the fast (2D) mode, its trapezoidal time mean over the cycle,
    and one Forward-Euler step of the slow (3D) mode with that mean.

    Parameters
    ----------
    n_substeps : int
        Number of barotropic substeps per baroclinic step.
    """

    n_substeps: int = eqx.field(static=True, default=50)

    def step(
        self,
        rhs_slow: Callable,
        rhs_fast: Callable,
        t: float,
        y_3d,
        y_2d,
        dt_slow: float,
        slow_forcing: Callable | None = None,
        couple_2d_to_3d: Callable | None = None,
    ):
        """Take one split-explicit step.

        Parameters
        ----------
        rhs_slow : Callable[[float, PyTree, PyTree], PyTree]
            Slow RHS: ``rhs_slow(t, y_3d, y_2d_avg) -> tendency_3d``.
        rhs_fast : Callable[[float, PyTree, PyTree], PyTree]
            Fast RHS: ``rhs_fast(t_sub, y_2d, y_3d) -> tendency_2d``, with
            ``t_sub`` the absolute time of the RK4 stage.
        t : float
            Current time.
        y_3d, y_2d : PyTree
            3D (slow) and 2D (fast) states.
        dt_slow : float
            Slow timestep.
        slow_forcing : Callable[[PyTree], PyTree] or None, optional
            ``slow_forcing(y_3d) -> forcing_2d``, evaluated once and added to
            every fast tendency.
        couple_2d_to_3d : Callable[[PyTree, PyTree], PyTree] or None, optional
            ``couple_2d_to_3d(new_y_3d, y_2d_avg) -> corrected_y_3d``,
            applied after the slow step.

        Returns
        -------
        tuple[PyTree, PyTree]
            ``(new_y_3d, new_y_2d)``; ``new_y_2d`` is the end-of-cycle fast
            state at ``t + dt_slow``.
        """
        terms = ixm.SplitExplicitTerms(
            slow=lambda t_, y3, mean_2d, args: rhs_slow(t_, y3, mean_2d),
            fast=lambda t_, y2, y3, args: rhs_fast(t_, y2, y3),
            slow_forcing=(
                None
                if slow_forcing is None
                else (lambda t_, y3, args: slow_forcing(y3))
            ),
            couple=(
                None
                if couple_2d_to_3d is None
                else (lambda y3, mean_2d, args: couple_2d_to_3d(y3, mean_2d))
            ),
        )
        stepper = ixm.SplitExplicitRKSolver(n_substeps=self.n_substeps)
        (new_3d, new_2d), _ = stepper.step(terms, t, (y_3d, y_2d), dt_slow, None, None)
        return new_3d, new_2d

step(rhs_slow, rhs_fast, t, y_3d, y_2d, dt_slow, slow_forcing=None, couple_2d_to_3d=None)

Take one split-explicit step.

Parameters:

Name	Type	Description	Default
rhs_slow	Callable[[float, PyTree, PyTree], PyTree]	Slow RHS: rhs_slow(t, y_3d, y_2d_avg) -> tendency_3d.	required
rhs_fast	Callable[[float, PyTree, PyTree], PyTree]	Fast RHS: rhs_fast(t_sub, y_2d, y_3d) -> tendency_2d, with t_sub the absolute time of the RK4 stage.	required
t	float	Current time.	required
y_3d	PyTree	3D (slow) and 2D (fast) states.	required
y_2d	PyTree	3D (slow) and 2D (fast) states.	required
dt_slow	float	Slow timestep.	required
slow_forcing	Callable[[PyTree], PyTree] or None	slow_forcing(y_3d) -> forcing_2d, evaluated once and added to every fast tendency.	None
couple_2d_to_3d	Callable[[PyTree, PyTree], PyTree] or None	couple_2d_to_3d(new_y_3d, y_2d_avg) -> corrected_y_3d, applied after the slow step.	None

Returns:

Type	Description
tuple[PyTree, PyTree]	(new_y_3d, new_y_2d); new_y_2d is the end-of-cycle fast state at t + dt_slow.

Source code in finitevolx/_src/timestepping/_deprecated.py

def step(
    self,
    rhs_slow: Callable,
    rhs_fast: Callable,
    t: float,
    y_3d,
    y_2d,
    dt_slow: float,
    slow_forcing: Callable | None = None,
    couple_2d_to_3d: Callable | None = None,
):
    """Take one split-explicit step.

    Parameters
    ----------
    rhs_slow : Callable[[float, PyTree, PyTree], PyTree]
        Slow RHS: ``rhs_slow(t, y_3d, y_2d_avg) -> tendency_3d``.
    rhs_fast : Callable[[float, PyTree, PyTree], PyTree]
        Fast RHS: ``rhs_fast(t_sub, y_2d, y_3d) -> tendency_2d``, with
        ``t_sub`` the absolute time of the RK4 stage.
    t : float
        Current time.
    y_3d, y_2d : PyTree
        3D (slow) and 2D (fast) states.
    dt_slow : float
        Slow timestep.
    slow_forcing : Callable[[PyTree], PyTree] or None, optional
        ``slow_forcing(y_3d) -> forcing_2d``, evaluated once and added to
        every fast tendency.
    couple_2d_to_3d : Callable[[PyTree, PyTree], PyTree] or None, optional
        ``couple_2d_to_3d(new_y_3d, y_2d_avg) -> corrected_y_3d``,
        applied after the slow step.

    Returns
    -------
    tuple[PyTree, PyTree]
        ``(new_y_3d, new_y_2d)``; ``new_y_2d`` is the end-of-cycle fast
        state at ``t + dt_slow``.
    """
    terms = ixm.SplitExplicitTerms(
        slow=lambda t_, y3, mean_2d, args: rhs_slow(t_, y3, mean_2d),
        fast=lambda t_, y2, y3, args: rhs_fast(t_, y2, y3),
        slow_forcing=(
            None
            if slow_forcing is None
            else (lambda t_, y3, args: slow_forcing(y3))
        ),
        couple=(
            None
            if couple_2d_to_3d is None
            else (lambda y3, mean_2d, args: couple_2d_to_3d(y3, mean_2d))
        ),
    )
    stepper = ixm.SplitExplicitRKSolver(n_substeps=self.n_substeps)
    (new_3d, new_2d), _ = stepper.step(terms, t, (y_3d, y_2d), dt_slow, None, None)
    return new_3d, new_2d

Convenience Wrapper

finitevolx.solve_ocean_pde(rhs_fn, solver, y0, t0, t1, dt0, saveat=None, bc_fn=None, args=None, **kwargs)

Integrate an ocean PDE using diffrax.

Parameters:

Name	Type	Description	Default
rhs_fn	Callable	Right-hand side rhs_fn(t, y, args) -> dy/dt.	required
solver	AbstractSolver	Time integration scheme (e.g. integrix.march.RK3SSP(), integrix.march.RK4Classic()).	required
y0	PyTree	Initial condition.	required
t0	float	Start and end times.	required
t1	float	Start and end times.	required
dt0	float	Initial (or fixed) timestep.	required
saveat	SaveAt	Output saving specification. Defaults to saving the final state.	None
bc_fn	Callable	Boundary condition function bc_fn(dydt) -> dydt_corrected applied to the tendency after each RHS evaluation.	None
args	PyTree	Static arguments forwarded to rhs_fn.	None
**kwargs	Any	Forwarded to :func:diffrax.diffeqsolve (stepsize_controller, adjoint, max_steps, throw, event, ...). For a fixed-step run (no stepsize_controller) with concrete times and no max_steps, max_steps defaults to ceil((t1 - t0) / dt0) + 16: diffrax's own default (4096) stops a long fixed-step run early. Otherwise diffrax's default applies.	{}

Returns:

Type	Description
Solution	Solution object containing the saved states.

Source code in finitevolx/_src/timestepping/_solve.py

def solve_ocean_pde(
    rhs_fn: Callable,
    solver: dfx.AbstractSolver,
    y0: PyTree,
    t0: float,
    t1: float,
    dt0: float,
    saveat: dfx.SaveAt | None = None,
    bc_fn: Callable | None = None,
    args: PyTree = None,
    **kwargs: Any,
) -> dfx.Solution:
    """Integrate an ocean PDE using diffrax.

    Parameters
    ----------
    rhs_fn : Callable
        Right-hand side ``rhs_fn(t, y, args) -> dy/dt``.
    solver : diffrax.AbstractSolver
        Time integration scheme (e.g. ``integrix.march.RK3SSP()``,
        ``integrix.march.RK4Classic()``).
    y0 : PyTree
        Initial condition.
    t0, t1 : float
        Start and end times.
    dt0 : float
        Initial (or fixed) timestep.
    saveat : diffrax.SaveAt, optional
        Output saving specification.  Defaults to saving the final state.
    bc_fn : Callable, optional
        Boundary condition function ``bc_fn(dydt) -> dydt_corrected`` applied
        to the tendency after each RHS evaluation.
    args : PyTree, optional
        Static arguments forwarded to ``rhs_fn``.
    **kwargs
        Forwarded to :func:`diffrax.diffeqsolve` (``stepsize_controller``,
        ``adjoint``, ``max_steps``, ``throw``, ``event``, ...).  For a
        fixed-step run (no ``stepsize_controller``) with concrete times and
        no ``max_steps``, ``max_steps`` defaults to
        ``ceil((t1 - t0) / dt0) + 16``: diffrax's own default (4096) stops a
        long fixed-step run early.  Otherwise diffrax's default applies.

    Returns
    -------
    diffrax.Solution
        Solution object containing the saved states.
    """
    if bc_fn is not None:

        def rhs_with_bc(t, y, args_):
            dydt = rhs_fn(t, y, args_)
            return bc_fn(dydt)

        term = dfx.ODETerm(rhs_with_bc)
    else:
        term = dfx.ODETerm(rhs_fn)

    if "max_steps" not in kwargs and "stepsize_controller" not in kwargs:
        # fixed steps: exactly ceil((t1 - t0) / dt0) of them
        try:
            n_steps = math.ceil((float(t1) - float(t0)) / float(dt0))
        except (TypeError, jax.errors.ConcretizationTypeError):
            pass  # traced times: keep diffrax's default
        else:
            kwargs["max_steps"] = max(n_steps, 0) + 16

    return dfx.diffeqsolve(
        term,
        solver,
        t0=t0,
        t1=t1,
        dt0=dt0,
        y0=y0,
        args=args,
        saveat=saveat if saveat is not None else dfx.SaveAt(t1=True),
        **kwargs,
    )