Safeguard steady-state steps during fast transient growth

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
 

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

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,13 +410,15 @@ 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)
 
     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,36 +481,63 @@ 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,
             # Both branches are evaluated, so keep the denominator nonzero.
             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)),
         )