tsdynamics 2.4.0__tar.gz → 2.5.0__tar.gz

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tsdynamics
-Version: 2.4.0
+Version: 2.5.0
 Summary: Compiled ODE/DDE integration, discrete maps, and chaos analysis for dynamical systems.
 Project-URL: Source, https://github.com/El3ssar/TSDynamics
 Project-URL: Documentation, https://el3ssar.github.io/TSDynamics/

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/src/tsdynamics/__init__.py

@@ -48,7 +48,7 @@ from .systems import continuous as _continuous
 from .systems import discrete as _discrete
 
 # Single source of truth for the package version; rewritten by python-semantic-release.
-__version__ = "2.4.0"
+__version__ = "2.5.0"
 
 # Re-export every system class at the top level so ``from tsdynamics import Lorenz``
 # works without users having to remember which submodule a system lives in.

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/src/tsdynamics/backends/rustcore.py

@@ -12,15 +12,17 @@ optional accelerator is installed::
 
     pip install tsdynamics-core    # or the [rustcore] extra
 
-Limitations (experimental): ODEs only; fixed-step explicit RK4 (no stiff
-support — use the ``jitcode``/``diffsol`` backends for stiff systems); the RHS
-must use functions the tape VM provides (the same set the DiffSL backend
-supports).  Unsupported constructs raise :class:`TapeCompileError`.
+Solvers (select with ``method=``): fixed-step ``RK4``, adaptive Dormand-Prince
+``RK45`` (default), and an L-stable linearly-implicit ``stiff`` kernel (aliases
+``Rosenbrock``/``LSODA``/``BDF``) that uses the system's analytic Jacobian —
+also lowered into the tape. Limitations (experimental): ODEs only; the RHS must
+use functions the tape VM provides (the same set the DiffSL backend supports).
+Unsupported constructs raise :class:`TapeCompileError`.
 """
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import numpy as np
@@ -101,6 +103,9 @@ class CompiledTape:
     n_state: int
     n_param: int
     control_names: list[str]
+    # Register index of each Jacobian entry, row-major dim×dim; empty unless the
+    # tape was compiled with_jacobian=True (only the stiff solver needs it).
+    jac_outputs: np.ndarray = field(default_factory=lambda: np.empty(0, dtype=np.int32))
 
 
 class _Emitter:
@@ -191,18 +196,24 @@ class _Emitter:
         return self._push(OP_POW, a=base_reg, b=self.emit(exp))
 
 
-def compile_tape(system: Any) -> CompiledTape:
+def compile_tape(system: Any, *, with_jacobian: bool = False) -> CompiledTape:
     """
     Lower a :class:`~tsdynamics.base.ContinuousSystem`'s RHS to a tape.
 
     Structural parameters are folded to constants; control parameters become
     inputs in ``control_names`` order (the layout the solve-time params vector
     must follow).
+
+    With ``with_jacobian=True`` the analytic Jacobian ``∂f_k/∂u_j`` is
+    differentiated symbolically and emitted into the *same* tape (sharing
+    common subexpressions with the RHS) — the stiff solver consumes it.
     """
     import symengine
     from jitcode import t as t_sym
     from jitcode import y
 
+    from tsdynamics.base.ode_base import _resolve_derivative_nodes
+
     dim = system.dim
     struct_vals = system._structural_vals()
     control_names = list(system._control_params())
@@ -212,11 +223,22 @@ def compile_tape(system: Any) -> CompiledTape:
     if len(exprs) != dim:
         raise ValueError(f"_equations must return {dim} expressions, got {len(exprs)}")
 
-    subs = {y(i): symengine.Symbol(f"u{i}") for i in range(dim)}
+    u_syms = [symengine.Symbol(f"u{i}") for i in range(dim)]
+    subs = {y(i): u_syms[i] for i in range(dim)}
     subs[t_sym] = symengine.Symbol("t")
+    rhs = [symengine.sympify(e).subs(subs) for e in exprs]
 
     em = _Emitter()
-    outputs = [em.emit(symengine.sympify(e).subs(subs)._sympy_()) for e in exprs]
+    outputs = [em.emit(e._sympy_()) for e in rhs]
+    jac_outputs: list[int] = []
+    if with_jacobian:
+        # Row-major dim×dim: ∂f_k/∂u_j. _resolve_derivative_nodes rewrites the
+        # unevaluated d|u|/du → sign(u) and d·sign/du → 0 a.e. (the same a.e.
+        # convention the jitcode/diffsol Jacobian autogen uses), so systems with
+        # abs/sign still lower cleanly.
+        for e in rhs:
+            for s in u_syms:
+                jac_outputs.append(em.emit(_resolve_derivative_nodes(e.diff(s))._sympy_()))
 
     return CompiledTape(
         ops=np.asarray(em.ops, dtype=np.int32),
@@ -227,6 +249,7 @@ def compile_tape(system: Any) -> CompiledTape:
         n_state=dim,
         n_param=len(control_names),
         control_names=control_names,
+        jac_outputs=np.asarray(jac_outputs, dtype=np.int32),
     )
 
 
@@ -253,23 +276,57 @@ def _args(tape: CompiledTape):
     return (tape.ops, tape.a, tape.b, tape.imm, tape.outputs, tape.n_state, tape.n_param)
 
 
-# Method names: fixed-step RK4, or the adaptive Dormand-Prince 5(4) kernel.
-# These are the only kernels the crate implements; an unknown name is rejected
-# (rather than silently downgraded) — there is no DOP853/LSODA kernel to reach.
+def _args_stiff(tape: CompiledTape):
+    # Stiff kernels take jac_outputs after outputs (mirrors the Rust signature).
+    return (
+        tape.ops,
+        tape.a,
+        tape.b,
+        tape.imm,
+        tape.outputs,
+        tape.jac_outputs,
+        tape.n_state,
+        tape.n_param,
+    )
+
+
+# Method names → kernel. RK4 = fixed step; RK45/dopri5 = adaptive Dormand-Prince
+# 5(4); stiff = L-stable linearly-implicit solver with the analytic Jacobian.
+# The SciPy-style stiff names (LSODA/BDF/Radau) route to the stiff kernel so
+# `method="LSODA"` does the expected thing rather than silently downgrading.
 _FIXED_RK4 = {"RK4", "rk4"}
 _ADAPTIVE = {"RK45", "rk45", "dopri5", "DP45"}
-_METHODS = _FIXED_RK4 | _ADAPTIVE
+_STIFF = {
+    "stiff",
+    "Rosenbrock",
+    "rosenbrock",
+    "ROS1",
+    "LSODA",
+    "lsoda",
+    "BDF",
+    "bdf",
+    "Radau",
+    "radau",
+}
+_METHODS = _FIXED_RK4 | _ADAPTIVE | _STIFF
 
 
 def _check_method(method: str) -> None:
     if method not in _METHODS:
         raise ValueError(
             f"rustcore: unknown method {method!r}; supported: {sorted(_METHODS)} "
-            "(RK4 = fixed step, RK45/dopri5 = adaptive). For stiff systems use the "
-            "jitcode/diffsol backends."
+            "(RK4 = fixed step, RK45/dopri5 = adaptive explicit, "
+            "stiff/Rosenbrock/LSODA/BDF = L-stable implicit)."
         )
 
 
+def _stiff_tape(system: Any, tape: CompiledTape | None) -> CompiledTape:
+    """Return a tape carrying the analytic Jacobian (recompiled if the given one lacks it)."""
+    if tape is not None and tape.jac_outputs.size:
+        return tape
+    return compile_tape(system, with_jacobian=True)
+
+
 def _as_state(x: Any, dim: int, name: str) -> np.ndarray:
     """Coerce to a contiguous float64 1-D vector of length ``dim`` (or raise)."""
     a = np.ascontiguousarray(x, dtype=np.float64).ravel()
@@ -303,14 +360,16 @@ def integrate_dense(
 
     ``method="RK45"`` (default) uses error-controlled adaptive Dormand-Prince
     5(4) with Hermite dense output (``rtol``/``atol`` set the tolerance);
-    ``method="RK4"`` uses fixed-step RK4 with internal step ``h``.
+    ``method="RK4"`` uses fixed-step RK4 with internal step ``h``;
+    ``method="stiff"`` (aliases ``Rosenbrock``/``LSODA``/``BDF``) uses an
+    L-stable linearly-implicit solver with the system's analytic Jacobian.
 
     Raises ``RuntimeError`` if the trajectory diverges or the step collapses
     before reaching the final time (matching the jitcode/diffsol backends).
     """
     _check_method(method)
     core = _require()
-    tape = tape or compile_tape(system)
+    tape = _stiff_tape(system, tape) if method in _STIFF else (tape or compile_tape(system))
     t_eval = np.ascontiguousarray(t_eval, dtype=np.float64)
     ic = _as_state(ic, tape.n_state, "ic")
     p = _params_vec(system, tape)
@@ -318,6 +377,12 @@ def integrate_dense(
         if h is None:
             h = float(np.min(np.diff(t_eval))) if t_eval.size > 1 else 1e-2
         y = np.asarray(core.integrate_dense_py(*_args(tape), ic, p, t_eval, float(h)))
+    elif method in _STIFF:
+        y = np.asarray(
+            core.integrate_dense_stiff_py(
+                *_args_stiff(tape), ic, p, t_eval, float(rtol), float(atol)
+            )
+        )
     else:
         y = np.asarray(
             core.integrate_dense_rk45_py(*_args(tape), ic, p, t_eval, float(rtol), float(atol))
@@ -346,8 +411,9 @@ def ensemble_final(
 
     ``u0_batch`` is ``(n, dim)``; each row is integrated from ``t0`` to ``t1``
     and its final state returned as a row of the ``(n, dim)`` result.
-    ``method`` selects the adaptive (``"RK45"``, default) or fixed-step
-    (``"RK4"``) kernel, as in :func:`integrate_dense`.
+    ``method`` selects the adaptive (``"RK45"``, default), fixed-step
+    (``"RK4"``), or stiff (``"stiff"``/``"LSODA"``) kernel, as in
+    :func:`integrate_dense`.
 
     A trajectory that diverges (escapes to infinity) yields a row of ``NaN``
     rather than aborting the batch — the basin/ensemble caller classifies those
@@ -355,7 +421,7 @@ def ensemble_final(
     """
     _check_method(method)
     core = _require()
-    tape = tape or compile_tape(system)
+    tape = _stiff_tape(system, tape) if method in _STIFF else (tape or compile_tape(system))
     u0_batch = np.ascontiguousarray(u0_batch, dtype=np.float64)
     if u0_batch.ndim != 2 or u0_batch.shape[1] != tape.n_state:
         raise ValueError(
@@ -368,6 +434,12 @@ def ensemble_final(
                 *_args(tape), u0_batch, p, float(t0), float(t1), float(h)
             )
         )
+    if method in _STIFF:
+        return np.asarray(
+            core.integrate_ensemble_final_stiff_py(
+                *_args_stiff(tape), u0_batch, p, float(t0), float(t1), float(rtol), float(atol)
+            )
+        )
     return np.asarray(
         core.integrate_ensemble_final_rk45_py(
             *_args(tape), u0_batch, p, float(t0), float(t1), float(rtol), float(atol)

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/src/tsdynamics/base/ode_base.py

@@ -153,6 +153,14 @@ class ContinuousSystem(SystemBase, ABC):
 
     _default_method: ClassVar[str] = "RK45"
 
+    #: Whether JiTCODE should run SymEngine's ``simplify(ratio=1)`` on each RHS
+    #: expression before emitting C.  ``None`` keeps JiTCODE's own default
+    #: (enabled for ``dim <= 10``).  Set ``False`` on systems whose RHS is a
+    #: large rational expression: the simplify pass is super-linear and can
+    #: effectively hang at compile time, while the C compiler optimises the
+    #: unsimplified code just as well (see ``BlinkingRotlet``).
+    _compile_simplify: ClassVar[bool | None] = None
+
     #: Parameters whose values affect the symbolic *structure* of _equations
     #: (e.g. integer loop bounds). These are baked in at compile time.
     _structural_params: ClassVar[frozenset[str]] = frozenset()
@@ -390,7 +398,7 @@ class ContinuousSystem(SystemBase, ABC):
         ode = cls_jitc(
             f_sym, n=self.dim, control_pars=control_par_list, verbose=False, **lyap_kwargs
         )
-        ode.generate_f_C()
+        ode.generate_f_C(simplify=type(self)._compile_simplify)
         so = pathlib.Path(ode.save_compiled(destination=str(dest), overwrite=True))
 
         if for_lyap:

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/src/tsdynamics/systems/continuous/climate_geophysics.py

@@ -75,6 +75,28 @@ class DoubleGyre(ContinuousSystem):
 
 
 class BlinkingRotlet(ContinuousSystem):
+    r"""Blinking-rotlet flow — a model of chaotic advection in Stokes mixing.
+
+    A passive tracer in a circular cell stirred by two off-centre rotlets at
+    radius ``+b`` and ``-b`` that alternate being active with period ``tau``,
+    the switch being a steep ``tanh`` ramp.  The periodic blinking folds and
+    stretches material lines, producing chaotic advection (Aref-type *blinking*
+    flow).  State is ``(r, theta, t)`` in polar coordinates with an explicit
+    clock; ``dtheta`` carries a ``1/r`` factor singular at the cell centre, so
+    ``default_ic`` sits on a bounded orbit away from ``r = 0``.
+
+    Compile note
+    ------------
+    The rotlet velocity field is a large rational expression.  JiTCODE's
+    default ``simplify(ratio=1)`` codegen pass is super-linear in expression
+    size and effectively hangs while *compiling* this RHS — which is what made
+    the system look un-integrable.  ``_compile_simplify = False`` skips that
+    pass (the C compiler optimises the unsimplified code), after which the flow
+    integrates in well under a millisecond at every tolerance, even at the
+    original steep ``tanh`` switch.
+    """
+
+    reference = "Meleshko & Aref (1996), Phys. Fluids 8, 3215-3217"
     params = {
         "a": 1.0,
         "b": 0.5298833894399929,
@@ -83,6 +105,15 @@ class BlinkingRotlet(ContinuousSystem):
         "tau": 3.0,
     }
     dim = 3
+    variables = ("r", "theta", "t")
+    #: On a bounded orbit (``r`` stays in ``[0.8, 0.91]``) away from the
+    #: ``r = 0`` angular singularity, so default sweeps stay bounded.
+    default_ic = [0.8, 4.887, 0.0]
+
+    # The rotlet RHS is a large rational expression; JiTCODE's default
+    # simplify(ratio=1) pass is super-linear on it and effectively hangs the
+    # compile.  Skip it — the C compiler optimises the unsimplified code.
+    _compile_simplify = False
 
     @staticmethod
     def _rotlet(r, theta, a, b, bc):

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/src/tsdynamics/systems/continuous/physical_systems.py

@@ -140,13 +140,52 @@ class JerkCircuit(ContinuousSystem):
 
 
 class WindmiReduced(ContinuousSystem):
+    r"""Reduced WINDMI model of the solar-wind–magnetosphere–ionosphere coupling.
+
+    Three-variable reduction (field-aligned current ``i``, voltage ``v``,
+    pressure ``p``) of the WINDMI energy-conserving model.  The pressure-loss
+    term carries a fractional power ``p**(5/4)`` gated by a steep
+    ``tanh(d1*(i-1))`` switch (``d1 = 2200`` makes it a near-step at ``i = 1``,
+    the substorm-onset threshold).
+
+    Numerical guards (no change to the dynamics)
+    --------------------------------------------
+    Two robustness guards let every backend reproduce the same orbit that
+    JiTCODE's explicit dopri5 already handles:
+
+    * ``abs(p)`` under the fractional power.  On the attractor ``p >= 0``, but
+      adaptive and implicit solvers probe trial states with ``p < 0`` where the
+      real ``p**(5/4)`` is complex.  ``abs(p)`` is identical to ``p`` on the
+      physical trajectory and real, ``C^1`` everywhere.
+
+    * The ``tanh`` argument is clamped to ``[-_TANH_CLAMP, _TANH_CLAMP]``.
+      ``tanh`` saturates to ``±1`` to machine precision past ``|arg| ≈ 20``, so
+      clamping leaves the switch value unchanged, but it stops the diffsol Rust
+      kernels' automatic differentiation from overflowing on ``exp`` of an
+      enormous ``d1*(i-1)`` (≈ ±1500 on the orbit) — the actual cause of their
+      Newton/error-test failures.  The clamp uses ``abs`` only, which both the
+      JiTCODE and DiffSL paths support.
+    """
+
+    reference = "Horton, Weigel & Sprott (2001), Phys. Plasmas 8, 2946-2952"
     params = {"a1": 0.247, "b1": 10.8, "b2": 0.0752, "b3": 1.06, "d1": 2200, "vsw": 5}
     dim = 3
+    variables = ("i", "v", "p")
+
+    #: Bound on the ``tanh`` switch argument; ``tanh(±20)`` already equals ``±1``
+    #: to ~1e-18, so clamping is dynamically invisible but autodiff-safe.
+    _TANH_CLAMP = 25.0
 
     @staticmethod
     def _equations(Y, t, *, a1, b1, b2, b3, d1, vsw):
         i, v, p = Y(0), Y(1), Y(2)
+        # Clamp d1*(i-1) to [-C, C] via the abs identity
+        # clamp(z, -C, C) = (|z + C| - |z - C|) / 2 — smooth-enough (abs only)
+        # and translatable by both the JiTCODE and DiffSL backends.
+        c = WindmiReduced._TANH_CLAMP
+        z = d1 * (i - 1)
+        z_clamped = (abs(z + c) - abs(z - c)) / 2
         idot = a1 * (vsw - v)
         vdot = b1 * i - b2 * p**1 / 2 - b3 * v
-        pdot = vsw**2 - p ** (5 / 4) * vsw ** (1 / 2) * (1 + tanh(d1 * (i - 1))) / 2
+        pdot = vsw**2 - abs(p) ** (5 / 4) * vsw ** (1 / 2) * (1 + tanh(z_clamped)) / 2
         return idot, vdot, pdot

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/tests/_sampling.py

@@ -56,25 +56,20 @@ INTEGRATION_SAMPLE: list[str] = [
 ]
 
 # --- ODE systems excluded from the exhaustive integration sweeps ------------
-# These cannot be integrated by adaptive solvers in bounded time, so they
-# would hang the nightly full sweep rather than fail it.
-HARD_TO_INTEGRATE: dict[str, str] = {
-    "BlinkingRotlet": (
-        "near-discontinuous blinking protocol (tanh of a steep sine) makes every "
-        "adaptive solver take vanishingly small steps and stall at any tolerance; "
-        "needs a dedicated integration recipe — tracked separately"
-    ),
-}
+# Systems that cannot be integrated by adaptive solvers in bounded time would
+# hang the nightly full sweep rather than fail it, so they are skipped here.
+# Currently empty: ``BlinkingRotlet`` (the former entry) only *looked*
+# un-integrable because JiTCODE's simplify(ratio=1) codegen pass hangs on its
+# large rational RHS — fixed with ``_compile_simplify = False`` on the class.
+HARD_TO_INTEGRATE: dict[str, str] = {}
 
 # Systems the diffsol backend specifically can't integrate (the default JiTCODE
 # backend handles them fine). Excluded from the diffsol full-catalogue sweep.
-DIFFSOL_SKIP: dict[str, str] = {
-    "WindmiReduced": (
-        "RHS has a near-discontinuous tanh(2200*(i-1)) switch and fractional "
-        "powers p**(5/4) that go complex when the solver probes p<0; diffsol's "
-        "Rust solvers fail to converge while JiTCODE's dopri5 tolerates it"
-    ),
-}
+# Currently empty: ``WindmiReduced`` (the former entry) now integrates on every
+# diffsol solver after guarding the fractional power with ``abs(p)`` and
+# clamping the steep ``tanh`` argument (autodiff-safe) — both no-ops on the
+# physical trajectory.
+DIFFSOL_SKIP: dict[str, str] = {}
 
 # --- Maps excluded from Lyapunov-spectrum shape tests -----------------------
 MAP_LYAPUNOV_EXCLUDE: dict[str, str] = {

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/tests/test_diffsol_backend.py

@@ -135,10 +135,10 @@ def test_diffsol_integrates_full_catalogue() -> None:
     (``-m full``) with the diffsol extra.
 
     Numeric agreement with JiTCODE is checked separately on the curated sample
-    (``test_cross_validation_over_sample``).  We deliberately do NOT cross-check
-    the JiTCODE reference here for the whole catalogue: a couple of systems
-    (see ``HARD_TO_INTEGRATE``) stall every adaptive solver and can't produce a
-    bounded reference, which would hang the sweep rather than measure diffsol.
+    (``test_cross_validation_over_sample``); here we only assert each diffsol
+    trajectory is finite.  ``HARD_TO_INTEGRATE | DIFFSOL_SKIP`` is honoured as
+    an escape hatch for systems that genuinely cannot be integrated, but both
+    sets are currently empty — the whole ODE catalogue integrates on diffsol.
     """
     import sys
     import zlib

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/tests/test_rustcore.py

@@ -47,6 +47,40 @@ class _Twin(ContinuousSystem):
         return [y(0), y(0)]
 
 
+class _LinStiff(ContinuousSystem):
+    """y' = -k y — large k stresses L-stability (exact decay to ~0)."""
+
+    params = {"k": 1000.0}
+    dim = 1
+
+    @staticmethod
+    def _equations(y, t, *, k):
+        return [-k * y(0)]
+
+
+class _VdPStiff(ContinuousSystem):
+    """Stiff van der Pol oscillator (μ = 1000)."""
+
+    params = {"mu": 1000.0}
+    dim = 2
+
+    @staticmethod
+    def _equations(y, t, *, mu):
+        return [y(1), mu * (1 - y(0) ** 2) * y(1) - y(0)]
+
+
+class _Robertson(ContinuousSystem):
+    """Robertson's stiff chemical-kinetics problem."""
+
+    params: dict = {}
+    dim = 3
+
+    @staticmethod
+    def _equations(y, t):
+        a, b, c = y(0), y(1), y(2)
+        return [-0.04 * a + 1e4 * b * c, 0.04 * a - 1e4 * b * c - 3e7 * b * b, 3e7 * b * b]
+
+
 @pytest.mark.parametrize("name", _EVAL_SAMPLE)
 def test_tape_eval_matches_symbolic_rhs(name: str) -> None:
     """The Rust tape evaluates the exact same RHS as the symbolic core."""
@@ -214,6 +248,84 @@ def test_empty_t_eval_returns_empty(method: str) -> None:
     assert y.shape[0] == 0
 
 
+# ---------------------------------------------------------------------------
+# Stiff solver — L-stable linearly-implicit kernel with the analytic Jacobian.
+# ---------------------------------------------------------------------------
+
+
+def test_stiff_is_l_stable() -> None:
+    """y' = -1000 y integrates accurately to ~0 (an explicit method would blow up)."""
+    sys = _LinStiff()
+    t_eval = np.linspace(0.0, 1.0, 101)
+    y = rc.integrate_dense(sys, [1.0], t_eval, method="stiff", rtol=1e-7, atol=1e-9)
+    assert np.all(np.isfinite(y))
+    np.testing.assert_allclose(y[:, 0], np.exp(-1000.0 * t_eval), atol=1e-6)
+
+
+@pytest.mark.parametrize(
+    ("sys_factory", "ic", "T", "ref_method"),
+    [
+        (_VdPStiff, [2.0, 0.0], 30.0, "Radau"),
+        (_Robertson, [1.0, 0.0, 0.0], 40.0, "Radau"),
+    ],
+)
+def test_stiff_matches_scipy(sys_factory, ic, T, ref_method) -> None:
+    """The stiff kernel reproduces SciPy's implicit solvers on classic stiff problems."""
+    from scipy.integrate import solve_ivp
+
+    sys = sys_factory()
+    t_eval = np.linspace(0.0, T, 200)
+    y = rc.integrate_dense(sys, ic, t_eval, method="stiff", rtol=1e-7, atol=1e-10)
+    assert np.all(np.isfinite(y))
+    f = sys._rhs_numeric()
+    ref = solve_ivp(
+        lambda t, u: f(u, t), (0.0, T), ic, t_eval=t_eval, rtol=1e-11, atol=1e-12, method=ref_method
+    ).y.T
+    np.testing.assert_allclose(y, ref, rtol=1e-3, atol=1e-6)
+
+
+def test_stiff_handles_abs_sign_systems() -> None:
+    """A piecewise-linear system (Chua) integrates via the resolved-sign Jacobian."""
+    from scipy.integrate import solve_ivp
+
+    sys = ts.Chua()
+    ic = sys.resolve_ic(0.1 * np.ones(sys.dim))
+    t_eval = np.linspace(0.0, 4.0, 201)
+    y = rc.integrate_dense(sys, ic, t_eval, method="stiff", rtol=1e-7, atol=1e-9)
+    assert np.all(np.isfinite(y))
+    f = sys._rhs_numeric()
+    ref = solve_ivp(
+        lambda t, u: f(u, t), (0.0, 4.0), ic, t_eval=t_eval, rtol=1e-10, atol=1e-12, method="LSODA"
+    ).y.T
+    early = t_eval <= 2.0
+    np.testing.assert_allclose(y[early], ref[early], atol=1e-3)
+
+
+def test_stiff_method_aliases_route_to_stiff_kernel() -> None:
+    """LSODA/BDF names reach the stiff kernel (no silent downgrade to explicit)."""
+    sys = _LinStiff()
+    t_eval = np.linspace(0.0, 1.0, 51)
+    for alias in ("LSODA", "BDF", "Rosenbrock"):
+        y = rc.integrate_dense(sys, [1.0], t_eval, method=alias, rtol=1e-7, atol=1e-9)
+        np.testing.assert_allclose(y[:, 0], np.exp(-1000.0 * t_eval), atol=1e-6)
+
+
+def test_stiff_ensemble_deterministic_and_matches_serial() -> None:
+    sys = _VdPStiff()
+    rng = np.random.default_rng(0)
+    batch = rng.uniform(-2.0, 2.0, size=(200, 2))
+    fin = rc.ensemble_final(sys, batch, 0.0, 5.0, method="stiff", rtol=1e-6, atol=1e-8)
+    assert fin.shape == batch.shape and np.all(np.isfinite(fin))
+    np.testing.assert_array_equal(
+        fin, rc.ensemble_final(sys, batch, 0.0, 5.0, method="stiff", rtol=1e-6, atol=1e-8)
+    )
+    for i in (0, 100, 199):
+        serial = rc.integrate_dense(
+            sys, batch[i], np.array([0.0, 5.0]), method="stiff", rtol=1e-6, atol=1e-8
+        )[-1]
+        np.testing.assert_array_equal(fin[i], serial)
+
+
 def test_version_is_exposed() -> None:
     import tsdynamics_core
 

{tsdynamics-2.4.0 → tsdynamics-2.5.0}/tests/test_rustcore_translation.py

@@ -47,3 +47,32 @@ def test_tape_coverage_is_total() -> None:
         except Exception as exc:  # noqa: BLE001 — record, don't abort
             failed.append((e.name, str(exc).splitlines()[0][:60]))
     assert not failed, f"{len(failed)} ODE systems no longer compile to a tape: {failed}"
+
+
+def test_every_ode_jacobian_compiles_to_tape(ode_entry) -> None:
+    """The analytic Jacobian must lower too (stiff solver) — incl. abs/sign systems."""
+    cls = ode_entry.cls
+    try:
+        tape = compile_tape(cls(), with_jacobian=True)
+    except TapeCompileError as exc:
+        raise AssertionError(
+            f"{ode_entry.name}: Jacobian no longer lowers to a tape — {exc}"
+        ) from exc
+    dim = cls().dim
+    # One register per Jacobian entry, row-major dim×dim.
+    assert tape.jac_outputs.size == dim * dim
+    n = tape.ops.size
+    assert tape.jac_outputs.size == 0 or (
+        tape.jac_outputs.max() < n and tape.jac_outputs.min() >= 0
+    )
+
+
+def test_jacobian_coverage_is_total() -> None:
+    """Every ODE's analytic Jacobian lowers — abs/sign derivatives resolved a.e."""
+    failed = []
+    for e in registry.all_systems(family="ode"):
+        try:
+            compile_tape(e.cls(), with_jacobian=True)
+        except Exception as exc:  # noqa: BLE001 — record, don't abort
+            failed.append((e.name, str(exc).splitlines()[0][:60]))
+    assert not failed, f"{len(failed)} ODE Jacobians no longer compile to a tape: {failed}"