From 157bb4755f2b2b4fb9c0d17bf5e14126abf132fd Mon Sep 17 00:00:00 2001 From: "Joshua C. Macdonald" Date: Tue, 13 Oct 2026 20:00:20 -0400 Subject: [PATCH 1/2] Safeguard steady-state continuation steps --- CHANGELOG.md | 6 ++ docs/guides/steady-states.md | 35 ++++++++- src/op_engine/steady_state.py | 137 +++++++++++++++++++++++++--------- 3 files changed, 141 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fc06162..91e6d3e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,12 @@ versioning while the project remains pre-1.0. in the initial state's array namespace, so native PyTorch tensors work with `CoreSolver.run`. CPU Torch CI covers fixed and adaptive explicit runs, float32/float64, and optional history (#208). +- `steady_state` rejects non-finite steps and excessive residual growth, + retaining the previous state and halving the pseudo-time step. Accepted + transient growth holds the step instead of repeatedly shrinking it. + Optional `nonnegative=True` or a boolean mask rejects negative candidates, + including the final Newton correction, without clipping conserved totals. + Seeded fast take-off is covered in eager and JAX jit/vmap solves (#209). ## [0.6.0] - 2026-10-09 diff --git a/docs/guides/steady-states.md b/docs/guides/steady-states.md index 0125236..8e633ae 100644 --- a/docs/guides/steady-states.md +++ b/docs/guides/steady-states.md @@ -127,15 +127,44 @@ to the returned state. Pseudo-transient continuation converges to *a* steady state: started next to an unstable one, such as the disease-free state, it can stop there. The `dt` -schedule holds the step while the residual grows moderately, as it does while -an epidemic takes off, so a small `dt0` (the default is 1 time unit) lets the -iterates follow the epidemic to the endemic state. Warm starts from a nearby +schedule holds the step through bounded residual growth, as it does while +an epidemic takes off. A small `dt0` (the default is 1 time unit) lets the +iterates follow the dynamics toward a stable equilibrium. Candidates with +non-finite states or residuals, or with scaled or unscaled residual growth +above `max_growth`, are rejected: the previous state and residual are retained, +and `dt` is halved for the next attempt. Rejections count toward +`max_iterations`; a smaller `dt0` or a short time integration can provide a +better starting point if the budget is insufficient. Warm starts from a nearby solution converge fastest; raise `dt0` for those. +For compartment populations or concentrations, pass `nonnegative=True` to +reject negative candidates, including the final Newton correction. A boolean +mask constrains selected states when the model also contains signed variables: + +```python +result = steady_state(rhs, y0, invariants=totals, nonnegative=True) +``` + +This guard is especially useful for fast take-off from a seeded disease-free +state, where a large pseudo-time step can overshoot into negative populations +or select the unstable equilibrium. Rejected steps preserve the invariant +totals; the solver does not clip states. Signs remain unrestricted by default. +The guard helps follow a physical trajectory but cannot guarantee convergence +to a particular equilibrium or avoid an iteration limit. + `SteadyStateConfig` collects the controls: `dt0`, `dt_max`, `min_growth`, `max_growth`, `max_iterations`, `residual_tol`, `step_tol`, and `atol`. The default tolerances assume double precision. +The safeguard retains pseudo-transient continuation for the coupled system. +[Brent's method](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brentq.html) +and [bisection](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.bisect.html) +need a scalar sign-changing bracket, which a general RHS does not provide. +The existing curvature-weighted grid utilities are spatial preprocessing tools. +For background on constraints and step control, see +[projected pseudo-transient continuation](https://ira.lib.polyu.edu.hk/bitstream/10397/6101/1/Kelley_Projected_Pseudotransient_Continuation.pdf) +and PETSc's [pseudo-time step verification](https://petsc.org/release/manualpages/TS/TSPseudoSetVerifyTimeStep/). + ## JAX: jit and vmap Pass an automatic-differentiation Jacobian and `loop=jax.lax.fori_loop`. The diff --git a/src/op_engine/steady_state.py b/src/op_engine/steady_state.py index 6e6fb56..4173b5a 100644 --- a/src/op_engine/steady_state.py +++ b/src/op_engine/steady_state.py @@ -64,13 +64,15 @@ class SteadyStateConfig: dt_max: Largest pseudo-time step. min_growth: Smallest factor by which ``dt`` grows when the unscaled residual did not increase. Without it, a slowly decaying mode - barely shrinks the residual and ``dt`` stalls. ``dt`` is held - while the residual grows by less than this factor per step, as - it does while an epidemic takes off, and shrinks beyond that. - max_growth: Largest factor by which ``dt`` grows, or shrinks, in one - iteration. Between the two bounds ``dt`` follows the ratio of - successive unscaled residual norms (switched evolution - relaxation). + barely shrinks the residual and ``dt`` stalls. + max_growth: Largest factor by which ``dt`` grows in one iteration. + While the unscaled residual decreases, ``dt`` follows its ratio + between the growth bounds (switched evolution relaxation). + Residual growth is allowed through a transient: ``dt`` is held + on accepted steps. Candidates whose scaled or unscaled residual + exceeds the current one by this factor are rejected, as are + non-finite or infeasible candidates. Rejection holds the state + and halves ``dt`` for the next attempt. max_iterations: Iteration budget; the loop is unrolled this many times when ``early_exit`` is false. residual_tol: Bound on the RMS of ``rhs / (atol + |y|)`` over the @@ -149,7 +151,8 @@ class SteadyStateResult: residual: ``rhs(t, state)``, zero at fixed states. converged: Whether both the residual and the Newton correction met their tolerances. - iterations: PTC iterations performed before convergence. + iterations: PTC attempts performed before convergence, including + rejected steps. residual_norm: Scaled residual RMS at ``state``. step_norm: Scaled RMS of the last Newton correction. dt: Final pseudo-time step. @@ -304,6 +307,21 @@ def _invariant_matrix(invariants: object, n: int, xp: Any, dtype: object) -> Any return xp.asarray(rows, dtype=dtype) +def _nonnegative_mask(nonnegative: object, n: int) -> np.ndarray: + """Return a host boolean mask of states constrained to be non-negative. + + Raises: + ValueError: If the input is not a boolean or a length-``n`` boolean mask. + """ + if isinstance(nonnegative, (bool, np.bool_)): + return np.full(n, nonnegative, dtype=np.bool_) + mask = np.asarray(nonnegative) + if mask.dtype != np.bool_ or mask.shape != (n,): + msg = f"nonnegative must be a boolean or a boolean mask of shape ({n},)" + raise ValueError(msg) + return mask + + def _scaled_rms(values: Any, scale: Any, free: Any, n_free: int, xp: Any) -> Any: # noqa: ANN401 """Return the RMS of ``values / scale`` over the unknowns. @@ -315,6 +333,16 @@ def _scaled_rms(values: Any, scale: Any, free: Any, n_free: int, xp: Any) -> Any return xp.sqrt(xp.divide(xp.sum(squared), n_free)) +def _solve_or_nan(matrix: Any, target: Any, xp: Any) -> Any: # noqa: ANN401 + """Return a dense solve, or a non-finite step for an eager singular system.""" + try: + return xp.linalg.solve(matrix, target) + except np.linalg.LinAlgError: + # NumPy raises for singular systems; traced namespaces generally + # return non-finite steps. Both paths use the same rejection test. + return xp.full_like(target, np.nan) + + def steady_state( # noqa: PLR0913, PLR0914, PLR0915 rhs: SteadyStateRhs, y0: Array, @@ -323,11 +351,19 @@ def steady_state( # noqa: PLR0913, PLR0914, PLR0915 jacobian: SteadyStateJacobian | None = None, fixed: Sequence[int] | Array | None = None, invariants: Array | Sequence[Sequence[float]] | None = None, + nonnegative: bool | Sequence[bool] | Array = False, config: SteadyStateConfig | None = None, loop: SteadyStateLoop | None = None, ) -> SteadyStateResult: """Find a steady state of ``rhs`` near ``y0`` by pseudo-transient continuation. + Non-finite candidates and excessive residual growth are rejected; + the state is held and the pseudo-time step is halved. Use + ``nonnegative=True`` for compartment or concentration states, especially + during fast take-off from a seeded unstable equilibrium. Rejection + preserves invariant totals without clipping. A short time integration + can supply a warm start if the iteration budget is insufficient. + Args: rhs: ``rhs(t, y) -> dy/dt`` for a one-dimensional state ``y``, such as ``lambda t, y: compiled.eval_fn(t, y, **params)``. @@ -342,6 +378,10 @@ def steady_state( # noqa: PLR0913, PLR0914, PLR0915 ``W @ y`` the dynamics conserve; held at ``W @ y0``. They are setup data, checked on the host, so they must be concrete values (not traced inside ``jax.jit``). + nonnegative: ``True`` constrains every state to be non-negative; + a boolean mask constrains selected states. ``False`` leaves + signs unrestricted. Like ``fixed``, this is concrete setup data. + The final Newton correction is also checked against the bounds. config: Iteration controls. loop: Optional ``loop(lower, upper, body, carry)`` driver for the iterations, such as ``jax.lax.fori_loop``. Under ``jax.jit`` it @@ -370,6 +410,7 @@ def steady_state( # noqa: PLR0913, PLR0914, PLR0915 free_host = _free_mask(fixed, n) n_free = int(free_host.sum()) free = xp.asarray(free_host) + bounded = xp.asarray(_nonnegative_mask(nonnegative, n)) w = _invariant_matrix(invariants, n, xp, dtype) k = 0 if w is None else int(w.shape[0]) totals = None if w is None else xp.matmul(w, y0) @@ -377,6 +418,7 @@ def steady_state( # noqa: PLR0913, PLR0914, PLR0915 eye = xp.eye(n, dtype=dtype) free_pair = xp.logical_and(xp.reshape(free, (n, 1)), xp.reshape(free, (1, n))) inv_dt_max = xp.asarray(1.0 / config.dt_max, dtype=dtype) + dt_min = xp.asarray(xp.finfo(dtype).tiny, dtype=dtype) def masked_rhs(y: Any) -> Any: # noqa: ANN401 return xp.where(free, rhs(t, y), xp.zeros_like(y)) @@ -388,17 +430,17 @@ def step(y: Any, value: Any, jac: Any, inv_dt: Any) -> Any: # noqa: ANN401 The step ``delta`` for the state. """ matrix = xp.where(free_pair, xp.subtract(xp.multiply(eye, inv_dt), jac), eye) - if w is None: - return xp.linalg.solve(matrix, value) - bordered = xp.concat( - ( - xp.concat((matrix, xp.matrix_transpose(w)), axis=1), - xp.concat((w, xp.zeros((k, k), dtype=dtype)), axis=1), - ), - axis=0, - ) - target = xp.concat((value, xp.subtract(totals, xp.matmul(w, y)))) - return xp.linalg.solve(bordered, target)[:n] + target = value + if w is not None: + matrix = xp.concat( + ( + xp.concat((matrix, xp.matrix_transpose(w)), axis=1), + xp.concat((w, xp.zeros((k, k), dtype=dtype)), axis=1), + ), + axis=0, + ) + target = xp.concat((value, xp.subtract(totals, xp.matmul(w, y)))) + return _solve_or_nan(matrix, target, xp)[:n] value0 = masked_rhs(y0) scale0 = xp.add(config.atol, xp.abs(y0)) @@ -439,9 +481,33 @@ def iterate(_index: int, carry: _Carry) -> _Carry: # noqa: PLR0914 ) # A converged iterate takes its Newton correction as a final polish. candidate = xp.add(y, xp.where(done_now, newton, ptc)) - candidate_value = masked_rhs(candidate) + feasible = xp.all(xp.logical_or(xp.logical_not(bounded), candidate >= 0)) + admissible = xp.logical_and(xp.all(xp.isfinite(candidate)), feasible) + # RHS callbacks can have a restricted domain. An inadmissible state + # is rejected without evaluating the callback outside that domain. + candidate_value = masked_rhs(xp.where(admissible, candidate, y)) candidate_scale = xp.add(config.atol, xp.abs(candidate)) + candidate_norm = _scaled_rms(candidate_value, candidate_scale, free, n_free, xp) + finite = xp.logical_and( + xp.all(xp.isfinite(candidate)), + xp.logical_and( + xp.all(xp.isfinite(candidate_value)), xp.isfinite(candidate_norm) + ), + ) candidate_raw = xp.sqrt(xp.mean(xp.multiply(candidate_value, candidate_value))) + bounded_growth = xp.logical_and( + xp.less_equal( + candidate_norm, + xp.maximum( + xp.multiply(config.max_growth, residual_norm), config.residual_tol + ), + ), + xp.less_equal(candidate_raw, xp.multiply(config.max_growth, raw_norm)), + ) + accepted = xp.logical_and( + xp.logical_and(finite, xp.isfinite(candidate_raw)), + xp.logical_and(feasible, bounded_growth), + ) positive = xp.greater(candidate_raw, 0) ratio = xp.where( positive, @@ -449,26 +515,29 @@ def iterate(_index: int, carry: _Carry) -> _Carry: # noqa: PLR0914 xp.divide(raw_norm, xp.where(positive, candidate_raw, 1.0)), xp.asarray(config.max_growth, dtype=dtype), ) - # Grow while the residual falls; hold through moderate rises, which - # ordinary dynamics produce (an epidemic taking off); shrink only when - # the residual jumps by more than min_growth in one step. + # Grow while the residual falls; hold during transient growth. + # An unacceptable candidate never becomes the next state. factor = xp.where( xp.greater_equal(ratio, 1.0), xp.clip(ratio, config.min_growth, config.max_growth), - xp.where( - xp.greater_equal(ratio, 1.0 / config.min_growth), - xp.ones_like(ratio), - xp.maximum(ratio, 1.0 / config.max_growth), - ), + xp.ones_like(ratio), + ) + done_now = xp.logical_and( + xp.logical_and(done_now, accepted), + xp.less_equal(candidate_norm, config.residual_tol), ) updated: _Carry = ( - candidate, - candidate_value, - candidate_scale, - _scaled_rms(candidate_value, candidate_scale, free, n_free, xp), - candidate_raw, + xp.where(accepted, candidate, y), + xp.where(accepted, candidate_value, value), + xp.where(accepted, candidate_scale, scale), + xp.where(accepted, candidate_norm, residual_norm), + xp.where(accepted, candidate_raw, raw_norm), newton_norm, - xp.minimum(xp.multiply(dt, factor), config.dt_max), + xp.where( + accepted, + xp.minimum(xp.multiply(dt, factor), config.dt_max), + xp.maximum(xp.multiply(dt, 0.5), dt_min), + ), xp.logical_or(converged, done_now), xp.add(iterations, xp.ones_like(iterations)), ) From 5aa8ba5533c9be1da3108b633f3a3addd22d4918 Mon Sep 17 00:00:00 2001 From: "Joshua C. Macdonald" Date: Tue, 13 Oct 2026 20:00:20 -0400 Subject: [PATCH 2/2] Cover steady-state rejection and fast take-off --- tests/op_engine/test_steady_state.py | 201 +++++++++++++++++++++++++++ 1 file changed, 201 insertions(+) diff --git a/tests/op_engine/test_steady_state.py b/tests/op_engine/test_steady_state.py index 14884cf..b89ef08 100644 --- a/tests/op_engine/test_steady_state.py +++ b/tests/op_engine/test_steady_state.py @@ -213,6 +213,180 @@ def test_config_rejects_invalid_controls(field: str, value: float) -> None: SteadyStateConfig(**{field: value}) +def _vector_host(b: Any, xp: Any = np) -> Any: # noqa: ANN401 + """Return the conserved vector-host SIRS reproduction from issue #209.""" + + def rhs(_t: float, y: Any) -> Any: # noqa: ANN401 + s, i, r, sv, iv = y + mu, gamma, omega, mu_v = 1 / 3650, 1 / 7, 1 / 300, 1 / 700 + lam_h, lam_v = b * iv / 1000, b * i / 1000 + return xp.stack([ + mu * 1000 - lam_h * s - mu * s + omega * r, + lam_h * s - (gamma + mu) * i, + gamma * i - (omega + mu) * r, + mu_v * 100_000 - lam_v * sv - mu_v * sv, + lam_v * sv - mu_v * iv, + ]) + + return rhs + + +def _vector_host_equilibrium(b: float) -> NDArray[np.float64]: + """Return the analytic endemic root, independently of the PTC iteration.""" + mu, gamma, omega, mu_v = 1 / 3650, 1 / 7, 1 / 300, 1 / 700 + recovery_ratio = gamma / (omega + mu) + i = (1000 - (gamma + mu) * 1000**2 * mu_v / (b**2 * 100_000)) / ( + 1 + recovery_ratio + (gamma + mu) * 1000 / (b * 100_000) + ) + denominator = 1000 * mu_v + b * i + iv = 100_000 * b * i / denominator + sv = 100_000 * 1000 * mu_v / denominator + s = (gamma + mu) * i * 1000 / (b * iv) + return np.array([s, i, recovery_ratio * i, sv, iv]) + + +VECTOR_HOST_INITIAL = np.array([1000.0, 0.0, 0.0, 99_990.0, 10.0]) +VECTOR_HOST_TOTALS = np.array([[1.0, 1, 1, 0, 0], [0.0, 0, 0, 1, 1]]) + + +@pytest.mark.parametrize("b", [0.02, 0.05, 0.2]) +def test_fast_takeoff_reaches_endemic_equilibrium(b: float) -> None: + """Seeded take-off converges within the default budget without losing mass.""" + result = require_steady_state( + steady_state( + _vector_host(b), + VECTOR_HOST_INITIAL, + invariants=VECTOR_HOST_TOTALS, + nonnegative=True, + ) + ) + np.testing.assert_allclose(result.state, _vector_host_equilibrium(b), rtol=1e-10) + assert np.all(result.state >= 0) + assert float(result.invariant_drift) < 1e-14 + assert int(result.iterations) < 100 + + +def test_moderate_takeoff_does_not_need_a_nonnegativity_constraint() -> None: + """Ordinary transient growth no longer shrinks the step after acceptance.""" + result = require_steady_state( + steady_state( + _vector_host(0.05), VECTOR_HOST_INITIAL, invariants=VECTOR_HOST_TOTALS + ) + ) + np.testing.assert_allclose(result.state, _vector_host_equilibrium(0.05), rtol=1e-10) + + +def test_excessive_residual_growth_rejects_without_changing_state() -> None: + """A finite overshoot is held; a smaller next attempt can recover.""" + result = steady_state( + lambda _t, y: 1 - y**3, + np.zeros(1), + config=SteadyStateConfig(dt0=10, max_iterations=1), + ) + np.testing.assert_array_equal(result.state, [0.0]) + np.testing.assert_array_equal(result.residual, [1.0]) + assert not bool(result.converged) + assert float(result.dt) == 5 + recovered = require_steady_state( + steady_state( + lambda _t, y: 1 - y**3, np.zeros(1), config=SteadyStateConfig(dt0=10) + ) + ) + np.testing.assert_allclose(recovered.state, [1.0]) + + +@pytest.mark.parametrize("bad_value", [np.inf, np.nan]) +def test_nonfinite_candidate_residual_is_rejected(bad_value: float) -> None: + """Non-finite residuals never replace a valid state or its diagnostics.""" + result = steady_state( + lambda _t, y: np.where(y <= 0.5, 1.0, bad_value), + np.zeros(1), + jacobian=lambda _t, _y: np.zeros((1, 1)), + config=SteadyStateConfig(max_iterations=1), + ) + np.testing.assert_array_equal(result.state, [0.0]) + np.testing.assert_array_equal(result.residual, [1.0]) + assert np.isfinite(result.residual_norm) + assert float(result.dt) == 0.5 + assert not bool(result.converged) + + +def test_singular_ptc_step_is_rejected_and_retried() -> None: + """At a singular shift NumPy raises; halving the step restores the solve.""" + rhs = lambda _t, y: y * (1 - y) # noqa: E731 + jacobian = lambda _t, y: np.diag(1 - 2 * y) # noqa: E731 + config = SteadyStateConfig(dt0=1.25, max_iterations=1) + y0 = np.array([0.1]) + result = steady_state(rhs, y0, jacobian=jacobian, config=config) + np.testing.assert_array_equal(result.state, y0) + assert float(result.dt) == 0.625 + recovered = require_steady_state( + steady_state( + rhs, + y0, + jacobian=jacobian, + nonnegative=True, + config=SteadyStateConfig(dt0=1.25), + ) + ) + np.testing.assert_allclose(recovered.state, [1.0]) + + +def test_nonnegative_mask_allows_unconstrained_signed_states() -> None: + """Only selected entries are constrained, and fixed states stay fixed.""" + result = require_steady_state( + steady_state( + lambda _t, y: np.array([1.0, -2.0, 0.0]) - y, + np.array([0.0, 0.0, -3.0]), + fixed=[2], + nonnegative=[True, False, False], + ) + ) + np.testing.assert_allclose(result.state, [1.0, -2.0, -3.0]) + + +def test_nonnegative_rejection_keeps_rhs_evaluations_in_bounds() -> None: + """A callback with a non-negative domain never sees a rejected overshoot.""" + + def rhs(_t: float, y: NDArray[np.float64]) -> NDArray[np.float64]: + assert np.all(np.isfinite(y)) + assert np.all(y >= 0) + return y * (1 - y) + + y0 = np.array([0.1]) + result = steady_state( + rhs, + y0, + jacobian=lambda _t, y: np.diag(1 - 2 * y), + nonnegative=True, + config=SteadyStateConfig(dt0=2, max_iterations=1), + ) + np.testing.assert_array_equal(result.state, y0) + assert not bool(result.converged) + assert float(result.dt) == 1 + + +@pytest.mark.parametrize("mask", [[True], [0, 1, 0, 0], np.ones((2, 2), dtype=bool)]) +def test_nonnegative_rejects_malformed_masks(mask: Any) -> None: # noqa: ANN401 + """Bound masks must have one boolean for each state.""" + with pytest.raises(ValueError, match="nonnegative"): + steady_state(_sirs, Y0, nonnegative=mask) + + +def test_nonnegative_constraint_checks_the_final_newton_correction() -> None: + """A small residual and correction cannot mark an infeasible root converged.""" + result = steady_state( + lambda _t, y: -y - 1e-12, + np.zeros(1), + jacobian=lambda _t, _y: -np.eye(1), + nonnegative=True, + config=SteadyStateConfig(residual_tol=1, step_tol=1, max_iterations=1), + ) + assert not bool(result.converged) + np.testing.assert_array_equal(result.state, [0.0]) + assert float(result.dt) == 0.5 + + def _enable_x64() -> contextlib.AbstractContextManager[None]: jax = pytest.importorskip("jax") enable = getattr(jax, "enable_x64", None) @@ -221,6 +395,33 @@ def _enable_x64() -> contextlib.AbstractContextManager[None]: return enable() +def test_fast_takeoff_with_rejection_under_jit_and_vmap() -> None: + """Array-valued accept/reject decisions remain valid in a batched JAX loop.""" + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + with _enable_x64(): + + def solve(b: Any) -> tuple[Any, Any, Any]: # noqa: ANN401 + rhs = _vector_host(b, jnp) + result = steady_state( + rhs, + jnp.asarray(VECTOR_HOST_INITIAL), + jacobian=lambda t, y: jax.jacfwd(lambda z: rhs(t, z))(y), + invariants=VECTOR_HOST_TOTALS, + nonnegative=True, + loop=jax.lax.fori_loop, + ) + return result.state, result.converged, result.invariant_drift + + parameters = np.array([0.02, 0.05, 0.2]) + states, converged, drifts = jax.jit(jax.vmap(solve))(jnp.asarray(parameters)) + assert bool(jnp.all(converged)) + assert bool(jnp.all(states >= 0)) + assert bool(jnp.all(drifts < 1e-14)) + for b, state in zip(parameters, np.asarray(states), strict=True): + np.testing.assert_allclose(state, _vector_host_equilibrium(b), rtol=1e-10) + + def test_jit_and_vmap_over_parameters_with_fori_loop() -> None: """A compiled, batched solve matches the analytic equilibria.""" jax = pytest.importorskip("jax")