Fix reaction propensity for from-side-pinned transitions - #191
Merged
Conversation
_in_order_wildcard_axes() only extracted WildcardTokens from a transition's from-side selector, so a transition pinned on the from side too (e.g. a vaccination-dose-progression transition written as S[age, vax=unvaccinated] -> S[age, vax=partial]) got a from_sub Subscript missing that axis entirely relative to S's true shape. compile-time lowering then failed and the reaction was silently dropped from CompiledRhs.reactions (opportunistic omission, no error) -- 4 of the real diphtheria config's 25 named transitions (vax_dose1/vax_dose2/wane_full/wane_partial) were affected. _reactions.py: added ReactionArtifactIR.full_axes (every axis of from_base's true template, wildcard or pinned) and build from_sub's indices from all from-side tokens, baking in pinned coordinates via AxisIndex(coord=...) instead of dropping them. Also widened `pinned` to include from-side-pinned axes (previously filtered to only to-side-pinned-and-from-wildcard axes), since a from/to double-pinned axis is a "point-to-point" shift that still needs a fixed scatter target coordinate, not a collapse. That surfaced a deeper, pre-existing gap: lower_subscript_to_buffer() (_ir_lower.py) only ever supported FREE-axis subscript indices -- any COORD-kind (pinned literal) index raised immediately, by design (the v1 docstring called this out explicitly). The deterministic RHS path never hits this because it fully expands pinned coordinates to per-cell scalars before vector lowering; the reaction-propensity path needs a mixed free+pinned Subscript to reach the vector lowerer directly. Extended lower_subscript_to_buffer to resolve COORD indices against axis_coords and slice them out of the buffer before the FREE-axis transpose/broadcast logic runs, so a Subscript can mix wildcard and pinned axes generally, not just for reactions. compile.py: threaded full_axes through to CompiledReaction (a consumer needs it to build a complete scatter-target index into to_base when an axis is pinned on the from side). Verified against the real diphtheria_outbreakvacc config: all 4 previously-missing from-pinned transitions now compile and their propensities match the deterministic pytree_eval_fn exactly.
jc-macdonald
added a commit
that referenced
this pull request
Aug 26, 2026
from_pinned records the from-side pinned coordinate for every axis a transition pins on its from-side selector (full_axes minus from_axes). This was missing from #191: that PR added full_axes and widened `pinned` to cover from-pinned axes too, but `pinned` only ever carries the TO-side coordinate for those axes. For a point-to-point transition (pinned on both sides, e.g. a vaccination-dose-progression transition pinned at vax=unvaccinated on from and vax=partial on to), the from-side and to-side coordinates differ, so a consumer building a depletion-target index into from_base cannot derive it from `pinned` alone -- confirmed by diphtheria_outbreakvacc's run_hybrid_ctmc consumer, which needs exactly this to deplete the correct source cell rather than the (wrong) to-side one. Also fixes ruff-format drift left on main by #191's merge (my own _reactions.py/compile.py edits weren't format-clean) and a pre-existing, unrelated formatting issue in README.md -- both were failing the "quality" CI check on main.
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.
Fixes #189.
CompiledRhs.reactionssilently omitted any named transition whosefrom-side selector pins a coordinate on an axis (e.g.S[age, vax=unvaccinated] -> S[age, vax=partial]) — 4 of the diphtheria_outbreakvacc config's 25 named transitions were affected.Root cause 1 —
_reactions._in_order_wildcard_axes()only extractedWildcardTokens from the from-side selector, dropping anyPinnedToken, so thefrom_subSubscript ended up missing that axis relative to the state buffer's true shape.Root cause 2 (deeper, pre-existing) —
lower_subscript_to_buffer()only ever supported FREE-axis subscript indices; any pinned-coordinate (COORD-kind) index raised immediately by design. The deterministic RHS path never hits this because it fully expands pinned coordinates to per-cell scalars before vector lowering; the reaction-propensity path needs a mixed free+pinnedSubscriptto reach the vector lowerer directly.Changes:
_reactions.py: addedReactionArtifactIR.full_axes, fixedfrom_subto bake in from-side pinned coordinates instead of dropping them, widenedpinnedto include from-side-pinned axes (a from/to double-pinned axis is a "point-to-point" shift, not a collapse, but still needs a fixed scatter-target coordinate)._ir_lower.py: extendedlower_subscript_to_bufferto resolve COORD-kind indices againstaxis_coordsand slice them out of the buffer before the FREE-axis transpose/broadcast logic — a general capability, not reaction-specific.compile.py: threadedfull_axesthrough toCompiledReaction.Testing: full suite (472 tests) + mypy clean; new tests for both the reaction-level fix and the underlying lowering capability; verified against the real diphtheria_outbreakvacc config (
vax_dose1/vax_dose2/wane_full/wane_partialnow compile and match the deterministicpytree_eval_fnexactly).