Support {state}_init symbols for initial values in equations#32
Merged
Conversation
Track and expose initial-value symbols so equations can reference a state's initial value via "{state}_init". Equation: add an inits set and collect/validate "*_init" symbols during parsing (ensure referenced state exists and do not register them as parameters). Simulation: forward y0 as an extra arg to the integrator so initial values are available at runtime. Stack: expose an "inits" symbol list (positionally matching state order), accept an optional inits array when evaluating modules, and propagate it through ODE/Reaction stacks. Tests: add an integration test verifying that "{state}_init" reuses the measurement initial value (not a parameter) and that rates behave accordingly.
Ensure initial-value symbols of the form "{state}_init" are resolved from live y0s and not treated as regular constants. Changes include: add target_accept_prob to MCMCConfig and pass it into HMC/NUTS; allow add_assignment to accept kwargs; exclude init symbols from get_constants_order via a new _is_init_symbol helper; rework surrogate-backed sim function to treat flattened y0s as per-measurement points, derive per-measurement inits from each measurement's t=0 state, and vmapped/jit the rate evaluation accordingly. Minor formatting tweak to a sigma prior call. Added integration tests covering surrogate init resolution, estimated-initials MCMC path, and init symbols declared as constants to prevent KeyError when building the constants matrix.
Remove an unused `import jax` from tests/integration/test_mcmc.py. This is a small cleanup to eliminate a leftover import and avoid linter/unused-import warnings; there are no behavioral changes.
Document the new `{state}_init` suffix for referencing a state's initial value inside ODE equations. Adds a brief example showing `S` vs `S_init`, explains use cases (e.g. normalization or driving terms by initial concentration), and notes that the referenced state must exist (raises ValueError if not), `_init` symbols are not registered as parameters, and current and initial values may be mixed in the same expression.
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
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.
This pull request introduces a robust mechanism for referencing the initial value of any state in model equations using the
{state}_initsymbol. It ensures that initial values are resolved dynamically from simulation initial conditions, both during standard simulations and when using surrogate models or MCMC with sampled initial conditions. Additionally, the documentation and tests have been updated to reflect and validate this new feature.Support for referencing initial values in equations:
{state}_initsuffix. These symbols are resolved from the simulation's initial conditions, not registered as parameters, and are validated to ensure the referenced state exists in the model. [1] [2] [3] [4]Simulation and surrogate model integration:
{state}_initsymbols are always resolved from the current initial conditions (y0) passed to the simulation, ensuring correct behavior when initial conditions are sampled or transformed (e.g., in MCMC or surrogate rate evaluation). [1] [2] [3] [4] [5] [6] [7]Testing and validation:
{state}_initbehave correctly, both in surrogate and full simulation paths, including when initial conditions are sampled (e.g., viaestimate_initials). This ensures the new mechanism works end-to-end and guards against regressions.Minor improvements and cleanup: