Fix reaction propensities that reference a kernel via apply_along - #187
Merged
Conversation
_build_reaction_artifacts compiled propensity_ir_full -- the pre-expanded (via expand_reduce_pointwise, empty lhs_assignment) template-symbolic form. For a propensity whose rate reduces over an axis coinciding with the target's own free axis (e.g. a spatial force-of-infection term, apply_along(k[age,age:ap]*S[age:ap], age=ap)), that expansion produces malformed AxisIndex(axis='', ...) subscripts, and compiling them raised "axes don't match array". Root cause, confirmed empirically: the deterministic equation path never hits this, because it compiles Reduce-bearing expressions via the Reduce-preserving form (propensity_ir_reduce's equivalent) through lower_to_vector_ast's native Reduce-node lowering -- not through expand_reduce_pointwise's template-symbolic expansion. Verified by compiling the exact same rate expression as a normal (non-reaction) transition and confirming it produces correct deterministic output, then switching _build_reaction_artifacts to compile propensity_ir_reduce instead of propensity_ir_full and confirming both the propensity value and its cross-check against the deterministic path now match a hand-computed expectation for a real (non-degenerate) kernel contraction. New regression test (test_reduce_bearing_kernel_rate_propensity): covers exactly this case, both propensity correctness and the deterministic-path cross-check. Found while testing the flepimop2-op_system adapter's reactions wiring against a spatial force-of-infection-style rate -- the pattern the diphtheria_outbreakvacc model's own S -> E transition uses, which would have silently produced no working reaction artifact for that transition (or a confusing runtime error) once named for stochastic treatment. Verified: full test suite (451 = 450 + 1 new) passes; mypy --strict clean; ruff clean (one genuinely-introduced unused-noqa, now removed; remaining findings confirmed pre-existing repo-wide noqa-comments style noise, unrelated).
3 tasks
jc-macdonald
added a commit
that referenced
this pull request
Aug 25, 2026
The kernel-reduction case (apply_along over a mixing kernel, matching a spatial force-of-infection term) needed the fix in #187 (propensity_ir_reduce instead of propensity_ir_full) to compile at all. Left as a documented known-limitation comment when this branch was first opened, since #187 hadn't merged yet and the test would have been red on this PR alone. Now that #187 is merged into main and this branch is rebased onto it, restore the real test: test_option_reactions_merges_mixing_kernels exercises the fix end to end through the adapter (mixing_kernels merged into propensity_fn params the same way the other steppers already do), not just at the op_system compile layer. Verified: flepimop2-op_system full suite (59 = 58 + 1) passes; ruff and mypy clean on the touched lines.
jc-macdonald
added a commit
that referenced
this pull request
Aug 25, 2026
* Expose CompiledRhs.reactions via the flepimop2-op_system adapter PR #185 added CompiledRhs.reactions but the flepimop2-op_system adapter (OpSystemSystem) builds a curated options dict listing specific compiled fields by name -- it doesn't auto-forward new CompiledRhs fields, so .reactions was compiled but unreachable by any engine plugin via stepper.option("reactions"). Also exports CompiledReaction/ReactionPropensityFn from op_system's public __init__ (missed in #185 -- they existed on compile.py but weren't re-exported, same gap CompiledRhs/compile_spec already have going the other way, see the pre-existing mypy note below). Adds a _make_reaction_steppers helper mirroring _maybe_make_pytree_stepper's exact pattern: wraps each CompiledReaction.propensity_fn so mixing_kernels get merged into params the same way every other stepper already does (via the shared OpSystemSystem._merged_params), returning a same-shape tuple of CompiledReaction with only propensity_fn replaced (dataclasses.replace). Wired into options["reactions"] alongside the other stepper options. Found and documented (not fixed, out of scope here) a real limitation while testing: a transition rate that references a mixing kernel through an apply_along reduction (e.g. a spatial force-of-infection term) isn't supported by the v1 reaction-artifact scope -- _compile_ir_expr raises "axes don't match array" trying to lower a Reduce-bearing propensity against a from-axes-only target shape. Rates that only reference from-side axis-indexed shaped params work fine. Left a detailed comment in tests/test_system.py where this was found; opening a tracking issue against op_system for it. Also noted, not touched: op_system's own __init__ re-export of CompiledRhs/compile_spec (now also CompiledReaction/ReactionPropensityFn) doesn't satisfy mypy's explicit-reexport check from a consuming package's perspective -- confirmed pre-existing (2 of the resulting 3 errors already existed before this commit, via git stash) as a harmless mismatch between this __all__-based re-export pattern and mypy's stricter check; not fixed here since a real fix touches op_system's export mechanism broadly, outside this PR's scope. Verified: flepimop2-op_system's full test suite (58 tests, 4 new) passes; ruff clean on all touched lines (pre-existing repo-wide noqa-comments style noise confirmed unrelated, same as op_system proper); mypy clean on all touched lines (pre-existing unrelated errors elsewhere in both touched files confirmed via git stash before/after comparison, not introduced here). Manually verified end-to-end: OpSystemSystem(spec=...).options["reactions"] surfaces a working, correctly-valued propensity_fn for a named collapse-to- pinned-target transition, matching the diphtheria model's real usage shape. * Restore reactions/mixing_kernels confirming test now that #187 is merged The kernel-reduction case (apply_along over a mixing kernel, matching a spatial force-of-infection term) needed the fix in #187 (propensity_ir_reduce instead of propensity_ir_full) to compile at all. Left as a documented known-limitation comment when this branch was first opened, since #187 hadn't merged yet and the test would have been red on this PR alone. Now that #187 is merged into main and this branch is rebased onto it, restore the real test: test_option_reactions_merges_mixing_kernels exercises the fix end to end through the adapter (mixing_kernels merged into propensity_fn params the same way the other steppers already do), not just at the op_system compile layer. Verified: flepimop2-op_system full suite (59 = 58 + 1) passes; ruff and mypy clean on the touched lines.
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.
Closes #186.
Summary
_build_reaction_artifactscompiledpropensity_ir_full— the pre-expanded (viaexpand_reduce_pointwise, emptylhs_assignment) template-symbolic form. For a propensity whose rate reduces over an axis coinciding with the target's own free axis (e.g. a spatial force-of-infection term:apply_along(k[age, age:ap] * S[age:ap], age=ap)), that expansion produces malformedAxisIndex(axis='', ...)subscripts, and compiling them raisedValueError: axes don't match array.Root cause
Confirmed empirically, not just by inspection: the deterministic equation path never hits this, because it compiles Reduce-bearing expressions via the Reduce-preserving form through
lower_to_vector_ast's native Reduce-node lowering — not throughexpand_reduce_pointwise's template-symbolic expansion. Verified by:_build_reaction_artifactsto compilepropensity_ir_reduceinstead ofpropensity_ir_full.* 0sanity check that happened to cancel the broken term.Fix
One-line change: compile
r.propensity_ir_reduceinstead ofr.propensity_ir_fullin_build_reaction_artifacts. No other code paths touched.Why this matters for the motivating use case
Found while testing
flepimop2-op_system's newreactionsadapter wiring against a spatial force-of-infection-style rate — the exact patterndiphtheria_outbreakvacc's ownS -> Etransition uses. Without this fix, that transition would have silently produced no working reaction artifact (or a confusing runtime error at whatever point downstream a consumer first calledpropensity_fn) once named for stochastic treatment.Verification
test_reduce_bearing_kernel_rate_propensityintests/op_system/test_op_system_reactions.py: covers both propensity correctness and the deterministic-path cross-check for a real kernel contraction.mypy --strictclean.ruffclean (one genuinely-introducedunused-noqafrom a no-longer-needed suppression, removed; remaining findings confirmed pre-existing repo-widenoqa-commentsstyle noise, unrelated to this change).Test plan
pytest— 451 passedmypy --strict— cleanruff check— clean (modulo pre-existing noise)pytree_eval_fn, for a real (non-degenerate) kernel matrix