Interpolate time-varying params in reaction propensity_fn - #196
Merged
jc-macdonald merged 1 commit intoAug 26, 2026
Conversation
_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
force-pushed
the
feature/reaction-propensity-time-varying-params
branch
from
August 26, 2026 20:16
150deba to
7250829
Compare
jc-macdonald
deleted the
feature/reaction-propensity-time-varying-params
branch
August 26, 2026 20:21
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
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
Companion to #195 (source-only transitions), and independently useful on
its own -- found while wiring
diphtheria_outbreakvacc's cross-districtcase-importation term, which needs a time-varying hazard (concentrated in
an initial window, much lower afterward).
_build_reaction_artifactscompiles each named transition's propensityindependently of
_wrap_eval_fn_for_time_varying/_wrap_pytree_eval_fn_for_time_varying-- the two wrappers that make arate like
rate: "lambda_import[time, loc]"work for the deterministicRHS, by interpolating the raw
(time, loc)grid down to the currenttimestep's
(loc,)slice before the compiled code runs (the expressionitself is already rewritten to the time-stripped
lambda_import[loc]form upstream by
_strip_time_axis_in_expr, so the compiled code expectsthe reduced shape). Reaction
propensity_fns never went through eitherwrapper, so the same rate reached a reaction's compiled code as the raw,
un-interpolated full grid array -- a shape mismatch. Confirmed via a
failing call before this fix:
Changes
_wrap_propensity_fn_for_time_varying, mirroring the existing twowrappers' exact interpolation logic (same
_interp_along_axiscall,same linear-interpolation/constant-extrapolation behavior), applied to
every compiled reaction in
_build_reaction_artifacts.Test plan
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 afractional, off-grid
t-- same correctness-oracle pattern as theexisting
test_reactions_reconstruct_deterministic_eval_fn.mypy --strictoncompile.py: cleanruff check: no new findings beyond this project's existingunrelated pre-existing
noqanotes