Simulator refactor

[mattlevine22]

Summary

This PR finishes the simulator/latent-path refactor by separating forward generation from explicit latent-state inference.

The main architectural change is that simulators are now generation-only. Forward rollout lives in the simulator stack, while joint state/parameter inference now goes through LatentPathBuilder, backed by shared pure-JAX state-path reconstruction and scoring utilities.

Along the way, this also adds new standalone APIs and broadens the missingness/inference workflows supported by the library.

Main changes

• Simulators are now generation-only.
• dsx.sample(..., obs_times=..., obs_values=...) with simulator handlers no longer serves as a latent-state inference route.
• The old monolithic dynestyx/simulators.py implementation was split into a new dynestyx/simulation/ package.
• Simulator execution is now pure JAX and returns SimulatedResult, with NumPyro site registration deferred to a _register_numpyro_sites step.
• Added standalone pure-JAX APIs in dynestyx/api.py:
  • dsx.simulate(...) for forward simulation
  • dsx.log_prob(...) for joint trajectory scoring
• Added LatentPathBuilder as the dedicated explicit latent-path inference handler.
• Added shared pure-JAX latent/state-path internals under:
  • dynestyx/inference/latent/
  • dynestyx/inference/state_paths/
• LatentPathBuilder is now NumPyro-facing only and supports dsx.sample(...), not dsx.condition(...).
• Simplified the latent-path execution flow so the builder now does a single reconstruction/scoring pass after dummy latent sites are created.
• Reworked missingness handling across latent-path inference:
  • partially missing DiracIdentityObservation data now automatically creates only the needed free state coordinates as latents
  • unsupported partially missing continuous observation models can fall back to explicit missing-observation augmentation through f_missing_obs_values
  • reconstructed completed observations are exposed through deterministic outputs such as f_completed_obs_values
  • ODE + Dirac inference is explicitly rejected
• Reorganized inference config modules under dynestyx/inference/configs/
• Moved inference helper utilities into dynestyx/inference/utils/
• Tightened the public API surface by removing top-level exports for low-level latent/missingness preparation helpers

New capabilities

• Pure-JAX forward simulation via dsx.simulate(...), without needing a NumPyro handler stack
• Pure-JAX trajectory scoring via dsx.log_prob(...), for evaluating log p(x, y | ...) on user-supplied latent paths
• A dedicated explicit latent-state inference route via LatentPathBuilder
• Deterministic ODE latent inference through LatentPathBuilder, where the latent path parameters are just the initial condition and the full path is reconstructed by solving the ODE
• Cleaner and broader missing-data support in latent-path inference:
  • exact-observation / Dirac partial missingness now infers only the missing state coordinates
  • unsupported correlated continuous observation families can use explicit observation augmentation
  • posterior samples of reconstructed missing observations are now directly available through f_completed_obs_values
• Cleaner composition of rollout with inference handlers, since simulators can now sit outside inference handlers as pure rollout/generation layers
• Unified simulator configuration objects (SimulatorConfig, ODESimulatorConfig, SDESimulatorConfig)
• Better support for custom pure-JAX workflows outside NumPyro, including manual optimization / inference loops

A frustration (caching)
One tricky design point in this refactor is latent-path layout preparation under NumPyro replay. LatentPathBuilder needs concrete observation missingness structure to determine latent dimensionality and indexing, but NumPyro’s MCMC/Predictive machinery may re-enter the model with traced values rather than ordinary arrays. That creates tension between three goals we want simultaneously: a simple fully hidden user API, compressed missingness-aware latent parameterizations, and replay-safe preparation under NumPyro tracing. The current cache-based solution is a pragmatic way to balance those goals: it preserves the simple API and compressed layouts without forcing users to precompute or pass bookkeeping objects themselves, at the cost of some extra internal complexity in this part of the implementation.

Docs and tutorials

• Updated simulator docs to reflect the generation-only contract
• Added dedicated public/developer docs for LatentPathBuilder
• Updated tutorials to use LatentPathBuilder instead of simulators for explicit latent inference
• Reworked 05_svi_no_numpyro to use the new dsx.simulate(...) generation pattern
• Replaced the old simulator-based missingness MCMC tutorial with a latent-path version:
  • 11b_missing_observations_latent_path_mcmc.ipynb
• Updated the missingness tutorial to show Dirac partial missingness and reconstructed observations under augmentation

Tests

• Added standalone coverage for:
  • dsx.simulate(...)
  • dsx.log_prob(...)
  • state-path assembly
  • LatentPathBuilder
  • simulator configs
• Updated simulator/inference/missingness smoke tests to the new API split

[DanWaxman]

I changed the base branch from main to dsx-infer-staging and took a quick look. The organization looks reasonable.

1. There's a few places in docs where files should be moved around.
2. It's surprising to me that latent has so much code now. I wonder if there's some boilerplate we can cut down on, or ways to unify the code. I get that missingness adds some complexity for us, but this amount of code is hard to digest and maintain.
3. I wonder the extent we need a type of callback machinery -- seems more complex than what goes on elsewhere.

[DanWaxman]

Will have to think if this is the best place

[DanWaxman]

No need for this here

[Copilot]

Pull request overview

This PR completes the simulator/latent-path refactor by making simulators generation-only, introducing a dedicated LatentPathBuilder for explicit latent inference, and adding pure-JAX public APIs (dsx.simulate, dsx.log_prob). It also reorganizes inference configuration/utilities and updates tests + documentation to match the new architecture.

Changes:

• Split simulator implementation into dynestyx/simulation/ and enforce generation-only semantics for simulators in inference workflows.
• Add pure-JAX APIs (dynestyx/api.py) for forward simulation and joint trajectory scoring, plus new pure-JAX state-path reconstruction/scoring internals.
• Introduce LatentPathBuilder and rework/extend missingness + hierarchical workflows; update tests and docs to the new API surface.

Reviewed changes

Copilot reviewed 129 out of 136 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/test_time_varying_linear_gaussian.py	Update imports to new inference.configs.* locations and simulator re-export.
tests/test_state_path_assembly.py	Add coverage for pure-JAX state-path assembly (discrete identity, ODE reconstruction, SDE rejection).
tests/test_smoothers.py	Update config imports and adjust expected error message module path.
tests/test_smoother_standalone.py	Add rejection test for unsupported Dirac+ODE inference; update config import path.
tests/test_simulator_configs.py	Add coverage for structured simulator config objects and routing/validation.
tests/test_simulate_standalone.py	Add coverage for pure-JAX dsx.simulate and deferred NumPyro site registration behavior.
tests/test_sgmcmc_inference_smoke.py	Update config imports to new inference.configs.* locations.
tests/test_science/test_ode_hierarchical_simulator_inference.py	Update simulator import to top-level re-export.
tests/test_science/test_nonlinear_discretized_hierarchical_multitraj_ekf.py	Update simulator + filter-config import paths.
tests/test_science/test_nonlinear_discrete_hierarchical_two_level_ekf.py	Update simulator + filter-config import paths.
tests/test_science/test_nonlinear_discrete_hierarchical_multitraj_ekf.py	Update simulator + filter-config import paths.
tests/test_science/test_lti_discrete_smoothing.py	Update simulator + smoother-config import paths.
tests/test_science/test_lti_discrete_simplified_sgmcmc.py	Update simulator + mcmc-config import paths.
tests/test_science/test_lti_discrete_hierarchical_multitraj_kf.py	Update simulator + filter-config import paths.
tests/test_science/test_lti_continuous_smoothing.py	Update simulator wrapper import and smoother-config path.
tests/test_science/test_lti_continuous_hierarchical_multitraj_kf.py	Update simulator wrapper import and filter-config path.
tests/test_science/test_l63_SDE_sgmcmc.py	Update simulator wrapper import and mcmc-config path.
tests/test_science/test_l63_ODE_sgmcmc.py	Update simulator wrapper import and mcmc-config path.
tests/test_science/test_hmm_hierarchical_multitraj.py	Update simulator + HMM config import paths.
tests/test_science/test_hierarchical_svi.py	Update simulator + filter-config import paths.
tests/test_science/test_discrete_hierarchical_simulator_inference.py	Update simulator import to top-level re-export.
tests/test_sample_input_matrix.py	Update behavior expectations: simulator conditioning paths now error (generation-only contract).
tests/test_predictive_filter_simulator_shapes.py	Update filter-config import path.
tests/test_plate_vector_initial_mean_continuous_filters.py	Update filter-config import path.
tests/test_plate_vector_initial_mean_continuous_drift_diffusion.py	Update filter/smoother config import paths and simulator util import (_slice_tree_for_plate_member).
tests/test_models_core.py	Update simulator import to top-level re-export.
tests/test_missing_observations.py	Update filter/smoother config import paths.
tests/test_mcmc_filter_smokes.py	Update simulator import + mcmc config import paths (reordered imports).
tests/test_mcmc_filter_smokes_discrete.py	Update simulator + mcmc-config import paths.
tests/test_log_prob_standalone.py	Add coverage for pure-JAX dsx.log_prob, including missingness and error cases.
tests/test_hierarchical_svi_smokes.py	Update filter-config import path.
tests/test_hierarchical_smokes.py	Update filter + mcmc config import paths.
tests/test_hierarchical_simulator_inference_smokes.py	Switch hierarchical explicit inference smoke tests from simulators to LatentPathBuilder.
tests/test_hierarchical_simulator_discretizer_smokes.py	Update plate-conditioning expectations (generation-only errors) and add LatentPathBuilder Dirac latent checks.
tests/test_filters.py	Update filter-config import path and simulator import to top-level re-export.
tests/test_filter_standalone.py	Add rejection test for unsupported Dirac+ODE inference; update config import path.
tests/test_filter_simulator.py	Update filter-config import paths (including local import in test).
tests/test_distribution_utils.py	Update import path to dynestyx.inference.utils.distribution_utils.
tests/test_discretizers.py	Update filter-config import path.
tests/test_desirable_shared_initial_mean_hierarchical.py	Update filter-config import path.
tests/test_blackjax_no_crn.py	Update simulator import + config imports and reorder integration imports.
tests/missingness/test_ode_simulator.py	Rework missingness tests to use LatentPathBuilder and joint log-prob outputs.
tests/missingness/test_observation_log_prob.py	Extend missingness log-prob tests and add explicit missing_observation_strategy usage.
tests/missingness/test_hmm.py	Update HMM config import path.
tests/missingness/test_hierarchical.py	Rework hierarchical missingness tests to latent-path route and update expected factorization.
tests/fixtures.py	Update fixtures to use LatentPathBuilder for data-conditioned models; update config imports.
mkdocs.yml	Update nav and API reference paths for new config modules and LatentPathBuilder docs.
dynestyx/utils.py	Update internal reference string for plate utils location (and related comments).
dynestyx/types.py	Add LatentStateResult, SimulatedResult, and a helper to chain deferred NumPyro registrations.
dynestyx/simulation/sde.py	New SDE simulator backend with pure-JAX rollout and structured config support.
dynestyx/simulation/ode.py	New ODE simulator backend with pure-JAX rollout and structured config support.
dynestyx/simulation/discrete.py	New discrete-time simulator backend (pure-JAX rollout).
dynestyx/simulation/init.py	Export simulation backends and public simulator handler classes/utilities.
dynestyx/inference/utils/plate_utils.py	New shared plate helpers (in_axes building, slicing distributions/arrays).
dynestyx/inference/utils/numpyro_sites.py	Update imports to new config modules and utils package structure.
dynestyx/inference/utils/distribution_utils.py	Update plate-utils import path to new inference.utils location.
dynestyx/inference/utils/init.py	New utils package initializer.
dynestyx/inference/state_paths/score.py	New pure-JAX trajectory scoring utilities (log-prob term decomposition).
dynestyx/inference/state_paths/reconstruct.py	New pure-JAX state-path reconstruction utilities (including ODE path reconstruction).
dynestyx/inference/state_paths/init.py	New public exports for state-path layout/reconstruct/score internals.
dynestyx/inference/smoothers.py	Rewire imports to new configs/utils; chain deferred NumPyro site registrations from forwarded handlers; add inference class validation.
dynestyx/inference/mcmc.py	Update to new mcmc config module path.
dynestyx/inference/latent/prepare.py	New latent-path request preparation with caching and canonicalization.
dynestyx/inference/latent/plate.py	New plate-splitting/restacking logic for latent-path inference.
dynestyx/inference/latent/init.py	New latent subpackage export (LatentPathBuilder).
dynestyx/inference/integrations/cuthbert/discrete_smoother.py	Update to new config modules and distribution utils location.
dynestyx/inference/integrations/cuthbert/discrete_filter.py	Update to new config modules and distribution utils location.
dynestyx/inference/integrations/cd_dynamax/discrete_smoother.py	Update to new config modules and distribution utils location.
dynestyx/inference/integrations/cd_dynamax/discrete_filter.py	Update to new config modules and distribution utils location.
dynestyx/inference/integrations/cd_dynamax/continuous_smoother.py	Update to new config modules and distribution utils location.
dynestyx/inference/integrations/cd_dynamax/continuous_filter.py	Update to new config modules and distribution utils location.
dynestyx/inference/integrations/blackjax/mcmc.py	Update to new mcmc config module path.
dynestyx/inference/hmm_filters.py	Update HMM config import path.
dynestyx/inference/filters.py	Rewire imports to new configs/utils; chain deferred NumPyro registrations; add inference class validation; update doc link.
dynestyx/inference/configs/smoother.py	Update base imports/docs to reference new inference.configs.filter module.
dynestyx/inference/configs/simulator.py	Add structured simulator configs (ODESimulatorConfig, SDESimulatorConfig).
dynestyx/inference/configs/mcmc.py	Add consolidated MCMC config dataclasses under inference.configs.mcmc.
dynestyx/inference/configs/init.py	New configs subpackage initializer.
dynestyx/inference/checkers.py	Add model-class validation for inference/scoring (reject ODE+Dirac identity observation).
dynestyx/inference/init.py	Update inference package exports to include new subpackages.
dynestyx/handlers.py	Refactor sample/condition flow to support deferred site registration and new handler result types.
dynestyx/discretizers.py	Widen handler return typing to support new latent-path result types.
dynestyx/api.py	New pure-JAX simulate and log_prob public APIs.
dynestyx/init.py	Re-export new APIs, simulator/config classes, and result types at top level.
docs/tutorials/quickstart.ipynb	Update simulator docs text to reflect generation-only simulator handlers.
docs/tutorials/gentle_intro/11c_missing_observations_hmms.ipynb	Update imports and cross-link text for new 11b tutorial.
docs/tutorials/gentle_intro/11_missing_observations.ipynb	Update config imports and add guidance for LatentPathBuilder augmentation fallback.
docs/tutorials/gentle_intro/10_continuous_smoothing.ipynb	Update config import paths.
docs/tutorials/gentle_intro/09_discrete_smoothing.ipynb	Update config import paths.
docs/tutorials/gentle_intro/08_hierarchical_inference.ipynb	Update config import paths.
docs/tutorials/gentle_intro/00_index.ipynb	Update narrative to reference LatentPathBuilder for explicit inference.
docs/math_intro.md	Update simulation section to mention dsx.simulate and new imports/paths; fix API reference link.
docs/faq.md	Update recommended explicit-inference route from simulators to LatentPathBuilder and adjust wording.
docs/deep_dives/mcmc_inference_algorithm_comparison.ipynb	Update imports to new config modules and simulator re-export.
docs/deep_dives/lv_uode.ipynb	Update filter config import paths.
docs/deep_dives/gp_drift.ipynb	Update filter config import paths.
docs/deep_dives/discrete_time_kf_associative_scaling.ipynb	Update filter config import paths and simulator import.
docs/api_reference/public/simulators/simulator_wrapper.md	Update doc target module path and examples to generation-only contract.
docs/api_reference/public/simulators/sde_simulator.md	Update doc target module path and examples to generation-only contract.
docs/api_reference/public/simulators/overview.md	Update overview text to generation-only simulators + new pure API split.
docs/api_reference/public/simulators/ode_simulator.md	Update doc target module path and examples to generation-only contract.
docs/api_reference/public/simulators/discrete_time_simulator.md	Update doc target module path and examples to generation-only contract.
docs/api_reference/public/models/specialized/lti_discrete.md	Update wording/link paths to new filter-config docs location.
docs/api_reference/public/models/specialized/lti_continuous.md	Update wording/link paths to new filter-config docs location.
docs/api_reference/public/models/specialized/linear_gaussian_state_evolution.md	Update link paths to new filter-config docs location.
docs/api_reference/public/models/specialized/linear_gaussian_observation.md	Update link paths to new filter-config docs location.
docs/api_reference/public/models/specialized/gaussian_state_evolution.md	Update link paths to new filter-config docs location.
docs/api_reference/public/models/specialized/gaussian_observation.md	Update link paths to new filter-config docs location.
docs/api_reference/public/models/specialized/affine_drift.md	Update wording/link paths to new filter-config docs location.
docs/api_reference/public/inference/smoothers.md	Update smoother config import path.
docs/api_reference/public/inference/latent_path_builder.md	Add public API docs for LatentPathBuilder.
docs/api_reference/public/inference/configs/smoother_configs.md	Update doc target module path and example imports.
docs/api_reference/public/inference/configs/mcmc_configs.md	Update doc target module path and add canonical import note.
docs/api_reference/public/inference/configs/filter_configs.md	Update doc target module path and add canonical import note.
docs/api_reference/public/handlers.md	Expand handlers doc to include new pure APIs and clarify primitive roles.
docs/api_reference/developer/simulators.md	Update developer simulator docs to new dynestyx.simulation package and generation-only semantics.
docs/api_reference/developer/inference/latent_path_builder.md	Add developer-facing LatentPathBuilder architecture reference.
docs/api_reference/developer/inference/configs/smoother_configs.md	Update doc target module path and add canonical module note.
docs/api_reference/developer/inference/configs/mcmc_configs.md	Update doc target module path and add canonical module note.
docs/api_reference/developer/inference/configs/filter_configs.md	Update doc target module path and add canonical module note.
docs/api_reference/developer/handlers.md	Update developer handler docs to include pure APIs and explicit member listing.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 129 out of 136 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

dynestyx/handlers.py:246

• _condition_intp is now interpreted by handlers that can return non-FunctionOfTime results (e.g. ConditionedResult, SimulatedResult, LatentStateResult). Keeping the return type annotated as FunctionOfTime is misleading for type-checking and IDE usage; other handlers in this PR already relaxed similar return annotations to Any.