Support source-only (from: null) transitions in reaction artifacts - #195
Merged
Conversation
build_reaction_artifacts_ir previously excluded every from: null transition from CompiledRhs.reactions -- the module docstring's stated reason was that a stochastic/CTMC consumer needs "how many independent source cells are firing, and where does each firing land", which a source-only transition supposedly didn't have a well-defined answer for. That reasoning doesn't actually hold: a source-only transition (an exogenous hazard with no compartment to deplete, e.g. cross-district case importation) has a perfectly well-defined firing-cell count -- one independent Poisson process per DESTINATION cell, with no source population to bound it against. - ReactionArtifactIR.from_base / CompiledReaction.from_base are now `str | None`; `None` marks a source-only reaction. - For a source-only reaction, from_axes/full_axes are taken from the to-side template (there is no from-side template to speak of), and the propensity is the bare rate itself (no from_state multiplication -- there's no source state). - _make_propensity_fn gained an explicit broadcast_shape param, used only for source-only reactions: lower_to_vector_ast's own contract is "broadcastABLE against target_axes", not "exactly shaped like target_axes" -- true by construction for a normal reaction (the rate*from_state multiply always broadcasts to from_state's full shape), but not guaranteed for a source-only reaction's bare-rate propensity when the rate doesn't reference every to-side axis (e.g. an age-independent importation hazard). Confirmed via a failing test before this fix: the compiled propensity came back under-shaped relative to its own declared from_axes. Consumers must still special-case from_base is None (skip depletion/ clamp steps) -- diphtheria_outbreakvacc's run_hybrid_ctmc updated accordingly in a companion change. Tests: replaced the old "excluded" assertion with coverage for inclusion, metadata, propensity correctness (including the broadcast fix), and agreement with the deterministic eval_fn oracle (same pattern already used for the from-pinned and alias-inlining fixes). Full suite (458 tests) and mypy --strict pass.
4 tasks
jc-macdonald
added a commit
that referenced
this pull request
Aug 26, 2026
_build_reaction_artifacts compiles each named transition's propensity independently of _wrap_eval_fn_for_time_varying / _wrap_pytree_eval_fn_for_time_varying (the two wrappers that make a rate like `rate: "lambda_import[time, loc]"` work for the deterministic RHS by interpolating the raw (time, loc) grid down to the current timestep's (loc,) slice before the compiled code runs). Reaction propensity_fns never went through either wrapper, so the same rate reached a reaction's compiled code as the raw, un-interpolated full grid array -- a shape mismatch against what that code actually expects (confirmed via a failing call before this fix: "cannot reshape array of size N into shape (...)"). Adds _wrap_propensity_fn_for_time_varying, mirroring the existing two wrappers' exact interpolation logic and applied to every compiled reaction in _build_reaction_artifacts. Motivating use case: diphtheria_outbreakvacc's cross-district case-importation term needs a time-varying hazard (concentrated in an initial window, much lower afterward) that's also addressable as a CTMC/tau-leap reaction via run_hybrid_ctmc -- this was blocked without this fix even after PR #195 (source-only transitions), since the two gaps are independent (this one applies to ANY reaction, source-only or not). Test: new test_time_varying_rate_propensity_matches_deterministic, cross-checking the reaction's propensity against the deterministic eval_fn's own (already-correct) interpolated inflow at a fractional, off-grid t -- same correctness-oracle pattern as the existing test_reactions_reconstruct_deterministic_eval_fn. Full suite (457 tests) and mypy --strict pass.
jc-macdonald
added a commit
that referenced
this pull request
Aug 26, 2026
_build_reaction_artifacts compiles each named transition's propensity independently of _wrap_eval_fn_for_time_varying / _wrap_pytree_eval_fn_for_time_varying (the two wrappers that make a rate like `rate: "lambda_import[time, loc]"` work for the deterministic RHS by interpolating the raw (time, loc) grid down to the current timestep's (loc,) slice before the compiled code runs). Reaction propensity_fns never went through either wrapper, so the same rate reached a reaction's compiled code as the raw, un-interpolated full grid array -- a shape mismatch against what that code actually expects (confirmed via a failing call before this fix: "cannot reshape array of size N into shape (...)"). Adds _wrap_propensity_fn_for_time_varying, mirroring the existing two wrappers' exact interpolation logic and applied to every compiled reaction in _build_reaction_artifacts. Motivating use case: diphtheria_outbreakvacc's cross-district case-importation term needs a time-varying hazard (concentrated in an initial window, much lower afterward) that's also addressable as a CTMC/tau-leap reaction via run_hybrid_ctmc -- this was blocked without this fix even after PR #195 (source-only transitions), since the two gaps are independent (this one applies to ANY reaction, source-only or not). Test: new test_time_varying_rate_propensity_matches_deterministic, cross-checking the reaction's propensity against the deterministic eval_fn's own (already-correct) interpolated inflow at a fractional, off-grid t -- same correctness-oracle pattern as the existing test_reactions_reconstruct_deterministic_eval_fn. Full suite (457 tests) and mypy --strict pass.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
build_reaction_artifacts_irpreviously excluded everyfrom: null(source-only) transition from
CompiledRhs.reactions. The moduledocstring's stated reason was that a stochastic/CTMC consumer needs "how
many independent source cells are firing, and where does each firing
land," which a source-only transition supposedly didn't have a
well-defined answer for.
That reasoning doesn't actually hold up: a source-only transition (an
exogenous hazard with no compartment to deplete -- e.g. cross-district
disease-case importation, immigration into a population, etc.) has a
perfectly well-defined firing-cell count: one independent Poisson process
per DESTINATION cell, with no source population to bound it against. This
PR adds that support.
Motivating use case:
diphtheria_outbreakvaccneeds a genuinelystochastic cross-district case-importation term
(
{to: "E[age, vax, loc]", rate: "lambda_import[loc]"}) for itsvaccination-campaign scenario grid work. The deterministic RHS path
(
_normalize.py) already supportsfrom: nullcorrectly; only thereaction-artifact path (for
run_hybrid_ctmc/discrete-CTMC consumers)didn't.
Changes
ReactionArtifactIR.from_base/CompiledReaction.from_baseare nowstr | None;Nonemarks a source-only reaction.from_axes/full_axesare taken from theto-side template (there is nofrom-side template to speak of), andthe propensity is the bare rate itself (no
from_statemultiplication-- there's no source state).
_make_propensity_fngained an explicitbroadcast_shapeparam, usedonly for source-only reactions:
lower_to_vector_ast's own contract is"broadcastABLE against
target_axes," not "exactly shaped liketarget_axes" -- true by construction for a normal reaction (therate*from_statemultiply always broadcasts tofrom_state's fullshape), but not guaranteed for a source-only reaction's bare-rate
propensity when the rate doesn't reference every to-side axis (e.g. an
age-independent importation hazard). Caught this via a failing test
before the fix: the compiled propensity came back under-shaped relative
to its own declared
from_axes.Consumers must still special-case
from_base is None(skipdepletion/clamp steps) --
diphtheria_outbreakvacc'srun_hybrid_ctmcupdated accordingly in a companion change on that side.
Test plan
test_source_only_transition_excluded_from_reactions_ir)with coverage for inclusion, metadata, propensity correctness
(including the broadcast fix), and agreement with the deterministic
eval_fnoracle -- same pattern already used for the from-pinnedand alias-inlining fixes in this file.
mypy --stricton both changed source files: cleanruff check: no new findings beyond this project's existingpreview-rule
noqa-vs-ruff:ignorestyle note, which firesuniformly across the whole file (pre-existing, not introduced here)