Missingness support for Simulators

[mattlevine22]

Summary

This PR adds missing-observation support for simulator-based inference with non-Dirac observation models, plus a new tutorial showing how to use Simulator + MCMC with missing data.

On the implementation side, DiscreteTimeSimulator and ODESimulator now support obs_values containing NaNs under an index-preserving interpretation: obs_values[k] always refers to latent index k, so missing observations do not remove latent states from the model. Full-row missingness is supported for all non-Dirac observation models by simply skipping the observation log-probability at those indices. Partial missingness is supported when the observation model has structure that can be marginalized cleanly, specifically:

• MultivariateNormal observations
• factorized Independent(..., 1) observations

To keep the simulator code clean, the missing-data logic is factored into a shared internal scorer layer in dynestyx/internal/observation_missingness.py. When NaNs are present, the simulator switches from obs=... sample sites to numpyro.factor(...)-based masked observation scoring while preserving the existing no-missing fast paths.

The PR also adds a focused test suite under tests/missingness/ covering:

• scorer math,
• discrete simulator behavior,
• ODE simulator behavior,
• hierarchical/plated behavior.

Finally, it adds a new gentle-intro tutorial notebook, 11b_missing_observations_simulator_mcmc.ipynb, which explains the full-vs-partial missingness distinction for simulator-based inference, compares the Gaussian case against KF smoothing, and illustrates additional non-Gaussian observation families supported by the simulator missingness path.

Questions / notes:

• do we really want the missingness utils in dynestyx/internal/observation_missingness.py?

[Copilot]

Pull request overview

This PR adds NaN-based missing-observation support for simulator-backed inference, with shared scorer utilities and tests covering discrete-time, ODE, hierarchical, and scorer behavior.

Changes:

• Added internal missing-observation preprocessing/scoring helpers for full-row and partial missingness.
• Integrated factor-based observation scoring into DiscreteTimeSimulator and ODESimulator.
• Added focused missingness test models/utilities and updated the gentle intro tutorial index.

Reviewed changes

Copilot reviewed 11 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
dynestyx/internal/__init__.py	Adds an internal helpers package marker.
dynestyx/internal/observation_missingness.py	Adds missing-observation preprocessing and masked likelihood scoring.
dynestyx/simulators.py	Wires missing-observation scoring into discrete and ODE simulators.
tests/missingness/__init__.py	Adds missingness test package marker.
tests/missingness/models.py	Adds shared discrete/ODE/plated test models.
tests/missingness/utils.py	Adds masking and manual log-prob test helpers.
tests/missingness/test_scorers.py	Tests scorer preprocessing and masked likelihood math.
tests/missingness/test_discrete_simulator.py	Tests discrete simulator missingness behavior and MCMC smoke path.
tests/missingness/test_ode_simulator.py	Tests ODE simulator missingness behavior and MCMC smoke path.
tests/missingness/test_hierarchical.py	Tests plated simulator missingness behavior.
docs/tutorials/gentle_intro/00_index.ipynb	Adds missing-observation simulator tutorial entry to the index.

---

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

[mattlevine22]

@copilot does the most recent commit address your comments/concerns?

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated 2 comments.

[DanWaxman]

I think this is directionally right, but I have a few concerns:

1. There's a lot of new complexity; observation_missingness has a dataclass, Protocol, datahelper, scorer class, and factory, along with some math. This is a lot of complexity when it feels like just the data helper and scorer class should suffice.
2. Simulators is getting quite long. It may be out of scope for this PR, but I do think it's time to break up that class if we're adding these new features.
3. I don't really want totally separate pathways for the missing data unless it's totally necessary (and I feel it probably isn't). The _simulate_missing_scan is doing almost the same thing as the normal simulation loop, with an abstracted scorer (which I would like to change the name of); we should just use that abstracted scorer directly, always, if we're going to already add the new complexity.

[DanWaxman]

It's true that we should use particle filtering/smoothing here, but that's because EKF is going to do really bad with Laplace observations.

[mattlevine22]

I changed to:

If we want to perform Filtering/Smoothing as before, we would need to use a variant that can handle these non-Gaussian observations (e.g. PF/PS or EKF/EKS); however, as noted in the previous tutorial, these methods are not yet set up to automatically support missing data. The EnKF, while configured to automatically support missing data, is not set up to support non-Gaussian observations.

Nevertheless, the MCMC + Simulator() strategy can handle all forms of missingness (partial and full) under non-Gaussian observation models as long as the observation model is independent per observation dimension, which is the case here.

[DanWaxman]

Why size 1?

[mattlevine22]

this appears in main as well...I think it is to keep a leading axis for consistency w.r.t. when we do multiple rollout simulations

[DanWaxman]

I'd like to avoid "score", it's somewhat overloaded for us

[mattlevine22]

renamed things to log_potential