Expose CompiledRhs.reactions via the flepimop2-op_system adapter - #188
Merged
Conversation
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.
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
force-pushed
the
feature/expose-reactions-via-adapter
branch
from
August 25, 2026 18:19
86b9aff to
b2d0296
Compare
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
PR #185 added
CompiledRhs.reactions, but theflepimop2-op_systemadapter (OpSystemSystem) builds a curatedoptionsdict listing specific compiled fields by name — it doesn't auto-forward newCompiledRhsfields, so.reactionswas compiled but unreachable by any engine plugin viastepper.option("reactions"). This PR closes that gap.What changed
op_system/__init__.py: exportsCompiledReaction/ReactionPropensityFnfrom the public API (missed in Expose per-transition compiled propensity + stoichiometry artifact on CompiledRhs #185 — they existed oncompile.pybut weren't re-exported, the same gapCompiledRhs/compile_specalready have solved the other way).flepimop2-op_system: new_make_reaction_steppershelper mirroring_maybe_make_pytree_stepper's exact pattern — wraps eachCompiledReaction.propensity_fnsomixing_kernelsget merged into params the same way every other stepper already does (via the sharedOpSystemSystem._merged_params), returning a same-shape tuple ofCompiledReactionwith onlypropensity_fnreplaced (dataclasses.replace). Wired intooptions["reactions"]alongside the other stepper options.A real limitation found while testing (now fixed)
Testing this against a rate that references a mixing kernel via
apply_along(a spatial force-of-infection term — the patterndiphtheria_outbreakvacc's ownS -> Etransition uses) surfaced a genuine op_system bug:_compile_ir_exprraisedValueError: axes don't match arrayfor that case. Opened #186, fixed in #187 (now merged) — compilepropensity_ir_reduceinstead ofpropensity_ir_full. This branch is rebased onto the merged fix, and includes a confirming end-to-end test through the adapter (test_option_reactions_merges_mixing_kernels), exercising both mixing_kernels merging and the kernel-reduction fix together.Also noted, not touched
op_system's__init__re-export pattern (from .compile import (X, ...)+__all__listing) doesn't satisfymypy's explicit-reexport check from a consuming package's perspective — confirmed pre-existing viagit stash(2 of the resulting 3 errors already existed before this PR's export additions). Not fixed here; a real fix touches the export mechanism broadly, outside this PR's scope.Verification
flepimop2-op_system's full test suite (59 tests, 5 new) passes.ruffclean on all touched lines (pre-existing repo-widenoqa-commentsstyle noise confirmed unrelated).mypyclean on all touched lines (pre-existing unrelated errors elsewhere in both touched files confirmed viagit stashbefore/after comparison).OpSystemSystem(spec=...).options["reactions"]surfaces a working, correctly-valuedpropensity_fnfor a named collapse-to-pinned-target transition and for a real kernel-contraction transition, matching the diphtheria model's real usage shapes.Test plan
pytest(flepimop2-op_system) — 59 passedruff check/mypy— clean on touched lines