Add axis-aware parameter values by TimothyWillard · Pull Request #189 · ACCIDDA/flepimop2

TimothyWillard · 2026-03-23T18:13:23Z

This changes the Simulator class to treat parameters as structured,
axis-aware values rather than raw scalars or arrays. Systems now declare
the parameter and model-state they need, parameters resolve against
named axes at runtime, and engines receive structured initial conditions
and parameters without hard-coded assumptions about state shape.

Added flepimop2.axis module to represent axis/collection of axis.
Changed ParameterABC.sample() to return frozen ParameterValue
objects with resolved shape and values.
Placed developer-facing parameter related types under the
flepimop2.parameter.abc subpackage.
Added SystemABC.requested_parameters()/model_state() methods to
request parameters needed for the stepper or declare internal model
state, namely initial conditions.
Kept initial conditions under parameters and represent model state as
ordered parameter keys via ModelStateSpecification.
ParameterValue objects are now passed through
simulator/engine/system so steppers can unwrap and manipulate the
parameter as needed.
Removed the legacy hard-coded s0/i0/r0 initial conditions in the
flepimop2 simulate command.
Added helpers for continuous bins, bin edits, points from continous
axis objects.
Updated wrapper assets, examples, docs, and tests to use the new
structured parameter contract.

Closes #115. Closes #133.

jc-macdonald · 2026-04-07T05:22:30Z

Here is my suggested chunking of this PR to make it more reviewable — no comments on code content yet. The 5 PRs follow the natural seams of the implementation and can land sequentially without blocking work:

PR 1 — flepimop2.axis module
axis.py, undeprecate axes in ConfigurationModel, tests/axis/
Purely additive. Introduces Axis, AxisCollection, ResolvedShape, and the continuous-axis helpers (bin_edges, bins, points). Nothing else in the codebase needs to change yet.

PR 2 — Structured parameter contract
parameter/abc/__init__.py (new ParameterValue, ParameterRequest, ModelStateSpecification dataclasses), updated parameter/fixed/__init__.py (shape + broadcast logic), tests/parameter/
Depends on PR 1. Updates ParameterABC.sample() to return ParameterValue. FixedParameter gains the shape field and conflict-detection. No changes to engine, simulator, or system yet.

PR 3 — System-declared contracts
system/abc/__init__.py (requested_parameters() + model_state() methods), system/wrapper/__init__.py (script loading), new wrapper_system_with_extras.py asset, updated tests/system/
Depends on PR 2. Systems can now declare what they need — but nothing calls those methods yet, so this is still safe to review in isolation.

PR 4 — Engine/Simulator wiring + CLI
engine/abc/__init__.py, simulator.py (resolve_inputs(), updated run()), _simulate_command.py (removes hard-coded s0/i0/r0), _checked_partial.py fix, tests/simulator/, updated tests/engine/
Depends on PR 3. This is the plumbing PR — wires the contracts declared in PRs 2–3 through the runtime. The hard-coded initial-condition logic in the CLI disappears here.

PR 5 — Examples, docs, and integration test
docs/assets/SIR.py, docs/assets/solve_ivp.py, docs/development/implementing-custom-engines-and-systems.md, tests/integration/ (config, euler.py, sir.py, linear.py, updated assertions), CHANGELOG.md
Depends on PR 4. No logic changes — just updates all examples, doc snippets, and integration fixtures to use the new contract. Easy final review pass.

Each PR is ~300–470 lines and touches a single concept. Happy to discuss the boundaries further if any of these feel awkward to split out.

jc-macdonald

see comment above for my suggestions on how to split this out/

pearsonca

Some notes. Still thinking thru this one.

pearsonca · 2026-04-07T17:00:36Z

i actually think its important here to avoid depending on flepimop

pearsonca · 2026-04-07T17:02:01Z

ibid previous comment; the point of the wrapper modules should be that I can use source which has no flepimop dependencies, but still works.

pearsonca · 2026-04-07T17:03:05Z

-        }
-
        simulator = Simulator.from_configuration_model(config_model, target=target)
+        initial_state, params = simulator.resolve_inputs()


definitely improvement in syntax here.

pearsonca · 2026-04-07T19:33:35Z

    @abstractmethod
-    def sample(self) -> Float64NDArray:
-        """Sample a value from the parameter.
+    def sample(


so in terms of MVP, i think this update + the definition of Axes is approximately all that's needed for the first stage.

i'm a little unclear on what request is doing here - feels like "axes" should define what the shape is?

say a parameter is a function of age and vaccination status. then the supplied axes should cover how the parameter is broken down.

pearsonca · 2026-04-07T19:37:02Z


    def __call__(
-        self, time: np.float64, state: Float64NDArray, **kwargs: Any
+        self, time: np.float64, state: Float64NDArray, **kwargs: ParameterValue


This approach seems fine for now, but i think we'd ultimately prefer ParameterValue to be more of an internal element, and let users have more flexibility in their parameter type choice.

This defines the API they have to meet; might be we want to place the shim here, but maybe not.

pearsonca · 2026-04-07T20:28:21Z

        """
        return self._stepper(time, state, **params)

+    def requested_parameters(


i think this becomes an abstract - this particular approach would work for adapter or wrapper modules, but probably something else for op_system

This changes the `Simulator` class to treat parameters as structured, axis-aware values rather than raw scalars or arrays. Systems now declare the parameter and model-state they need, parameters resolve against named axes at runtime, and engines receive structured initial conditions and parameters without hard-coded assumptions about state shape. - Added `flepimop2.axis` module to represent axis/collection of axis. - Changed `ParameterABC.sample()` to return frozen `ParameterValue` objects with resolved shape and values. - Placed developer-facing parameter related types under the `flepimop2.parameter.abc` subpackage. - Added `SystemABC.requested_parameters()/model_state()` methods to request parameters needed for the stepper or declare internal model state, namely initial conditions. - Kept initial conditions under parameters and represent model state as ordered parameter keys via `ModelStateSpecification`. - `ParameterValue` objects are now passed through simulator/engine/system so steppers can unwrap and manipulate the parameter as needed. - Removed the legacy hard-coded s0/i0/r0 initial conditions in the `flepimop2 simulate` command. - Added helpers for continuous bins, bin edits, points from continous axis objects. - Updated wrapper assets, examples, docs, and tests to use the new structured parameter contract. Closes #115. Closes #133.

jc-macdonald · 2026-04-09T13:43:55Z

@pearsonca regarding your point about no flepi dependence I've been thinking about this kind of thing for a lot of my earthsytems and solver calibration work. if helpful here's this package I've started working on https://github.com/jcm-sci/trade-study. Probably also relevant to your scenarios work.

TimothyWillard marked this pull request as draft March 23, 2026 18:23

TimothyWillard force-pushed the axis-aware-param-values branch 2 times, most recently from 5d72ef6 to a3a1711 Compare March 27, 2026 12:30

jc-macdonald mentioned this pull request Apr 5, 2026

Add scenarios top-level configuration key #199

Closed

jc-macdonald self-requested a review April 6, 2026 16:27

TimothyWillard force-pushed the axis-aware-param-values branch from a3a1711 to 0fbccdf Compare April 6, 2026 19:14

jc-macdonald reviewed Apr 7, 2026

View reviewed changes

pearsonca reviewed Apr 8, 2026

View reviewed changes

TimothyWillard force-pushed the axis-aware-param-values branch from 0fbccdf to 79cfa09 Compare April 8, 2026 18:01

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add axis-aware parameter values#189

Add axis-aware parameter values#189
TimothyWillard wants to merge 1 commit intomainfrom
axis-aware-param-values

TimothyWillard commented Mar 23, 2026 •

edited

Loading

Uh oh!

jc-macdonald commented Apr 7, 2026

Uh oh!

jc-macdonald left a comment

Uh oh!

pearsonca left a comment

Uh oh!

pearsonca Apr 7, 2026

Uh oh!

pearsonca Apr 7, 2026

Uh oh!

pearsonca Apr 7, 2026

Uh oh!

pearsonca Apr 7, 2026

Uh oh!

pearsonca Apr 7, 2026

Uh oh!

pearsonca Apr 7, 2026

Uh oh!

jc-macdonald commented Apr 9, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

TimothyWillard commented Mar 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jc-macdonald commented Apr 7, 2026

Uh oh!

jc-macdonald left a comment

Choose a reason for hiding this comment

Uh oh!

pearsonca left a comment

Choose a reason for hiding this comment

Uh oh!

pearsonca Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

pearsonca Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

pearsonca Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

pearsonca Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

pearsonca Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

pearsonca Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

jc-macdonald commented Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TimothyWillard commented Mar 23, 2026 •

edited

Loading

jc-macdonald commented Apr 9, 2026 •

edited

Loading