State-Centered Temporal Processes

[cdc-mitzimorris]

Added state-centered parameterizations for all three temporal-process
classes in pyrenew.latent:

• AR1 — stationary AR(1) on log-Rt levels
• DifferencedAR1 — AR(1) on first differences of log-Rt (the production
  process)
• RandomWalk — unconstrained drift on log-Rt

Each class now takes a constructor argument parameterization: Literal["innovation", "state"], defaulting to "innovation" to preserve current behavior. Setting "state" switches the internal sampling from standardized increments to the latent state path directly.

All three state-centered variants are custom numpyro.distributions.Distribution subclasses in pyrenew/latent/state_centered_distributions.py:

• StateRandomWalk: samples the post-initial path as accumulated Normal increments from the deterministic initial value.
• StateAR1 and StateDifferencedAR1: the AR(1) process on the level and on first differences, respectively.

Each has a vectorized log_prob (slice arithmetic, no scan) that carries the MCMC gradient, and a sample method (cumsum for the walk, lax.scan for the AR(1) variants) used only for forward simulation such as prior and posterior predictive, not on the gradient path.

Both parameterizations encode the same prior over the state path; they differ only in sampler geometry, that is, which latent variables HMC operates on.

Code added

File	Type	Purpose
pyrenew/latent/state_centered_distributions.py	new	StateAR1, StateDifferencedAR1
pyrenew/latent/temporal_processes.py	modified	parameterization flag on all three classes; _prepare_initial_value helper
test/test_temporal_processes.py	modified	+31 unit tests (parameterization flag, state-centered shape/site/prior-equivalence)
test/test_helpers.py	modified	fixed_ar1_state, fixed_differenced_ar1_state factories
test/integration/conftest.py	modified	he_model_state_centered, he_weekly_rt_model_state_centered, he_weekly_model_state_centered fixtures
test/integration/test_population_infections_he_state_centered.py	new	5 end-to-end tests, daily Rt
test/integration/test_population_infections_he_weekly_rt_state_centered.py	new	5 end-to-end tests, weekly Rt via WeeklyTemporalProcess
_typos.toml	modified	Whitelist reparametrized_params (NumPyro upstream attribute name)

[dylanhmorris]

Thanks @cdc-mitzimorris! Two requested changes:

• Change the StateDifferencedAR1 Distribution to sample only the post-init path (https://github.com/CDCgov/PyRenew/pull/828/changes#r3421572974) and adjust the sampling logic for the downstream StateDifferencedAR1 TemporalProcess accordingly
• Ensure tests of parameterization-agnostic TemporalProcess properties are conducted in one place for all parameterizations https://github.com/CDCgov/PyRenew/pull/828/changes#r3421572974

Also left some minor suggestions (mostly docs rendering and variable naming)

[cdc-mitzimorris]

@dylanhmorris made all changes per requests and suggestions.

[dylanhmorris]

Thanks, @cdc-mitzimorris! A few suggestions for test readability and coverage but happy for you to merge without re-review.