From 33c59ed02dc08d36ec14572649efb1cece6e08f9 Mon Sep 17 00:00:00 2001 From: jam-sudo Date: Sat, 4 Jul 2026 18:34:52 -0400 Subject: [PATCH] chore(engine): remove the dead, broken JAX MC backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The engine's JAX Monte-Carlo backend was dead and doubly broken: - solve_mc_jax built a batched JaxParams with only 18 of the dataclass's 20 required fields (missing node_fu_correction_applicable, drug_fu_correction_liver) → TypeError on construction. - jax / diffrax are not project dependencies (absent from pyproject and requirements-lock), so importing solver_jax raised ImportError first. - No caller ever selected backend="jax" (production defaults to scipy); every JAX test skipped via importorskip since jax isn't installed. It was labelled "UDE roadmap Phase 0", but UDE is an enumerated dead-end, so the scaffold has no live consumer. Removed the whole engine JAX subsystem: - delete engine/solver_jax.py, params_jax.py, rhs_jax.py - drop the backend=="jax" branch in uncertainty.py (scipy + surrogate remain) - delete test_solver_jax.py, test_jax_scipy_parity.py, test_rhs_jax_unsupported_flux.py - remove the two skipped JAX-vs-SciPy parity tests in test_active_transport*.py The surrogate backend (a separate subsystem) is untouched. SciPy is and remains the only production MC backend. tests/unit/test_active_transport*.py + test_uncertainty.py pass (15). --- src/sisyphus/engine/params_jax.py | 271 ----------- src/sisyphus/engine/rhs_jax.py | 428 ------------------ src/sisyphus/engine/solver_jax.py | 363 --------------- src/sisyphus/engine/uncertainty.py | 30 +- tests/unit/test_active_transport.py | 27 -- tests/unit/test_active_transport_direction.py | 21 - tests/unit/test_jax_scipy_parity.py | 179 -------- tests/unit/test_rhs_jax_unsupported_flux.py | 55 --- tests/unit/test_solver_jax.py | 167 ------- 9 files changed, 2 insertions(+), 1539 deletions(-) delete mode 100644 src/sisyphus/engine/params_jax.py delete mode 100644 src/sisyphus/engine/rhs_jax.py delete mode 100644 src/sisyphus/engine/solver_jax.py delete mode 100644 tests/unit/test_jax_scipy_parity.py delete mode 100644 tests/unit/test_rhs_jax_unsupported_flux.py delete mode 100644 tests/unit/test_solver_jax.py diff --git a/src/sisyphus/engine/params_jax.py b/src/sisyphus/engine/params_jax.py deleted file mode 100644 index 09d80488..00000000 --- a/src/sisyphus/engine/params_jax.py +++ /dev/null @@ -1,271 +0,0 @@ -"""JAX-compatible parameter container for the ODE engine. - -Flattens ``ResolvedParams`` (Python dicts with method dispatch) into -flat ``jnp.ndarray`` fields that are compatible with ``jax.jit``, -``jax.grad``, and ``jax.vmap``. - -Called once per MC sample to produce a ``JaxParams`` instance. -All dict lookups, enzyme aggregation, and IVIVE scaling happen at -construction time -- the RHS never touches Python dicts. - -Lazy-imports JAX so the module can be imported even when JAX is not -installed (the rest of the engine still works with SciPy). -""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -try: - import jax - - jax.config.update("jax_enable_x64", True) - import jax.numpy as jnp - - JAX_AVAILABLE = True -except ImportError: # pragma: no cover - JAX_AVAILABLE = False - jnp = None # type: ignore[assignment] - -if TYPE_CHECKING: - from sisyphus.engine.compiler import CompiledODE, ResolvedParams - - -# --------------------------------------------------------------------------- -# JaxParams — flat numeric container -# --------------------------------------------------------------------------- - - -@dataclass -class JaxParams: - """Flat, JAX-compatible parameter set for one ODE evaluation. - - Every field is either a ``jnp.ndarray`` or a Python float. - No dicts, no strings, no method dispatch -- pure numerics. - - Constructed via :func:`resolve_to_jax`, never by hand. - """ - - # -- Per-node arrays (indexed by state_index order) --------------------- - node_volumes: jnp.ndarray # (n_nodes,) - node_kp: jnp.ndarray # (n_nodes,) - node_is_blood: jnp.ndarray # (n_nodes,) 1.0 if blood_pool, else 0.0 - node_fu_correction_applicable: jnp.ndarray # (n_nodes,) 1.0 if flagged - - # -- Scalar drug properties --------------------------------------------- - drug_fup: float - drug_rbp: float - drug_mw: float - drug_renal_cl: float - drug_peff: float - drug_particle_radius: float - drug_fu_correction_liver: float - - # -- Per-edge arrays (indexed by edge_id) ------------------------------- - edge_flow_rates: jnp.ndarray # (n_edges,) - edge_transit_rates: jnp.ndarray # (n_edges,) - edge_ka_fractions: jnp.ndarray # (n_edges,) - edge_ps_products: jnp.ndarray # (n_edges,) - - # -- Pre-computed per-node clearance ------------------------------------ - node_clint_organ: jnp.ndarray # (n_nodes,) total CLint for well-stirred/PT - node_total_inflow: jnp.ndarray # (n_nodes,) sum of flow inflows (L/h) - node_ivive_scaling: jnp.ndarray # (n_nodes,) - - # -- Pre-computed per-node active transport ----------------------------- - node_transport_vmax: jnp.ndarray # (n_nodes,) sum(abundance * jmax) - node_transport_km: jnp.ndarray # (n_nodes,) effective Km per node - - -# Register JaxParams as a JAX pytree so it can be passed through jit/vmap. -if JAX_AVAILABLE: - def _jaxparams_flatten(p): - children = ( - p.node_volumes, p.node_kp, p.node_is_blood, - p.node_fu_correction_applicable, - p.drug_fup, p.drug_rbp, p.drug_mw, p.drug_renal_cl, - p.drug_peff, p.drug_particle_radius, p.drug_fu_correction_liver, - p.edge_flow_rates, p.edge_transit_rates, - p.edge_ka_fractions, p.edge_ps_products, - p.node_clint_organ, p.node_total_inflow, p.node_ivive_scaling, - p.node_transport_vmax, p.node_transport_km, - ) - return children, None - - def _jaxparams_unflatten(aux, children): - return JaxParams(*children) - - jax.tree_util.register_pytree_node( - JaxParams, _jaxparams_flatten, _jaxparams_unflatten, - ) - - -# --------------------------------------------------------------------------- -# resolve_to_jax — conversion from ResolvedParams -# --------------------------------------------------------------------------- - - -def resolve_to_jax(compiled: CompiledODE, params: ResolvedParams) -> JaxParams: - """Convert a ``ResolvedParams`` to a ``JaxParams``. - - Called once per MC sample. All dict lookups, enzyme aggregation, - and IVIVE scaling happen here so the RHS is pure numerics. - - Args: - compiled: The compiled ODE skeleton (provides topology). - params: Point-value resolved parameters for this MC sample. - - Returns: - A ``JaxParams`` ready for the JAX RHS function. - - Raises: - RuntimeError: If JAX is not installed. - """ - if not JAX_AVAILABLE: - raise RuntimeError( - "JAX is not installed. Install with: pip install jax jaxlib" - ) - - n_nodes = compiled.n_states # noqa: F841 - n_edges = len(compiled.flux_specs) # noqa: F841 - - # -- Build the sorted node name list (same order as state_index) -------- - # state_index was built from sorted(graph.nodes.keys()) - node_names_sorted = sorted(compiled.state_index.keys(), key=lambda n: compiled.state_index[n]) - - # -- Per-node arrays ---------------------------------------------------- - volumes = [] - kps = [] - is_blood = [] - fu_correction_applicable = [] - ivive_scaling = [] - total_inflow = [] - clint_organ = [] - transport_vmax = [] - transport_km = [] - - for name in node_names_sorted: - volumes.append(params.node_param(name, "volume")) - kps.append(params.drug_kp(name)) - is_blood.append(1.0 if params.is_blood_pool(name) else 0.0) - fu_correction_applicable.append(params.node_param(name, "fu_correction_applicable")) - ivive_scaling.append(params.node_param(name, "ivive_scaling")) - total_inflow.append(params.total_inflow(name)) - - # Pre-compute organ-level CLint = sum(abundance_i * affinity_i * ivive) - node_ivive = params.node_param(name, "ivive_scaling") - enzymes = params.node_enzymes(name) - cl_sum = 0.0 - for tag, abundance in enzymes.items(): - affinity = params.drug_enzyme_affinity(tag) - if affinity > 0 and abundance > 0: - cl_sum += abundance * affinity * node_ivive - clint_organ.append(cl_sum) - - # Pre-compute active transport: aggregate Vmax and effective Km - # Vmax_total = sum(abundance_i * jmax_i) - # Effective Km: weighted harmonic mean, or simple average if only one. - # For Michaelis-Menten with multiple transporters summed: - # total_rate = sum_i( abundance_i * jmax_i * C / (km_i + C) ) - # We cannot collapse to single Vmax/Km exactly when Km differs. - # Approximation: store aggregate Vmax and abundance-weighted Km. - transporters = params.node_transporters(name) - active_kms = set() - for tag, abundance in transporters.items(): - j = params.drug_transporter_jmax(tag) - km = params.drug_transporter_km(tag) - if j > 0 and km > 0 and abundance > 0: - active_kms.add(round(km, 9)) - if len(active_kms) > 1: - raise NotImplementedError( - f"node {name!r} has {len(active_kms)} active transporters with " - f"distinct Km {sorted(active_kms)}; the JAX aggregate-Vmax / " - f"weighted-Km approximation diverges from SciPy's exact " - f"per-transporter Michaelis-Menten sum. Use backend='scipy' " - f"(default), or implement exact padded per-transporter MM in JAX." - ) - vmax_sum = 0.0 - km_weighted_num = 0.0 # sum(abundance * jmax * km) - vmax_weight = 0.0 # sum(abundance * jmax) -- same as vmax_sum - for tag, abundance in transporters.items(): - j = params.drug_transporter_jmax(tag) - km = params.drug_transporter_km(tag) - if j > 0 and km > 0 and abundance > 0: - contrib = abundance * j - vmax_sum += contrib - km_weighted_num += contrib * km - vmax_weight += contrib - - transport_vmax.append(vmax_sum) - # Weighted average Km (weight = abundance * jmax). - # Falls back to 1.0 if no transporters (never used since vmax=0). - transport_km.append(km_weighted_num / vmax_weight if vmax_weight > 0 else 1.0) - - # -- Per-edge arrays ---------------------------------------------------- - flow_rates = [] - transit_rates = [] - ka_fractions = [] - ps_products = [] - - for spec in compiled.flux_specs: - edge_id = spec.edge_id - - # Flow - try: - flow_rates.append(params.edge_param(edge_id, "flow_rate")) - except KeyError: - flow_rates.append(0.0) - - # Transit - try: - transit_rates.append(params.edge_param(edge_id, "transit_rate")) - except KeyError: - transit_rates.append(0.0) - - # Absorption - try: - ka_fractions.append(params.edge_param(edge_id, "ka_fraction")) - except KeyError: - ka_fractions.append(0.0) - - # Diffusion (PS product) - # For diffusion edges, prefer drug-specific PS override, fall back to edge param - try: - ps = params.drug_ps(spec.source_name) - if ps <= 0: - ps = params.edge_param(edge_id, "ps_product") - ps_products.append(ps) - except KeyError: - ps_products.append(0.0) - - # -- Scalar drug properties --------------------------------------------- - drug_fup = params.drug_param("fup") - drug_rbp = params.drug_param("rbp") - drug_mw = params.drug_mw() - drug_renal_cl = params.drug_param("renal_clearance") - drug_peff = params.drug_param("peff") - drug_particle_radius = params.drug_param("particle_radius_um") - - return JaxParams( - node_volumes=jnp.array(volumes, dtype=jnp.float64), - node_kp=jnp.array(kps, dtype=jnp.float64), - node_is_blood=jnp.array(is_blood, dtype=jnp.float64), - node_fu_correction_applicable=jnp.array(fu_correction_applicable, dtype=jnp.float64), - drug_fup=drug_fup, - drug_rbp=drug_rbp, - drug_mw=drug_mw, - drug_renal_cl=drug_renal_cl, - drug_peff=drug_peff, - drug_particle_radius=drug_particle_radius, - drug_fu_correction_liver=params.drug_param("fu_correction_liver"), - edge_flow_rates=jnp.array(flow_rates, dtype=jnp.float64), - edge_transit_rates=jnp.array(transit_rates, dtype=jnp.float64), - edge_ka_fractions=jnp.array(ka_fractions, dtype=jnp.float64), - edge_ps_products=jnp.array(ps_products, dtype=jnp.float64), - node_clint_organ=jnp.array(clint_organ, dtype=jnp.float64), - node_total_inflow=jnp.array(total_inflow, dtype=jnp.float64), - node_ivive_scaling=jnp.array(ivive_scaling, dtype=jnp.float64), - node_transport_vmax=jnp.array(transport_vmax, dtype=jnp.float64), - node_transport_km=jnp.array(transport_km, dtype=jnp.float64), - ) diff --git a/src/sisyphus/engine/rhs_jax.py b/src/sisyphus/engine/rhs_jax.py deleted file mode 100644 index e27ed0e4..00000000 --- a/src/sisyphus/engine/rhs_jax.py +++ /dev/null @@ -1,428 +0,0 @@ -"""Pure-JAX RHS function builder for the ODE engine. - -Builds a JIT-compilable right-hand side function from the compiled ODE -topology. At **compile time**, static index arrays are extracted from -``CompiledODE.flux_specs`` (which edges are flow, which are clearance, -etc.). At **runtime**, all fluxes are computed functionally with -vectorized operations -- no Python loops, no dict lookups, no -if/else branching. - -The resulting ``rhs(t, y, params)`` is compatible with ``jax.jit``, -``jax.grad``, and ``jax.vmap``. - -Lazy-imports JAX so the module can be imported even when JAX is not -installed (the rest of the engine still works with SciPy). -""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import TYPE_CHECKING - -try: - import jax - - jax.config.update("jax_enable_x64", True) - import jax.numpy as jnp - - JAX_AVAILABLE = True -except ImportError: # pragma: no cover - JAX_AVAILABLE = False - jnp = None # type: ignore[assignment] - -if TYPE_CHECKING: - from sisyphus.engine.compiler import CompiledODE - from sisyphus.engine.params_jax import JaxParams - - -# --------------------------------------------------------------------------- -# Internal: flux-type detection helpers -# --------------------------------------------------------------------------- - -# Import concrete FluxSpec subclasses for isinstance checks. -# These are only needed at build time (not at JAX runtime). -from sisyphus.engine.flux import ( - AbsorptionFluxSpec, - ActiveTransportFluxSpec, - ClearanceFluxSpec, - DiffusionFluxSpec, - FlowFluxSpec, - TransitFluxSpec, -) - -# Concrete FluxSpec subclasses the JAX RHS can vectorize. ProdrugActivation and -# OneCompartmentElimination are intentionally absent — they have no JAX branch. -# Before _unsupported_flux_specs() guarded make_jax_rhs, such specs fell through -# the build loop silently, omitting their flux from the RHS. -_JAX_SUPPORTED_FLUX = ( - FlowFluxSpec, - ClearanceFluxSpec, - TransitFluxSpec, - AbsorptionFluxSpec, - DiffusionFluxSpec, - ActiveTransportFluxSpec, -) - - -def _unsupported_flux_specs(flux_specs): - """Return the flux specs whose type the JAX RHS cannot vectorize. - - Pure Python (imports no JAX) so the "no silent drop" invariant is testable - even when JAX is not installed. - """ - return [s for s in flux_specs if not isinstance(s, _JAX_SUPPORTED_FLUX)] - - -# --------------------------------------------------------------------------- -# make_jax_rhs — the public API -# --------------------------------------------------------------------------- - - -def make_jax_rhs( - compiled: CompiledODE, -) -> Callable: - """Build a pure-JAX RHS from compiled ODE topology. - - Extracts static topology (edge index arrays per flux type) at build - time. Returns a function ``rhs(t, y, params: JaxParams) -> dydt`` - that uses only JAX operations and is compatible with ``jax.jit``, - ``jax.grad``, and ``jax.vmap``. - - Args: - compiled: A ``CompiledODE`` compiled from a ``BodyGraph``. - - Returns: - A callable ``rhs(t, y, params) -> dydt`` where ``y`` and ``dydt`` - are ``jnp.ndarray`` of shape ``(n_states,)`` and ``params`` is a - ``JaxParams``. - - Raises: - RuntimeError: If JAX is not installed. - """ - if not JAX_AVAILABLE: - raise RuntimeError( - "JAX is not installed. Install with: pip install jax jaxlib" - ) - - n_states = compiled.n_states - - # Fail loudly rather than silently dropping a flux the JAX RHS cannot - # vectorize (e.g. ProdrugActivation / OneCompartmentElimination). The - # SciPy backend handles them; backend="jax" must not corrupt those graphs. - unsupported = _unsupported_flux_specs(compiled.flux_specs) - if unsupported: - names = sorted({type(s).__name__ for s in unsupported}) - raise NotImplementedError( - f"JAX RHS cannot vectorize flux type(s) {names}; they would be " - f"silently dropped. Use backend='scipy' (default) or add a " - f"vmap-safe branch to rhs_jax.py." - ) - - # -- Extract static topology per flux type (Python, not JAX) ------------ - # Each list collects (source_idx, target_idx, edge_id) for that flux type. - - flow_src = [] - flow_tgt = [] - flow_eid = [] - - # Clearance: separate by model type - cl_ws_src = [] # well_stirred source indices - cl_ws_tgt = [] - cl_gfr_src = [] # gfr_filtration source indices - cl_gfr_tgt = [] - - transit_src = [] - transit_tgt = [] - transit_eid = [] - - absorption_src = [] - absorption_tgt = [] - absorption_eid = [] - - diffusion_src = [] - diffusion_tgt = [] - diffusion_eid = [] - - # Active transport, split by direction. The transporter-bearing node differs: - # uptake → target node; efflux → source node. Substrate conc is always the source. - transport_up_src, transport_up_node, transport_up_tgt = [], [], [] - transport_ef_src, transport_ef_node, transport_ef_tgt = [], [], [] - - for spec in compiled.flux_specs: - if isinstance(spec, FlowFluxSpec): - flow_src.append(spec.source_idx) - flow_tgt.append(spec.target_idx) - flow_eid.append(spec.edge_id) - - elif isinstance(spec, ClearanceFluxSpec): - if spec.model == "well_stirred": - cl_ws_src.append(spec.source_idx) - cl_ws_tgt.append(spec.target_idx) - elif spec.model == "gfr_filtration": - cl_gfr_src.append(spec.source_idx) - cl_gfr_tgt.append(spec.target_idx) - elif spec.model == "extended": - # ECM (2026-04-20 migration) is not yet vectorised for JAX. - # The scipy backend handles it correctly; fail loudly here so - # experiments with backend="jax" do not silently zero hepatic - # clearance for OATP substrates. - raise NotImplementedError( - "ClearanceFluxSpec model='extended' (ECM) is not yet " - "implemented in the JAX RHS. Use backend='scipy' (default) " - "or add a vmap-safe ECM branch to rhs_jax.py." - ) - else: - raise NotImplementedError( - f"ClearanceFluxSpec model={spec.model!r} not supported in JAX RHS." - ) - - elif isinstance(spec, TransitFluxSpec): - transit_src.append(spec.source_idx) - transit_tgt.append(spec.target_idx) - transit_eid.append(spec.edge_id) - - elif isinstance(spec, AbsorptionFluxSpec): - absorption_src.append(spec.source_idx) - absorption_tgt.append(spec.target_idx) - absorption_eid.append(spec.edge_id) - - elif isinstance(spec, DiffusionFluxSpec): - diffusion_src.append(spec.source_idx) - diffusion_tgt.append(spec.target_idx) - diffusion_eid.append(spec.edge_id) - - elif isinstance(spec, ActiveTransportFluxSpec): - if getattr(spec, "direction", "uptake") == "uptake": - transport_up_src.append(spec.source_idx) - transport_up_node.append(spec.target_idx) # transporter at target - transport_up_tgt.append(spec.target_idx) - else: - transport_ef_src.append(spec.source_idx) - transport_ef_node.append(spec.source_idx) # transporter at source - transport_ef_tgt.append(spec.target_idx) - - # Convert to JAX arrays (static, traced as constants by jit) - _flow_src = jnp.array(flow_src, dtype=jnp.int32) - _flow_tgt = jnp.array(flow_tgt, dtype=jnp.int32) - _flow_eid = jnp.array(flow_eid, dtype=jnp.int32) - - _cl_ws_src = jnp.array(cl_ws_src, dtype=jnp.int32) - _cl_ws_tgt = jnp.array(cl_ws_tgt, dtype=jnp.int32) - _cl_gfr_src = jnp.array(cl_gfr_src, dtype=jnp.int32) - _cl_gfr_tgt = jnp.array(cl_gfr_tgt, dtype=jnp.int32) - - _transit_src = jnp.array(transit_src, dtype=jnp.int32) - _transit_tgt = jnp.array(transit_tgt, dtype=jnp.int32) - _transit_eid = jnp.array(transit_eid, dtype=jnp.int32) - - _absorption_src = jnp.array(absorption_src, dtype=jnp.int32) - _absorption_tgt = jnp.array(absorption_tgt, dtype=jnp.int32) - _absorption_eid = jnp.array(absorption_eid, dtype=jnp.int32) - - _diffusion_src = jnp.array(diffusion_src, dtype=jnp.int32) - _diffusion_tgt = jnp.array(diffusion_tgt, dtype=jnp.int32) - _diffusion_eid = jnp.array(diffusion_eid, dtype=jnp.int32) - - _t_up_src = jnp.array(transport_up_src, dtype=jnp.int32) - _t_up_node = jnp.array(transport_up_node, dtype=jnp.int32) - _t_up_tgt = jnp.array(transport_up_tgt, dtype=jnp.int32) - _t_ef_src = jnp.array(transport_ef_src, dtype=jnp.int32) - _t_ef_node = jnp.array(transport_ef_node, dtype=jnp.int32) - _t_ef_tgt = jnp.array(transport_ef_tgt, dtype=jnp.int32) - - # -- Boolean flags for which flux types are present --------------------- - has_flow = len(flow_src) > 0 - has_cl_ws = len(cl_ws_src) > 0 - has_cl_gfr = len(cl_gfr_src) > 0 - has_transit = len(transit_src) > 0 - has_absorption = len(absorption_src) > 0 - has_diffusion = len(diffusion_src) > 0 - has_t_up = len(transport_up_src) > 0 - has_t_ef = len(transport_ef_src) > 0 - - # -- The RHS function --------------------------------------------------- - - def rhs(t: float, y: jnp.ndarray, params: JaxParams) -> jnp.ndarray: - """Pure-JAX ODE right-hand side. - - All operations use ``jnp`` -- no Python control flow on data, - no in-place mutation. Uses ``jnp.where`` instead of ``if/else`` - and ``.at[].add()`` for scatter-add. - - Args: - t: Current time (h). - y: State vector, shape ``(n_states,)`` -- amounts in mg. - params: A ``JaxParams`` instance. - - Returns: - Derivative vector ``dydt``, shape ``(n_states,)``. - """ - dydt = jnp.zeros(n_states, dtype=jnp.float64) - - # Shorthand for frequently used params - fup = params.drug_fup - rbp = params.drug_rbp - - # --------------------------------------------------------------- - # 1. Flow fluxes (convective transport) - # C_out = A[src] * RBP / (V[src] * Kp[src]) (tissue) - # C_out = A[src] / V[src] (blood pool, Kp=1) - # flux = Q * C_out - # --------------------------------------------------------------- - if has_flow: - v_src = params.node_volumes[_flow_src] - kp_src = params.node_kp[_flow_src] - q = params.edge_flow_rates[_flow_eid] - - # RBP-2: convective edge carries WHOLE BLOOD. Blood-pool source A/V is - # already blood (gate RBP->1); tissue source outlet is A*RBP/(V*Kp). - rbp_flow = jnp.where(params.node_is_blood[_flow_src] > 0.5, 1.0, rbp) - c_out = jnp.where( - v_src > 0.0, - y[_flow_src] * rbp_flow / (v_src * kp_src), - 0.0, - ) - flow_flux = q * c_out - - dydt = dydt.at[_flow_src].add(-flow_flux) - dydt = dydt.at[_flow_tgt].add(flow_flux) - - # --------------------------------------------------------------- - # 2. Clearance fluxes - # --------------------------------------------------------------- - - # 2a. Well-stirred model - # FLUX-1: intrinsic clearance applied to C_out. The convective - # Q*c_out outflow edge supplies the flow limitation, so realized - # extraction is fup*CLint/(Q+fup*CLint) -> 1.0. Applying the - # whole-organ CL_h here (embedding Q) double-counts flow (E->0.5). - # rate = (fup * CLint) * C_out ; C_out = A[src]*RBP/(V[src]*Kp[src]) - if has_cl_ws: - clint = params.node_clint_organ[_cl_ws_src] - v_ws = params.node_volumes[_cl_ws_src] - kp_ws = params.node_kp[_cl_ws_src] - - # WS-4: hepatic intracellular fu correction at flagged nodes (parity - # with the SciPy well_stirred branch). - fu_corr = params.drug_fu_correction_liver - applicable = params.node_fu_correction_applicable[_cl_ws_src] - fup_eff = jnp.where(applicable > 0.5, fup * fu_corr, fup) - cl_intrinsic_ws = fup_eff * clint - # RBP-2: plasma-basis sink → emergent fu_b extraction (see numpy flux). - c_plasma_ws = jnp.where( - v_ws > 0.0, - y[_cl_ws_src] / (v_ws * kp_ws), - 0.0, - ) - rate_ws = cl_intrinsic_ws * c_plasma_ws - - dydt = dydt.at[_cl_ws_src].add(-rate_ws) - dydt = dydt.at[_cl_ws_tgt].add(rate_ws) - - # 2c. GFR filtration - # rate = renal_cl * C_plasma - # C_plasma = A[src] * RBP / (V[src] * Kp[src]) - if has_cl_gfr: - renal_cl = params.drug_renal_cl - v_gfr = params.node_volumes[_cl_gfr_src] - kp_gfr = params.node_kp[_cl_gfr_src] - - # RBP-2: renal_cl = GFR*fup is plasma-basis → filter PLASMA A/(V*Kp). - c_plasma_gfr = jnp.where( - v_gfr > 0.0, - y[_cl_gfr_src] / (v_gfr * kp_gfr), - 0.0, - ) - rate_gfr = renal_cl * c_plasma_gfr - - dydt = dydt.at[_cl_gfr_src].add(-rate_gfr) - dydt = dydt.at[_cl_gfr_tgt].add(rate_gfr) - - # --------------------------------------------------------------- - # 3. Transit fluxes (first-order) - # rate = k_transit * A[src] - # --------------------------------------------------------------- - if has_transit: - kt = params.edge_transit_rates[_transit_eid] - transit_flux = kt * y[_transit_src] - - dydt = dydt.at[_transit_src].add(-transit_flux) - dydt = dydt.at[_transit_tgt].add(transit_flux) - - # --------------------------------------------------------------- - # 4. Absorption fluxes - # ka = 2.88 * Peff * ka_fraction / particle_radius - # rate = ka * A[src] - # --------------------------------------------------------------- - if has_absorption: - ka_frac = params.edge_ka_fractions[_absorption_eid] - peff = params.drug_peff - radius = params.drug_particle_radius - - # ka = 2.88 * Peff * ka_fraction / radius - ka = jnp.where( - (ka_frac > 0.0) & (peff > 0.0) & (radius > 0.0), - 2.88 * peff * ka_frac / radius, - 0.0, - ) - absorption_flux = ka * y[_absorption_src] - - dydt = dydt.at[_absorption_src].add(-absorption_flux) - dydt = dydt.at[_absorption_tgt].add(absorption_flux) - - # --------------------------------------------------------------- - # 5. Diffusion fluxes (PS-limited) - # cu_vasc = fup * C_vasc / RBP - # cu_tissue = fup * C_tissue / Kp - # flux = PS * (cu_vasc - cu_tissue) - # --------------------------------------------------------------- - if has_diffusion: - ps = params.edge_ps_products[_diffusion_eid] - - v_vasc = params.node_volumes[_diffusion_src] - v_tissue = params.node_volumes[_diffusion_tgt] - kp_tgt = params.node_kp[_diffusion_tgt] - - c_vasc = jnp.where(v_vasc > 0.0, y[_diffusion_src] / v_vasc, 0.0) - c_tissue = jnp.where(v_tissue > 0.0, y[_diffusion_tgt] / v_tissue, 0.0) - - cu_vasc = jnp.where(rbp > 0.0, fup * c_vasc / rbp, 0.0) - cu_tissue = jnp.where(kp_tgt > 0.0, fup * c_tissue / kp_tgt, 0.0) - - diff_flux = jnp.where(ps > 0.0, ps * (cu_vasc - cu_tissue), 0.0) - - dydt = dydt.at[_diffusion_src].add(-diff_flux) - dydt = dydt.at[_diffusion_tgt].add(diff_flux) - - # --------------------------------------------------------------- - # 6. Active transport (Michaelis-Menten), split by direction. - # Substrate conc from the SOURCE; Vmax/Km/ivive from the transporter node - # (target for uptake, source for efflux). mass moves source → edge target. - # --------------------------------------------------------------- - def _transport_mass(src_idx, node_idx): - mw = params.drug_mw - v_src_t = params.node_volumes[src_idx] - vmax = params.node_transport_vmax[node_idx] - km = params.node_transport_km[node_idx] - ivive_t = params.node_ivive_scaling[node_idx] - c_mg_l = jnp.where(v_src_t > 0.0, y[src_idx] / v_src_t, 0.0) - c_um = jnp.where(mw > 0.0, c_mg_l * 1000.0 / mw, 0.0) - mm_rate = jnp.where( - (vmax > 0.0) & (km > 0.0) & (c_um > 0.0), - vmax * c_um / (km + c_um), - 0.0, - ) - return mm_rate * ivive_t - - if has_t_up: - mass = _transport_mass(_t_up_src, _t_up_node) - dydt = dydt.at[_t_up_src].add(-mass) - dydt = dydt.at[_t_up_tgt].add(mass) - if has_t_ef: - mass = _transport_mass(_t_ef_src, _t_ef_node) - dydt = dydt.at[_t_ef_src].add(-mass) - dydt = dydt.at[_t_ef_tgt].add(mass) - - return dydt - - return rhs diff --git a/src/sisyphus/engine/solver_jax.py b/src/sisyphus/engine/solver_jax.py deleted file mode 100644 index f8205939..00000000 --- a/src/sisyphus/engine/solver_jax.py +++ /dev/null @@ -1,363 +0,0 @@ -"""JAX/Diffrax-based ODE solver — differentiable alternative to SciPy LSODA. - -Phase 0 of the UDE roadmap (see docs/breakthrough_path.md). -Provides a drop-in replacement for ``solve()`` that uses Diffrax Kvaerno5 -for stiff ODE integration, with support for autograd through the solver. - -Two code paths: - 1. ``solve_jax_wrapped``: wraps an existing NumPy RHS via jax.pure_callback. - Numerical equivalence to LSODA but NO gradients (numpy boundary breaks - autograd). Used for Phase 0 validation. - - 2. ``solve_jax_pure``: takes a pure-JAX vector field. Supports - jax.grad / jax.jvp / jax.vjp through the adjoint method. - Used by Phase 1 UDE training. - -The adjoint method (RecursiveCheckpointAdjoint) gives O(log N) memory -for backprop through N solver steps — essential for end-to-end training -of neural closures against observed Cmax. -""" - -from __future__ import annotations - -import logging -from collections.abc import Callable -from typing import Any - -import jax - -# PBPK requires float64 for mass balance accuracy over long timescales. -jax.config.update("jax_enable_x64", True) - -import jax.numpy as jnp # noqa: E402 -import numpy as np # noqa: E402 -from diffrax import ( # noqa: E402 - RESULTS, - Kvaerno3, - Kvaerno5, - ODETerm, - PIDController, - RecursiveCheckpointAdjoint, - SaveAt, - Tsit5, - diffeqsolve, -) - -from sisyphus.core import SimResult # noqa: E402 -from sisyphus.engine.compiler import CompiledODE, ResolvedParams # noqa: E402 -from sisyphus.engine.params_jax import JaxParams, resolve_to_jax # noqa: E402 -from sisyphus.engine.rhs_jax import make_jax_rhs # noqa: E402 - -logger = logging.getLogger(__name__) - -# Default solver tolerances — match SciPy LSODA defaults used in solver.py -_DEFAULT_RTOL = 1e-8 -_DEFAULT_ATOL = 1e-10 -_DEFAULT_DT0 = 0.01 -_DEFAULT_MAX_STEPS = 16384 - - -def solve_jax_pure( - vector_field: Callable[[float, jnp.ndarray, Any], jnp.ndarray], - y0: jnp.ndarray, - t_span: tuple[float, float], - args: Any = None, - t_eval: jnp.ndarray | None = None, - rtol: float = _DEFAULT_RTOL, - atol: float = _DEFAULT_ATOL, - stiff: bool = True, - max_steps: int = _DEFAULT_MAX_STEPS, - allow_jvp: bool = True, -) -> tuple[jnp.ndarray, jnp.ndarray, bool]: - """Solve an ODE system using Diffrax Kvaerno5 (5th-order stiff). - - Args: - vector_field: pure JAX function ``f(t, y, args) -> dy/dt``. - y0: initial state vector (JAX array, shape ``(n_states,)``). - t_span: ``(t0, t1)`` integration interval in hours. - args: extra arguments passed to ``vector_field``. - t_eval: time points to return solution at. If ``None``, uses 500 - evenly-spaced points. - rtol: relative tolerance. - atol: absolute tolerance. - stiff: if True, use Kvaerno5 (implicit). If False, use explicit RK. - max_steps: maximum solver steps. - allow_jvp: if False, use an explicit solver (Tsit5) that does not - require JVP. Required when the vector field wraps a - jax.pure_callback (no JVP support). - - Returns: - Tuple of (t, y, success) where: - - t: shape (T,) - - y: shape (T, n_states) - - success: bool indicating solver success - """ - t0, t1 = t_span - if t_eval is None: - t_eval = jnp.linspace(t0, t1, 500) - - term = ODETerm(vector_field) - if not allow_jvp: - # Explicit solver for non-differentiable vector fields (numpy wrap). - solver = Tsit5() - else: - solver = Kvaerno5() if stiff else Kvaerno3() - saveat = SaveAt(ts=t_eval) - controller = PIDController(rtol=rtol, atol=atol) - - sol = diffeqsolve( - term, - solver, - t0=t0, - t1=t1, - dt0=_DEFAULT_DT0, - y0=y0, - args=args, - saveat=saveat, - stepsize_controller=controller, - max_steps=max_steps, - adjoint=RecursiveCheckpointAdjoint(), - throw=False, # return result even on solver failure (check result.result) - ) - - success = sol.result == RESULTS.successful if hasattr(sol, "result") else True - return sol.ts, sol.ys, bool(success) - - -def _wrap_numpy_rhs(rhs_np: Callable[[float, np.ndarray], np.ndarray], n: int): - """Wrap a NumPy RHS for Diffrax via jax.pure_callback. - - NOTE: This breaks autograd at the numpy boundary. Used for validation - only — Phase 1 requires a pure-JAX vector field. - """ - out_shape = jax.ShapeDtypeStruct((n,), jnp.float64) - - def _cb(tt: jnp.ndarray, yy: jnp.ndarray) -> jnp.ndarray: - # pure_callback calls expect host-side numpy - dydt = rhs_np(float(tt), np.asarray(yy, dtype=np.float64)) - return np.asarray(dydt, dtype=np.float64) - - def vf(t, y, _args): - return jax.pure_callback(_cb, out_shape, t, y) - - return vf - - -def solve_jax_wrapped( - rhs_np: Callable[[float, np.ndarray], np.ndarray], - y0: np.ndarray, - t_span: tuple[float, float], - t_eval: np.ndarray | None = None, - rtol: float = _DEFAULT_RTOL, - atol: float = _DEFAULT_ATOL, - stiff: bool = True, -) -> tuple[np.ndarray, np.ndarray, bool]: - """Validation path: wrap a NumPy RHS for Diffrax solver. - - Returns numerical-equivalent results to SciPy LSODA but no autograd. - Use ``solve_jax_pure`` with a JAX-compatible RHS for gradient flow. - """ - n = len(y0) - vf = _wrap_numpy_rhs(rhs_np, n) - y0_j = jnp.asarray(y0, dtype=jnp.float64) - t_eval_j = jnp.asarray(t_eval, dtype=jnp.float64) if t_eval is not None else None - - # Wrapped path uses explicit solver (Tsit5) because jax.pure_callback - # does not support JVP needed by implicit solvers. - ts, ys, success = solve_jax_pure( - vf, y0_j, t_span, args=None, t_eval=t_eval_j, rtol=rtol, atol=atol, - stiff=stiff, allow_jvp=False, max_steps=65536, - ) - return np.asarray(ts), np.asarray(ys), success - - -def solve( - compiled: CompiledODE, - params: ResolvedParams, - y0: np.ndarray, - t_span: tuple[float, float], - t_eval: np.ndarray | None = None, -) -> SimResult: - """Drop-in replacement for ``engine.solver.solve`` using Diffrax Kvaerno5. - - Produces a ``SimResult`` identical in structure to the SciPy-based solver. - Concentration conversion, mass-balance checking, and output formatting - match the existing implementation bit-for-bit. - - Uses the numpy-wrapped RHS (jax.pure_callback) — validation path. - """ - rhs_np = compiled.make_rhs(params) - - if t_eval is None: - t_eval = np.linspace(t_span[0], t_span[1], 500) - - ts, ys, success = solve_jax_wrapped( - rhs_np, y0, t_span, t_eval=t_eval, rtol=_DEFAULT_RTOL, atol=_DEFAULT_ATOL, - ) - # Diffrax returns ys as (T, n_states). SciPy returns as (n_states, T). - # Transpose to match existing interface. - y = ys.T # (n_states, T) - - # Build concentration and amount dicts — same logic as solver.py - concentrations: dict[str, np.ndarray] = {} - amounts: dict[str, np.ndarray] = {} - for name, idx in compiled.state_index.items(): - amounts[name] = y[idx] - v = params.node_param(name, "volume") - if v > 0: - if params.is_blood_pool(name): - concentrations[name] = y[idx] / v - else: - kp = params.drug_kp(name) - concentrations[name] = y[idx] / (v * kp) if kp > 0 else y[idx] / v - else: - concentrations[name] = y[idx] - - # Mass balance check - total = np.zeros_like(ts) - for idx in range(compiled.n_states): - total += y[idx] - dose = params.drug_param("dose_mg") - mbe = float(np.max(np.abs(total - dose) / dose)) if dose > 0 else 0.0 - - return SimResult( - time_h=ts, - concentrations=concentrations, - amounts=amounts, - mass_balance_error=mbe, - solver_success=success, - ) - - -# --------------------------------------------------------------------------- -# vmap MC — batch ODE solve over N parameter sets -# --------------------------------------------------------------------------- - -_MC_RTOL = 1e-4 -_MC_ATOL = 1e-6 -_MC_MAX_STEPS = 8192 - - -def _solve_single_mc( - rhs_fn, - y0: jnp.ndarray, - t_span: tuple[float, float], - params: JaxParams, - obs_idx: int, - obs_volume: float, - obs_kp: float, -) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: - """Solve one MC sample, return (cmax, tmax, auc, success).""" - t0, t1 = t_span - # Adaptive grid for speed — use SaveAt with dense output - t_eval = jnp.linspace(t0, t1, 200) - term = ODETerm(lambda t, y, args: rhs_fn(t, y, args)) - solver = Kvaerno5() - saveat = SaveAt(ts=t_eval) - controller = PIDController(rtol=_MC_RTOL, atol=_MC_ATOL) - - sol = diffeqsolve( - term, solver, - t0=t0, t1=t1, dt0=0.01, - y0=y0, args=params, - saveat=saveat, - stepsize_controller=controller, - max_steps=_MC_MAX_STEPS, - adjoint=RecursiveCheckpointAdjoint(), - throw=False, - ) - - success = jnp.where(sol.result == RESULTS.successful, 1.0, 0.0) - - # Extract observation node concentration: C = A / (V * Kp) - amounts_obs = sol.ys[:, obs_idx] - conc = amounts_obs / (obs_volume * obs_kp) - conc = jnp.maximum(conc, 0.0) - - cmax = jnp.max(conc) - tmax = sol.ts[jnp.argmax(conc)] - auc = jnp.trapezoid(conc, sol.ts) - - return cmax, tmax, auc, success - - -def solve_mc_jax( - compiled: CompiledODE, - params_list: list[ResolvedParams], - y0_template: np.ndarray, - t_span: tuple[float, float], - observation_node: str = "venous_blood", -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Batch MC solve over N parameter sets using JAX vmap. - - Args: - compiled: Compiled ODE skeleton. - params_list: List of N ResolvedParams (one per MC sample). - y0_template: Initial state vector (same for all samples). - t_span: Integration interval (t0, t1). - observation_node: Node to extract Cmax/AUC from. - - Returns: - (cmax[N], tmax[N], auc[N], success[N]) as numpy arrays. - """ - n_samples = len(params_list) - if n_samples == 0: - empty = np.array([]) - return empty, empty, empty, empty - - # Build the pure-JAX RHS (topology is static, shared across samples) - rhs_fn = make_jax_rhs(compiled) - - # Convert all ResolvedParams to JaxParams and stack into batched arrays - jax_params_all = [resolve_to_jax(compiled, p) for p in params_list] - - # Stack into a single JaxParams with leading N dimension - stacked = JaxParams( - node_volumes=jnp.stack([p.node_volumes for p in jax_params_all]), - node_kp=jnp.stack([p.node_kp for p in jax_params_all]), - node_is_blood=jnp.stack([p.node_is_blood for p in jax_params_all]), - drug_fup=jnp.array([p.drug_fup for p in jax_params_all]), - drug_rbp=jnp.array([p.drug_rbp for p in jax_params_all]), - drug_mw=jnp.array([p.drug_mw for p in jax_params_all]), - drug_renal_cl=jnp.array([p.drug_renal_cl for p in jax_params_all]), - drug_peff=jnp.array([p.drug_peff for p in jax_params_all]), - drug_particle_radius=jnp.array([p.drug_particle_radius for p in jax_params_all]), - edge_flow_rates=jnp.stack([p.edge_flow_rates for p in jax_params_all]), - edge_transit_rates=jnp.stack([p.edge_transit_rates for p in jax_params_all]), - edge_ka_fractions=jnp.stack([p.edge_ka_fractions for p in jax_params_all]), - edge_ps_products=jnp.stack([p.edge_ps_products for p in jax_params_all]), - node_clint_organ=jnp.stack([p.node_clint_organ for p in jax_params_all]), - node_total_inflow=jnp.stack([p.node_total_inflow for p in jax_params_all]), - node_ivive_scaling=jnp.stack([p.node_ivive_scaling for p in jax_params_all]), - node_transport_vmax=jnp.stack([p.node_transport_vmax for p in jax_params_all]), - node_transport_km=jnp.stack([p.node_transport_km for p in jax_params_all]), - ) - - obs_idx = compiled.state_index[observation_node] - # Use first sample's volume and kp for observation node - # (these vary per sample, but we pass them as arrays for vmap) - obs_volumes = stacked.node_volumes[:, obs_idx] - obs_kps = stacked.node_kp[:, obs_idx] - - y0_jax = jnp.asarray(y0_template, dtype=jnp.float64) - - # vmap over samples: each sample gets its own JaxParams slice - @jax.jit - def _batch_solve(stacked_params, obs_vols, obs_kps): - def _single(params_i, vol_i, kp_i): - return _solve_single_mc( - rhs_fn, y0_jax, t_span, params_i, obs_idx, vol_i, kp_i, - ) - return jax.vmap(_single)(stacked_params, obs_vols, obs_kps) - - logger.info("solve_mc_jax: vmapping %d samples...", n_samples) - cmax_arr, tmax_arr, auc_arr, success_arr = _batch_solve( - stacked, obs_volumes, obs_kps, - ) - - return ( - np.asarray(cmax_arr), - np.asarray(tmax_arr), - np.asarray(auc_arr), - np.asarray(success_arr), - ) diff --git a/src/sisyphus/engine/uncertainty.py b/src/sisyphus/engine/uncertainty.py index 4441eeed..86f0006c 100644 --- a/src/sisyphus/engine/uncertainty.py +++ b/src/sisyphus/engine/uncertainty.py @@ -202,10 +202,9 @@ def propagate_fast( observation_node: Node to extract PK from. t_min_h: Minimum observation time in hours. When > 0, skips the t=0 IV-bolus spike for Cmax extraction. scipy backend only — - JAX and surrogate backends ignore this parameter. Default 0.0 + the surrogate backend ignores this parameter. Default 0.0 (V2-compatible). - backend: Solver backend to use (``"scipy"``, ``"jax"``, or - ``"surrogate"``). + backend: Solver backend to use (``"scipy"`` or ``"surrogate"``). Returns: MCResult with aggregated PKEndpoints and 90% PI. @@ -249,31 +248,6 @@ def propagate_fast( logger.info("MC-fast surrogate: %d/%d successful", n_ok, n_samples) return self._build_mc_result(cmax_arr, tmax_arr, auc_arr, n_ok, n_failures) - # ── JAX backend: vmap over all samples in one JIT call ── - if backend == "jax": - from sisyphus.engine.solver_jax import solve_mc_jax - - params_list = [] - for i in range(n_samples): - rng = np.random.default_rng(seed + i) - realized_graph = graph.sample(rng) - realized_drug = drug.sample(rng) - params_list.append(ResolvedParams(realized_graph, realized_drug)) - - cmax_arr, tmax_arr, auc_arr, success_arr = solve_mc_jax( - compiled, params_list, y0_template, t_span, observation_node, - ) - - mask = (success_arr > 0.5) & (cmax_arr > 0) - cmax_arr = cmax_arr[mask] - tmax_arr = tmax_arr[mask] - auc_arr = auc_arr[mask] - n_ok = int(np.sum(mask)) - n_failures = n_samples - n_ok - - logger.info("MC-fast JAX: %d/%d successful", n_ok, n_samples) - return self._build_mc_result(cmax_arr, tmax_arr, auc_arr, n_ok, n_failures) - # ── SciPy backend: sequential loop (default) ── cmax_samples: list[float] = [] tmax_samples: list[float] = [] diff --git a/tests/unit/test_active_transport.py b/tests/unit/test_active_transport.py index 70c3fe47..1a992fbe 100644 --- a/tests/unit/test_active_transport.py +++ b/tests/unit/test_active_transport.py @@ -8,7 +8,6 @@ from __future__ import annotations import numpy as np -import pytest import sisyphus.engine.flux # noqa: F401 — registers all FluxSpec implementations from sisyphus.core import Distribution, DrugOnGraph, TransporterKinetics @@ -135,32 +134,6 @@ def test_active_transport_uses_target_ivive_and_transporters(): assert abs(transport_contrib_src + transport_contrib_tgt) < 1e-12 -def test_active_transport_scipy_jax_parity(): - """SciPy and JAX backends produce the same dydt for an active-transport edge.""" - pytest.importorskip("jax") - import jax.numpy as jnp - - from sisyphus.engine.params_jax import resolve_to_jax - from sisyphus.engine.rhs_jax import make_jax_rhs - - graph = _minimal_uptake_graph() - drug = _drug_with_oatp() - - compiled = ODECompiler().compile(graph) - params_scipy = ResolvedParams(graph, drug) - rhs_scipy = compiled.make_rhs(params_scipy) - - y = np.zeros(compiled.n_states) - y[compiled.state_index["src_blood"]] = 10.0 - dydt_scipy = rhs_scipy(0.0, y) - - jax_params = resolve_to_jax(compiled, params_scipy) - rhs_jax_fn = make_jax_rhs(compiled) - dydt_jax = np.asarray(rhs_jax_fn(0.0, jnp.asarray(y), jax_params)) - - np.testing.assert_allclose(dydt_jax, dydt_scipy, rtol=1e-6, atol=1e-9) - - def test_active_transport_mass_conservation(): """Sum of source + target dydt contributions from active transport is zero.""" graph = _minimal_uptake_graph() diff --git a/tests/unit/test_active_transport_direction.py b/tests/unit/test_active_transport_direction.py index dca9135f..f5bf12e9 100644 --- a/tests/unit/test_active_transport_direction.py +++ b/tests/unit/test_active_transport_direction.py @@ -61,24 +61,3 @@ def test_uptake_default_unchanged(): dydt = rhs(0.0, y) # Uptake reads target (liver) transporters; mass moves blood → liver via transport. assert dydt[compiled.state_index["liver"]] > 0 - - -def test_efflux_scipy_jax_parity(): - import pytest - pytest.importorskip("jax") - import jax.numpy as jnp - - from sisyphus.engine.params_jax import resolve_to_jax - from sisyphus.engine.rhs_jax import make_jax_rhs - - g = _efflux_graph() - compiled = ODECompiler().compile(g) - params = ResolvedParams(g, _drug()) - rhs_scipy = compiled.make_rhs(params) - y = np.zeros(compiled.n_states) - y[compiled.state_index["gut_wall"]] = 10.0 - dydt_scipy = rhs_scipy(0.0, y) - - jax_params = resolve_to_jax(compiled, params) - dydt_jax = np.asarray(make_jax_rhs(compiled)(0.0, jnp.asarray(y), jax_params)) - np.testing.assert_allclose(dydt_jax, dydt_scipy, rtol=1e-6, atol=1e-9) diff --git a/tests/unit/test_jax_scipy_parity.py b/tests/unit/test_jax_scipy_parity.py deleted file mode 100644 index 76a898f2..00000000 --- a/tests/unit/test_jax_scipy_parity.py +++ /dev/null @@ -1,179 +0,0 @@ -"""WS-4: JAX↔SciPy RHS-level parity per flux branch.""" -from __future__ import annotations - -import numpy as np -import pytest - -pytest.importorskip("jax") -import jax.numpy as jnp # noqa: E402 - -import sisyphus.engine.flux # noqa: F401,E402 -from sisyphus.core import Distribution, DrugOnGraph # noqa: E402 -from sisyphus.engine.compiler import ODECompiler, ResolvedParams # noqa: E402 -from sisyphus.engine.params_jax import resolve_to_jax # noqa: E402 -from sisyphus.engine.rhs_jax import make_jax_rhs # noqa: E402 -from sisyphus.graph.body import BodyGraph # noqa: E402 -from sisyphus.graph.types import ( # noqa: E402 - AbsorptionEdge, - ClearanceEdge, - DiffusionEdge, - FlowEdge, - Node, - TransitEdge, -) - - -def _drug(fu_corr: float = 1.0) -> DrugOnGraph: - return DrugOnGraph( - name="d", smiles="CCO", dose_mg=100.0, route="oral", - administration_node="blood", mw=300.0, pka=4.5, compound_type="acid", - fup=Distribution(0.3), rbp=Distribution(1.0), kp_method="provided", - kp_overrides={"organ": Distribution(1.0)}, peff=Distribution(1.0), - solubility=Distribution(1.0), enzyme_affinity={"CYP3A4": Distribution(5.0)}, - renal_clearance=Distribution(0.0), fu_correction_liver=Distribution(fu_corr), - ) - - -def _ws_flagged_graph(flagged: bool = True) -> BodyGraph: - g = BodyGraph() - g.add_node(Node(name="blood", node_type="blood_pool", volume=Distribution(5.0))) - g.add_node(Node(name="organ", node_type="organ", volume=Distribution(1.0), - enzymes={"CYP3A4": Distribution(1.0e6)}, ivive_scaling=1.0e-6, - fu_correction_applicable=1.0 if flagged else 0.0, lookup_name="organ")) - g.add_node(Node(name="sink", node_type="sink", volume=Distribution(1.0))) - g.add_edge(FlowEdge(source="blood", target="organ", flow_rate=Distribution(10.0))) - g.add_edge(FlowEdge(source="organ", target="blood", flow_rate=Distribution(10.0))) - g.add_edge(ClearanceEdge(source="organ", target="sink", model="well_stirred")) - return g - - -def _rhs_pair(graph, drug): - compiled = ODECompiler().compile(graph) - params = ResolvedParams(graph, drug) - y = np.zeros(compiled.n_states) - y[compiled.state_index["organ"]] = 3.0 - y[compiled.state_index["blood"]] = 7.0 - dydt_scipy = compiled.make_rhs(params)(0.0, y) - dydt_jax = np.asarray( - make_jax_rhs(compiled)(0.0, jnp.asarray(y), resolve_to_jax(compiled, params)) - ) - return dydt_scipy, dydt_jax - - -def test_ws_fu_correction_parity(): - """Flagged well_stirred node with non-identity fu_correction: JAX == SciPy.""" - s, j = _rhs_pair(_ws_flagged_graph(), _drug(fu_corr=1.5)) - np.testing.assert_allclose(j, s, rtol=1e-9, atol=1e-12) - - -def test_jax_fails_loud_on_distinct_km_multitransporter(): - """≥2 active transporters with distinct Km at one node: the JAX aggregate - Vmax/weighted-Km approximation diverges from SciPy → fail loud, not silent.""" - from sisyphus.core import TransporterKinetics - from sisyphus.graph.types import ActiveTransportEdge - - g = BodyGraph() - g.add_node(Node(name="blood", node_type="blood_pool", volume=Distribution(1.0))) - g.add_node(Node(name="organ", node_type="organ", volume=Distribution(1.0), - transporters={"A": Distribution(1e10), "B": Distribution(1e10)}, - ivive_scaling=1e-4)) - g.add_edge(FlowEdge(source="blood", target="organ", flow_rate=Distribution(10.0))) - g.add_edge(FlowEdge(source="organ", target="blood", flow_rate=Distribution(10.0))) - g.add_edge(ActiveTransportEdge(source="blood", target="organ")) - drug = DrugOnGraph( - name="d", smiles="CCO", dose_mg=100.0, route="oral", administration_node="blood", - mw=300.0, pka=4.5, compound_type="acid", fup=Distribution(0.3), rbp=Distribution(1.0), - kp_method="provided", kp_overrides={"organ": Distribution(1.0)}, peff=Distribution(1.0), - solubility=Distribution(1.0), enzyme_affinity={}, renal_clearance=Distribution(0.0), - transporter_kinetics={ - "A": TransporterKinetics(jmax=Distribution(100.0), km=Distribution(10.0)), - "B": TransporterKinetics(jmax=Distribution(100.0), km=Distribution(80.0)), - }, - ) - compiled = ODECompiler().compile(g) - with pytest.raises(NotImplementedError, match="distinct Km"): - resolve_to_jax(compiled, ResolvedParams(g, drug)) - - -# --------------------------------------------------------------------------- -# Per-branch RHS-level parity (SciPy make_rhs == make_jax_rhs at fixed t, y). -# Each test builds a minimal valid graph exercising ONLY one flux branch and -# asserts dydt agreement to rtol=1e-9, atol=1e-12. RHS-level (not integrated -# trajectories): isolates flux math from integrator differences. -# --------------------------------------------------------------------------- - - -def _assert_parity(graph, drug, fill): - compiled = ODECompiler().compile(graph) - params = ResolvedParams(graph, drug) - y = np.zeros(compiled.n_states) - for name, amt in fill.items(): - y[compiled.state_index[name]] = amt - s = compiled.make_rhs(params)(0.0, y) - j = np.asarray(make_jax_rhs(compiled)(0.0, jnp.asarray(y), resolve_to_jax(compiled, params))) - np.testing.assert_allclose(j, s, rtol=1e-9, atol=1e-12) - - -def test_ws_no_fu_correction_parity(): - g = _ws_flagged_graph(flagged=False) # well_stirred WITHOUT the fu_correction flag - _assert_parity(g, _drug(fu_corr=1.0), {"organ": 3.0, "blood": 7.0}) - - -def test_gfr_parity(): - import dataclasses - g = BodyGraph() - g.add_node(Node(name="blood", node_type="blood_pool", volume=Distribution(5.0))) - g.add_node(Node(name="kidney", node_type="organ", volume=Distribution(1.0), - lookup_name="kidney")) - g.add_node(Node(name="urine", node_type="sink", volume=Distribution(1.0))) - g.add_edge(FlowEdge(source="blood", target="kidney", flow_rate=Distribution(10.0))) - g.add_edge(FlowEdge(source="kidney", target="blood", flow_rate=Distribution(10.0))) - g.add_edge(ClearanceEdge(source="kidney", target="urine", model="gfr_filtration")) - drug = dataclasses.replace(_drug(), renal_clearance=Distribution(5.0)) - _assert_parity(g, drug, {"kidney": 2.0, "blood": 8.0}) - - -def test_flow_parity(): - # FlowEdge between a blood pool (C_out = A/V) and an organ - # (C_out = A*RBP/(V*Kp)); both convective forms exercised, balanced flow. - g = BodyGraph() - g.add_node(Node(name="blood", node_type="blood_pool", volume=Distribution(5.0))) - g.add_node(Node(name="organ", node_type="organ", volume=Distribution(1.0), - lookup_name="organ")) - g.add_edge(FlowEdge(source="blood", target="organ", flow_rate=Distribution(10.0))) - g.add_edge(FlowEdge(source="organ", target="blood", flow_rate=Distribution(10.0))) - _assert_parity(g, _drug(), {"organ": 3.0, "blood": 7.0}) - - -def test_transit_parity(): - # First-order transit (rate = k * A_source) between two flow-exempt - # lumen nodes — no flow-conservation requirement. - g = BodyGraph() - g.add_node(Node(name="lumen1", node_type="lumen", volume=Distribution(1.0))) - g.add_node(Node(name="lumen2", node_type="lumen", volume=Distribution(1.0))) - g.add_edge(TransitEdge(source="lumen1", target="lumen2", transit_rate=Distribution(0.5))) - _assert_parity(g, _drug(), {"lumen1": 4.0}) - - -def test_absorption_parity(): - # ka = 2.88 * peff * ka_fraction / particle_radius_um; _drug() has - # peff=1.0 and the default radius=25.0 → ka>0. Lumen source is flow-exempt; - # the organ target has no flow edges, so conservation is satisfied trivially. - g = BodyGraph() - g.add_node(Node(name="lumen", node_type="lumen", volume=Distribution(1.0))) - g.add_node(Node(name="gut", node_type="organ", volume=Distribution(1.0), - lookup_name="gut")) - g.add_edge(AbsorptionEdge(source="lumen", target="gut", ka_fraction=Distribution(1.0))) - _assert_parity(g, _drug(), {"lumen": 5.0}) - - -def test_diffusion_parity(): - # PS-limited exchange: flux = PS * (fup*C_vasc/RBP - fup*C_tissue/Kp). - # ps_product on the edge (no drug ps_override) drives the flux; mass in - # vasc with the tissue defaulting Kp=1.0 yields a nonzero gradient. - g = BodyGraph() - g.add_node(Node(name="vasc", node_type="blood_pool", volume=Distribution(2.0))) - g.add_node(Node(name="tissue", node_type="organ", volume=Distribution(1.0), - lookup_name="tissue")) - g.add_edge(DiffusionEdge(source="vasc", target="tissue", ps_product=Distribution(20.0))) - _assert_parity(g, _drug(), {"vasc": 6.0, "tissue": 1.0}) diff --git a/tests/unit/test_rhs_jax_unsupported_flux.py b/tests/unit/test_rhs_jax_unsupported_flux.py deleted file mode 100644 index 25633896..00000000 --- a/tests/unit/test_rhs_jax_unsupported_flux.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Audit follow-up: the JAX RHS must not silently drop unsupported flux types. - -``ProdrugActivationFluxSpec`` and ``OneCompartmentEliminationFluxSpec`` have no -branch in ``make_jax_rhs``'s isinstance chain. Before the -``_unsupported_flux_specs`` guard they fell through the build loop silently, so a -graph with a prodrug-activation or active-metabolite edge would have that flux -omitted from the JAX RHS with no error (the SciPy backend handles them -correctly). These tests pin a loud failure instead of silent corruption. - -The guard is pure Python (no JAX import) so it runs even though JAX is optional -and absent from ``requirements-lock.txt``. -""" -from __future__ import annotations - -from sisyphus.engine.flux import ( - FlowFluxSpec, - OneCompartmentEliminationFluxSpec, - ProdrugActivationFluxSpec, -) -from sisyphus.engine.rhs_jax import _unsupported_flux_specs - - -def _prodrug() -> ProdrugActivationFluxSpec: - return ProdrugActivationFluxSpec( - 0, 0, 1, "parent", "active", frozenset({"CYP3A4"}), 1.0 - ) - - -def _one_compartment() -> OneCompartmentEliminationFluxSpec: - return OneCompartmentEliminationFluxSpec(0, 0, 1, "active", "sink") - - -def test_flags_prodrug_activation(): - spec = _prodrug() - assert _unsupported_flux_specs([spec]) == [spec] - - -def test_flags_one_compartment_elimination(): - spec = _one_compartment() - assert _unsupported_flux_specs([spec]) == [spec] - - -def test_supported_flux_not_flagged(): - flow = FlowFluxSpec(0, 0, 1, "a", "b") - assert _unsupported_flux_specs([flow]) == [] - - -def test_mixed_returns_only_unsupported(): - flow = FlowFluxSpec(0, 0, 1, "a", "b") - prod = _prodrug() - one = _one_compartment() - out = _unsupported_flux_specs([flow, prod, one]) - assert flow not in out - assert prod in out - assert one in out diff --git a/tests/unit/test_solver_jax.py b/tests/unit/test_solver_jax.py deleted file mode 100644 index 81a42341..00000000 --- a/tests/unit/test_solver_jax.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Tests for JAX/Diffrax solver — Phase 0 of UDE roadmap. - -Validates: -1. Numerical equivalence to SciPy LSODA on the wrapped path. -2. Autograd support through Diffrax Kvaerno5 on pure-JAX path. -""" - -from __future__ import annotations - -import os -from pathlib import Path - -import numpy as np -import pytest - -os.environ.setdefault("JAX_PLATFORMS", "cpu") - -# Skip entire module if jax/diffrax not installed (optional dependency). -pytest.importorskip("jax") -pytest.importorskip("diffrax") - -import jax # noqa: E402 - -jax.config.update("jax_enable_x64", True) -import jax.numpy as jnp # noqa: E402 - -from sisyphus.engine.solver_jax import solve as solve_jax # noqa: E402 -from sisyphus.engine.solver_jax import solve_jax_pure # noqa: E402 - -# --------------------------------------------------------------------------- -# Phase 0b: pure-JAX RHS + gradient verification on 2-compartment model -# --------------------------------------------------------------------------- - - -def _two_compartment_rhs(t, y, args): - """2-compartment oral absorption: A_gut → A_central → sink.""" - ka, CL, V = args - A_gut, A_cent = y - return jnp.array([-ka * A_gut, ka * A_gut - (CL / V) * A_cent]) - - -def _analytical_cmax(dose: float, ka: float, CL: float, V: float) -> float: - """Closed-form Cmax for 1st-order oral absorption, 1-compartment elim.""" - k_el = CL / V - tmax = float(np.log(ka / k_el) / (ka - k_el)) - return float(dose / V * ka / (ka - k_el) * (np.exp(-k_el * tmax) - np.exp(-ka * tmax))) - - -def _cmax_fn(args, dose: float): - """Compute Cmax via Diffrax forward solve.""" - y0 = jnp.array([dose, 0.0]) - t_eval = jnp.linspace(0.01, 24.0, 241) - _, ys, _ = solve_jax_pure( - _two_compartment_rhs, y0, (0.0, 24.0), args=args, t_eval=t_eval, - allow_jvp=True, stiff=False, - ) - C_cent = ys[:, 1] / args[2] - return jnp.max(C_cent) - - -class TestJAXSolverForward: - """Forward solve accuracy against analytical solution.""" - - def test_two_compartment_matches_analytical(self): - dose, ka, CL, V = 10.0, 1.5, 5.0, 50.0 - cmax_analytical = _analytical_cmax(dose, ka, CL, V) - cmax_diffrax = float(_cmax_fn((ka, CL, V), dose)) - rel_err = abs(cmax_diffrax - cmax_analytical) / cmax_analytical - assert rel_err < 1e-3, f"rel_err={rel_err:.3e}, expected < 1e-3" - - def test_high_clearance_drug(self): - """High CL drug: short t½, low Cmax.""" - dose, ka, CL, V = 100.0, 2.0, 100.0, 70.0 - cmax_an = _analytical_cmax(dose, ka, CL, V) - cmax_dx = float(_cmax_fn((ka, CL, V), dose)) - assert abs(cmax_dx - cmax_an) / cmax_an < 1e-3 - - def test_low_clearance_drug(self): - """Low CL drug: long t½, plateau behavior.""" - dose, ka, CL, V = 50.0, 0.5, 2.0, 100.0 - cmax_an = _analytical_cmax(dose, ka, CL, V) - cmax_dx = float(_cmax_fn((ka, CL, V), dose)) - assert abs(cmax_dx - cmax_an) / cmax_an < 1e-3 - - -class TestJAXSolverGradients: - """Verify jax.grad through Diffrax matches finite differences.""" - - @staticmethod - def _fd(f, x, eps=1e-5): - return (f(x + eps) - f(x - eps)) / (2 * eps) - - def test_grad_clearance(self): - dose, ka, CL, V = 10.0, 1.5, 5.0, 50.0 - g = jax.grad(lambda cl: _cmax_fn((ka, cl, V), dose))(CL) - fd = self._fd(lambda cl: float(_cmax_fn((ka, float(cl), V), dose)), CL) - assert abs(float(g) - fd) / abs(fd) < 1e-3, f"grad={float(g):.6f} fd={fd:.6f}" - - def test_grad_ka(self): - dose, ka, CL, V = 10.0, 1.5, 5.0, 50.0 - g = jax.grad(lambda k: _cmax_fn((k, CL, V), dose))(ka) - fd = self._fd(lambda k: float(_cmax_fn((float(k), CL, V), dose)), ka) - assert abs(float(g) - fd) / abs(fd) < 1e-3 - - def test_grad_volume(self): - dose, ka, CL, V = 10.0, 1.5, 5.0, 50.0 - g = jax.grad(lambda v: _cmax_fn((ka, CL, v), dose))(V) - fd = self._fd(lambda v: float(_cmax_fn((ka, CL, float(v)), dose)), V) - assert abs(float(g) - fd) / abs(fd) < 1e-3 - - -# --------------------------------------------------------------------------- -# Phase 0a: numerical equivalence between SciPy LSODA and Diffrax (wrapped) -# --------------------------------------------------------------------------- - - -class TestJAXSolverMatchesSciPy: - """Validate solve_jax wrapped path matches existing SciPy-based solver.""" - - @pytest.fixture(scope="class") - def setup(self): - import sisyphus.engine.flux # noqa: F401 -- register flux specs - from sisyphus.engine.compiler import ODECompiler - from sisyphus.graph.builder import build_from_yaml - - physiology = Path(__file__).resolve().parent.parent.parent / "data" / "physiology" / "reference_man.yaml" # noqa: E501 - graph = build_from_yaml(physiology) - compiler = ODECompiler() - compiled = compiler.compile(graph) - return graph, compiled - - def _run_both(self, setup, smiles, dose, route="oral"): - from sisyphus.engine.compiler import ResolvedParams - from sisyphus.engine.solver import solve as solve_scipy - from sisyphus.predict.adme import predict_adme - from sisyphus.predict.chemistry import compute_profile - from sisyphus.predict.ivive import build_drug_on_graph - - graph, compiled = setup - profile = compute_profile(smiles) - adme = predict_adme(profile) - drug = build_drug_on_graph(profile, adme, dose, route) - rng = np.random.default_rng(42) - params = ResolvedParams(graph.sample(rng), drug.sample(rng)) - y0 = np.zeros(compiled.n_states) - y0[compiled.state_index[drug.administration_node]] = dose - sol_s = solve_scipy(compiled, params, y0, t_span=(0, 24)) - sol_j = solve_jax(compiled, params, y0, t_span=(0, 24)) - return sol_s, sol_j - - @pytest.mark.slow - def test_midazolam_venous_blood(self, setup): - smi = "Cc1ncc2n1-c1ccc(Cl)cc1C(c1ccccc1F)=NC2" - sol_s, sol_j = self._run_both(setup, smi, 15.0) - c_s = sol_s.concentrations["venous_blood"] - c_j = sol_j.concentrations["venous_blood"] - rel_err = np.abs(c_s - c_j).max() / max(c_s.max(), 1e-10) - assert rel_err < 1e-3, f"venous_blood rel_err={rel_err:.3e}" - assert sol_j.solver_success - - @pytest.mark.slow - def test_warfarin_mass_balance(self, setup): - smi = "CC(=O)CC(c1ccccc1)c1c(O)c2ccccc2oc1=O" - _, sol_j = self._run_both(setup, smi, 10.0) - # Mass balance must be tight (same threshold as SciPy path) - assert sol_j.mass_balance_error < 1e-4, f"MBE={sol_j.mass_balance_error:.3e}" - assert sol_j.solver_success