Skip to content

Interpolate time-varying params in reaction propensity_fn - #196

Merged
jc-macdonald merged 1 commit into
mainfrom
feature/reaction-propensity-time-varying-params
Aug 26, 2026
Merged

Interpolate time-varying params in reaction propensity_fn#196
jc-macdonald merged 1 commit into
mainfrom
feature/reaction-propensity-time-varying-params

Conversation

@jc-macdonald

Copy link
Copy Markdown
Member

Summary

Companion to #195 (source-only transitions), and independently useful on
its own -- found while wiring diphtheria_outbreakvacc's cross-district
case-importation term, which needs a time-varying hazard (concentrated in
an initial window, much lower afterward).

_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 (the expression
itself is already rewritten to the time-stripped lambda_import[loc]
form upstream by _strip_time_axis_in_expr, so the compiled code expects
the reduced shape). 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. Confirmed via a
failing call before this fix:

ValueError: cannot reshape array of size N into shape (...)

Changes

  • Adds _wrap_propensity_fn_for_time_varying, mirroring the existing two
    wrappers' exact interpolation logic (same _interp_along_axis call,
    same linear-interpolation/constant-extrapolation behavior), applied to
    every compiled reaction in _build_reaction_artifacts.

Test plan

  • 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 test suite: 457 passed
  • mypy --strict on compile.py: clean
  • ruff check: no new findings beyond this project's existing
    unrelated pre-existing noqa notes

_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
jc-macdonald force-pushed the feature/reaction-propensity-time-varying-params branch from 150deba to 7250829 Compare August 26, 2026 20:16
@jc-macdonald
jc-macdonald merged commit 604609a into main Aug 26, 2026
3 of 4 checks passed
@jc-macdonald
jc-macdonald deleted the feature/reaction-propensity-time-varying-params branch August 26, 2026 20:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant