BoundedSolver

experimental.network_dynamics.solvers.BoundedSolver(
    base_solver,
    low=-jnp.inf,
    high=jnp.inf,
)

Wrapper that enforces hard bounds on solver output via clipping.

Wraps any native solver and clips the output state to specified bounds. Useful for ensuring states remain in valid ranges (e.g., firing rates ≥ 0).

The bounds support flexible broadcasting: - Scalar: same bound for all states/nodes - [n_states]: different bounds per state variable - [n_states, n_nodes]: different bounds per state per node - State-matching PyTree: different broadcastable bounds for each state leaf

Use -jnp.inf or jnp.inf to disable clipping for specific states/nodes.

Args: base_solver: The underlying solver to wrap low: Lower bound(s) for state clipping (default: -inf, no clipping) high: Upper bound(s) for state clipping (default: +inf, no clipping)

Example: # Ensure all states stay in [0, 1] solver = BoundedSolver(Euler(), low=0.0, high=1.0)

# Different bounds per state variable
solver = BoundedSolver(
    Heun(),
    low=jnp.array([0.0, -5.0]),  # state 0: ≥0, state 1: ≥-5
    high=jnp.array([1.0, 5.0])   # state 0: ≤1, state 1: ≤5
)

# Different bounds for heterogeneous state groups
solver = BoundedSolver(
    Heun(),
    low=Bunch(cortex=jnp.array([[0.0]]), relay=-jnp.inf),
    high=Bunch(cortex=jnp.array([[1.0]]), relay=jnp.inf),
)

Methods

Name	Description
step	Integration step with state clipping.

step

experimental.network_dynamics.solvers.BoundedSolver.step(
    dynamics_fn,
    t,
    state,
    dt,
    params,
    noise_sample=0.0,
)

Integration step with state clipping.

Delegates to base solver, then clips output state to bounds. Input state is not clipped (already clipped from previous iteration).

Args: dynamics_fn: Dynamics function (t, state, params) -> (derivatives, auxiliaries) t: Current time state: Current state array or PyTree of state arrays dt: Time step params: Parameters noise_sample: Pre-scaled scalar, array, or matching PyTree increment

Returns: Tuple of (next_state, auxiliaries): - next_state: Clipped state array or PyTree - auxiliaries: Auxiliary-variable array or PyTree