From 5daeceda3372fe70530de69bf263bde73daf1163 Mon Sep 17 00:00:00 2001
From: Mitzi Morris <bf06@cdc.gov>
Date: Tue, 15 Sep 2026 09:40:02 -0400
Subject: [PATCH 1/6] prelim AGENTS.md instructions for codex

---
 AGENTS.md | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 AGENTS.md

diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 00000000..51fe97c3
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,14 @@
+Functions, classes, and packages designs should be straightforward.
+The code needs to be clean and human-readable and human-maintainable.
+Factor repeated computation into general functions.
+Avoid overly general and extensible frameworks.
+Always ask if a design or code change requires backwards compatibility.
+In most cases, backwards compatibility is not a concern.
+
+coding style:   always use MyPy type hints.   all functions should have a docstring; public-facing API functions should be user-focussed.  users want to know what the function does, ditto maintainers.  omit design discussion details.
+  inline comments only for critical, non-obvious details.
+
+Design and implementation plans should be written in plain style.
+describe math in terms of python computation.
+
+I will commit and push edited files to GitHub myself - do not commit or push unless I tell you to.

From 5e164c6cabf08bd96a271d18aaf9f8f8f74569a4 Mon Sep 17 00:00:00 2001
From: Mitzi Morris <bf06@cdc.gov>
Date: Tue, 15 Sep 2026 15:58:59 -0400
Subject: [PATCH 2/6] fix state-centered innovation process

---
 pyrenew/latent/__init__.py           |   8 +-
 pyrenew/latent/temporal_processes.py | 140 +++++++-----
 test/test_temporal_processes.py      | 323 ++++++++++++++++++++++++++-
 3 files changed, 409 insertions(+), 62 deletions(-)

diff --git a/pyrenew/latent/__init__.py b/pyrenew/latent/__init__.py
index c1b9caa5..2fbe4837 100644
--- a/pyrenew/latent/__init__.py
+++ b/pyrenew/latent/__init__.py
@@ -28,14 +28,16 @@
 from pyrenew.latent.infections import Infections
 from pyrenew.latent.infectionswithfeedback import InfectionsWithFeedback
 from pyrenew.latent.population_infections import PopulationInfections
+from pyrenew.latent.state_centered_distributions import (
+    StateAR1,
+    StateDifferencedAR1,
+    StateRandomWalk,
+)
 from pyrenew.latent.subpopulation_infections import SubpopulationInfections
 from pyrenew.latent.temporal_processes import (
     AR1,
     DifferencedAR1,
     RandomWalk,
-    StateAR1,
-    StateDifferencedAR1,
-    StateRandomWalk,
     StepwiseTemporalProcess,
     TemporalProcess,
     WeeklyTemporalProcess,
diff --git a/pyrenew/latent/temporal_processes.py b/pyrenew/latent/temporal_processes.py
index b4960023..48f72e08 100644
--- a/pyrenew/latent/temporal_processes.py
+++ b/pyrenew/latent/temporal_processes.py
@@ -54,17 +54,14 @@
 
 from typing import Literal, Protocol, runtime_checkable
 
+import jax
 import jax.numpy as jnp
 import numpyro
 import numpyro.distributions as dist
 from jax.typing import ArrayLike
+from numpyro.contrib.control_flow import scan
 
 from pyrenew.deterministic import DeterministicVariable
-from pyrenew.latent.state_centered_distributions import (
-    StateAR1,
-    StateDifferencedAR1,
-    StateRandomWalk,
-)
 from pyrenew.metaclass import RandomVariable
 from pyrenew.process import ARProcess, DifferencedProcess
 from pyrenew.process.randomwalk import RandomWalk as ProcessRandomWalk
@@ -171,7 +168,7 @@ def _validate_parameterization(parameterization: str) -> None:
 def _prepare_initial_value(
     initial_value: float | ArrayLike | None,
     n_processes: int,
-) -> ArrayLike:
+) -> jax.Array:
     """
     Resolve a per-process initial value to a 1D array of length n_processes.
 
@@ -303,9 +300,9 @@ def sample(
         """
         initial_value = _prepare_initial_value(initial_value, n_processes)
 
-        autoreg = self.autoreg_rv()
-        innovation_sd = self.innovation_sd_rv()
-        autoreg_broadcast = jnp.broadcast_to(jnp.asarray(autoreg), (n_processes,))
+        autoreg = jnp.asarray(self.autoreg_rv())
+        innovation_sd = jnp.asarray(self.innovation_sd_rv())
+        autoreg_broadcast = jnp.broadcast_to(autoreg, (n_processes,))
 
         stationary_sd = innovation_sd / jnp.sqrt(1 - autoreg**2)
         with numpyro.plate(f"{name_prefix}_init_plate", n_processes):
@@ -326,19 +323,31 @@ def sample(
         if n_timepoints == 1:
             return init_states[jnp.newaxis, :]
 
-        scale_broadcast = jnp.broadcast_to(jnp.asarray(innovation_sd), (n_processes,))
-        post_init = numpyro.sample(
-            f"{name_prefix}_state",
-            StateAR1(
-                autoreg=autoreg_broadcast,
-                scale=scale_broadcast,
-                initial_loc=init_states,
-                num_steps=n_timepoints - 1,
-            ),
+        def transition(
+            previous_state: jax.Array, unused: None
+        ) -> tuple[jax.Array, jax.Array]:
+            """
+            Sample the next AR(1) state.
+
+            Returns
+            -------
+            tuple[jax.Array, jax.Array]
+                The next carry and recorded state.
+            """
+            transition_mean = autoreg_broadcast * previous_state
+            next_state = numpyro.sample(
+                f"{name_prefix}_state",
+                dist.Normal(transition_mean, innovation_sd),
+            )
+            return next_state, next_state
+
+        _, post_init = scan(
+            transition,
+            init_states,
+            xs=None,
+            length=n_timepoints - 1,
         )
-        x = jnp.concatenate([init_states[:, jnp.newaxis], post_init], axis=-1)
-        # ensure time is the leading axis (length n_timepoints)
-        return jnp.moveaxis(x, -1, 0)
+        return jnp.concatenate([init_states[jnp.newaxis, :], post_init], axis=0)
 
 
 class DifferencedAR1(TemporalProcess):
@@ -472,9 +481,9 @@ def sample(
         """
         initial_value = _prepare_initial_value(initial_value, n_processes)
 
-        autoreg = self.autoreg_rv()
-        innovation_sd = self.innovation_sd_rv()
-        autoreg_broadcast = jnp.broadcast_to(jnp.asarray(autoreg), (n_processes,))
+        autoreg = jnp.asarray(self.autoreg_rv())
+        innovation_sd = jnp.asarray(self.innovation_sd_rv())
+        autoreg_broadcast = jnp.broadcast_to(autoreg, (n_processes,))
 
         stationary_sd = innovation_sd / jnp.sqrt(1 - autoreg**2)
         with numpyro.plate(f"{name_prefix}_init_rate_plate", n_processes):
@@ -499,25 +508,38 @@ def sample(
         x1 = initial_value + init_rates
 
         if n_timepoints == 2:
-            # ensure time is the leading axis (length n_timepoints)
-            return jnp.moveaxis(jnp.stack([initial_value, x1], axis=-1), -1, 0)
-
-        scale_broadcast = jnp.broadcast_to(jnp.asarray(innovation_sd), (n_processes,))
-        post_init = numpyro.sample(
-            f"{name_prefix}_state",
-            StateDifferencedAR1(
-                autoreg=autoreg_broadcast,
-                scale=scale_broadcast,
-                initial_loc=initial_value,
-                initial_diff=init_rates,
-                num_steps=n_timepoints - 2,
-            ),
+            return jnp.stack([initial_value, x1], axis=0)
+
+        def transition(
+            carry: tuple[jax.Array, jax.Array], unused: None
+        ) -> tuple[tuple[jax.Array, jax.Array], jax.Array]:
+            """
+            Sample the next differenced AR(1) state.
+
+            Returns
+            -------
+            tuple[tuple[jax.Array, jax.Array], jax.Array]
+                The updated two-state carry and recorded state.
+            """
+            state_before_previous, previous_state = carry
+            previous_difference = previous_state - state_before_previous
+            transition_mean = previous_state + autoreg_broadcast * previous_difference
+            next_state = numpyro.sample(
+                f"{name_prefix}_state",
+                dist.Normal(transition_mean, innovation_sd),
+            )
+            return (previous_state, next_state), next_state
+
+        _, post_init = scan(
+            transition,
+            (initial_value, x1),
+            xs=None,
+            length=n_timepoints - 2,
         )
-        full_path = jnp.concatenate(
-            [initial_value[:, jnp.newaxis], x1[:, jnp.newaxis], post_init], axis=-1
+        return jnp.concatenate(
+            [initial_value[jnp.newaxis, :], x1[jnp.newaxis, :], post_init],
+            axis=0,
         )
-        # ensure time is the leading axis (length n_timepoints)
-        return jnp.moveaxis(full_path, -1, 0)
 
 
 class RandomWalk(TemporalProcess):
@@ -626,7 +648,7 @@ def sample(
         """
         initial_value = _prepare_initial_value(initial_value, n_processes)
 
-        innovation_sd = self.innovation_sd_rv()
+        innovation_sd = jnp.asarray(self.innovation_sd_rv())
 
         if self.parameterization == "innovation":
             rw = ProcessRandomWalk(
@@ -648,18 +670,30 @@ def sample(
         if n_timepoints == 1:
             return initial_value[jnp.newaxis, :]
 
-        scale_broadcast = jnp.broadcast_to(jnp.asarray(innovation_sd), (n_processes,))
-        post_init = numpyro.sample(
-            f"{name_prefix}_state",
-            StateRandomWalk(
-                scale=scale_broadcast,
-                initial_loc=initial_value,
-                num_steps=n_timepoints - 1,
-            ),
+        def transition(
+            previous_state: jax.Array, unused: None
+        ) -> tuple[jax.Array, jax.Array]:
+            """
+            Sample the next random-walk state.
+
+            Returns
+            -------
+            tuple[jax.Array, jax.Array]
+                The next carry and recorded state.
+            """
+            next_state = numpyro.sample(
+                f"{name_prefix}_state",
+                dist.Normal(previous_state, innovation_sd),
+            )
+            return next_state, next_state
+
+        _, post_init = scan(
+            transition,
+            initial_value,
+            xs=None,
+            length=n_timepoints - 1,
         )
-        x = jnp.concatenate([initial_value[:, jnp.newaxis], post_init], axis=-1)
-        # ensure time is the leading axis (length n_timepoints)
-        return jnp.moveaxis(x, -1, 0)
+        return jnp.concatenate([initial_value[jnp.newaxis, :], post_init], axis=0)
 
 
 class StepwiseTemporalProcess(TemporalProcess):
diff --git a/test/test_temporal_processes.py b/test/test_temporal_processes.py
index c51e261e..377ad0f2 100644
--- a/test/test_temporal_processes.py
+++ b/test/test_temporal_processes.py
@@ -2,6 +2,8 @@
 Unit tests for temporal processes.
 """
 
+from typing import Literal
+
 import jax
 import jax.numpy as jnp
 import numpyro
@@ -16,6 +18,7 @@
     DifferencedAR1,
     RandomWalk,
     StepwiseTemporalProcess,
+    TemporalProcess,
     WeeklyTemporalProcess,
 )
 from pyrenew.latent.state_centered_distributions import (
@@ -165,6 +168,112 @@ def _differenced_ar1_expected_path(
 
 PARAMETERIZATIONS = ["innovation", "state"]
 
+ProcessName = Literal["random_walk", "ar1", "differenced_ar1"]
+Parameterization = Literal["innovation", "state"]
+
+
+def _make_forecast_process(
+    process_name: ProcessName,
+    parameterization: Parameterization,
+) -> TemporalProcess:
+    """
+    Build a temporal process with deterministic forecast parameters.
+
+    Returns
+    -------
+    TemporalProcess
+        The configured process.
+    """
+    if process_name == "random_walk":
+        return RandomWalk(
+            **fixed_rw_kwargs(innovation_sd=0.2),
+            parameterization=parameterization,
+        )
+    return (AR1 if process_name == "ar1" else DifferencedAR1)(
+        **fixed_ar1_kwargs(autoreg=0.6, innovation_sd=0.2),
+        parameterization=parameterization,
+    )
+
+
+def _state_forecast_posterior(
+    process_name: ProcessName,
+    fitted_states: jax.Array,
+) -> tuple[dict[str, jax.Array], jax.Array, int]:
+    """
+    Build posterior samples and trajectory metadata for a state process.
+
+    Returns
+    -------
+    tuple[dict[str, jax.Array], jax.Array, int]
+        Posterior samples, initial value, and state trajectory offset.
+    """
+    initial_value = jnp.array([0.25, -0.5])
+    if process_name == "random_walk":
+        return {"forecast_state": fitted_states}, initial_value, 1
+    if process_name == "ar1":
+        init_states = jnp.array([[0.1, -0.2], [0.4, 0.3]])
+        return (
+            {
+                "forecast_init": init_states,
+                "forecast_state": fitted_states,
+            },
+            initial_value,
+            1,
+        )
+
+    init_rates = jnp.array([[0.15, -0.1], [-0.2, 0.05]])
+    return (
+        {
+            "forecast_init_rate": init_rates,
+            "forecast_state": fitted_states,
+        },
+        initial_value,
+        2,
+    )
+
+
+def _transition_mean(
+    process_name: ProcessName,
+    fitted_states: jax.Array,
+) -> jax.Array:
+    """
+    Compute the next state-transition mean from fitted terminal states.
+
+    Returns
+    -------
+    jax.Array
+        The next transition mean for every posterior draw and process.
+    """
+    if process_name == "random_walk":
+        return fitted_states[:, -1, :]
+    if process_name == "ar1":
+        return 0.6 * fitted_states[:, -1, :]
+    return fitted_states[:, -1, :] + 0.6 * (
+        fitted_states[:, -1, :] - fitted_states[:, -2, :]
+    )
+
+
+def _forecast_sample_sites(
+    process_name: ProcessName,
+    parameterization: Parameterization,
+) -> tuple[str, ...]:
+    """
+    Return the fitted latent sites required to replay a trajectory.
+
+    Returns
+    -------
+    tuple[str, ...]
+        Names of the posterior sample sites required for prediction.
+    """
+    transition_site = "state" if parameterization == "state" else "noise_decentered"
+    if process_name == "random_walk":
+        transition_site = "state" if parameterization == "state" else "step"
+        return (f"forecast_{transition_site}",)
+    if process_name == "ar1":
+        return ("forecast_init", f"forecast_{transition_site}")
+    return ("forecast_init_rate", f"forecast_{transition_site}")
+
+
 DETERMINISTIC_INIT_PROCESS_PARAMS = [
     (DifferencedAR1, fixed_ar1_kwargs()),
     (RandomWalk, fixed_rw_kwargs()),
@@ -911,8 +1020,8 @@ def model():
         ).get_trace()
         state_site = traced["rw_state"]["value"]
         path = traced["path"]["value"]
-        assert state_site.shape == (2, 5)
-        assert jnp.allclose(state_site, path[1:].T)
+        assert state_site.shape == (5, 2)
+        assert jnp.allclose(state_site, path[1:])
 
     @pytest.mark.parametrize(
         "innovation_sd",
@@ -993,12 +1102,12 @@ def test_trace_has_init_and_state_sites_not_noise(self):
         assert "ar1_noise" not in traced
 
     def test_state_site_shape(self):
-        """The state site holds the post-initial path of shape ``(n_processes, n_timepoints - 1)``."""
+        """The state site has time leading over the process dimension."""
         ar1 = AR1(**fixed_ar1_kwargs(), parameterization="state")
         traced = numpyro.handlers.trace(
             numpyro.handlers.seed(ar1.sample, rng_seed=0)
         ).get_trace(n_timepoints=12, n_processes=3, name_prefix="ar1")
-        assert traced["ar1_state"]["value"].shape == (3, 11)
+        assert traced["ar1_state"]["value"].shape == (11, 3)
 
     @pytest.mark.parametrize("autoreg,innovation_sd", [(0.5, 0.05), (0.9, 0.1)])
     def test_prior_moments_match_innovation_parameterization(
@@ -1105,12 +1214,12 @@ def test_trace_has_init_rate_and_state_sites_not_noise(self):
         assert "diff_noise" not in traced
 
     def test_state_site_shape(self):
-        """The state site holds the post-initial path of shape ``(n_processes, n_timepoints - 2)``."""
+        """The state site has time leading over the process dimension."""
         d = DifferencedAR1(**fixed_ar1_kwargs(), parameterization="state")
         traced = numpyro.handlers.trace(
             numpyro.handlers.seed(d.sample, rng_seed=0)
         ).get_trace(n_timepoints=12, n_processes=3, name_prefix="diff")
-        assert traced["diff_state"]["value"].shape == (3, 10)
+        assert traced["diff_state"]["value"].shape == (10, 3)
 
     def test_second_row_equals_initial_value_plus_init_rate(self):
         """``x[1]`` equals ``initial_value`` plus the unbundled initial difference."""
@@ -1186,6 +1295,208 @@ def model_innov():
         )
 
 
+class TestStatePosteriorExtension:
+    """Posterior prediction extends time-leading state-centered sites."""
+
+    @pytest.mark.parametrize(
+        "process_name",
+        ["random_walk", "ar1", "differenced_ar1"],
+    )
+    def test_fitted_states_are_preserved_as_forecast_prefix(
+        self,
+        process_name: ProcessName,
+    ) -> None:
+        """A longer forecast replays fitted states and samples only its suffix."""
+        n_draws = 2
+        n_processes = 2
+        fitted_n_timepoints = 5
+        forecast_n_timepoints = 8
+        state_offset = 2 if process_name == "differenced_ar1" else 1
+        fitted_scanned_timepoints = fitted_n_timepoints - state_offset
+        forecast_scanned_timepoints = forecast_n_timepoints - state_offset
+        fitted_states = jnp.arange(
+            n_draws * fitted_scanned_timepoints * n_processes,
+            dtype=jnp.asarray(0.2).dtype,
+        ).reshape(n_draws, fitted_scanned_timepoints, n_processes)
+        fitted_states /= 10.0
+        posterior_samples, initial_value, posterior_state_offset = (
+            _state_forecast_posterior(process_name, fitted_states)
+        )
+        assert posterior_state_offset == state_offset
+        process = _make_forecast_process(process_name, "state")
+
+        def model(n_timepoints: int) -> None:
+            """Record the public state-process trajectory."""
+            trajectory = process.sample(
+                n_timepoints=n_timepoints,
+                n_processes=n_processes,
+                initial_value=initial_value,
+                name_prefix="forecast",
+            )
+            numpyro.deterministic("trajectory", trajectory)
+
+        forecast = Predictive(
+            model,
+            posterior_samples=posterior_samples,
+            return_sites=("trajectory", "forecast_state"),
+        )(jax.random.PRNGKey(10), forecast_n_timepoints)
+
+        assert forecast["trajectory"].shape == (
+            n_draws,
+            forecast_n_timepoints,
+            n_processes,
+        )
+        assert forecast["forecast_state"].shape == (
+            n_draws,
+            forecast_scanned_timepoints,
+            n_processes,
+        )
+        assert jnp.array_equal(
+            forecast["forecast_state"][:, :fitted_scanned_timepoints, :],
+            fitted_states,
+        )
+        assert jnp.array_equal(
+            forecast["trajectory"][:, state_offset:fitted_n_timepoints, :],
+            fitted_states,
+        )
+        assert (
+            forecast["forecast_state"][:, fitted_scanned_timepoints:, :].shape[1]
+            == forecast_n_timepoints - fitted_n_timepoints
+        )
+
+    @pytest.mark.parametrize(
+        "process_name",
+        ["random_walk", "ar1", "differenced_ar1"],
+    )
+    def test_first_new_state_uses_fitted_terminal_carry(
+        self,
+        process_name: ProcessName,
+    ) -> None:
+        """The first extension transition starts from the fitted terminal state."""
+        n_draws = 2
+        n_processes = 2
+        fitted_n_timepoints = 5
+        forecast_n_timepoints = 6
+        state_offset = 2 if process_name == "differenced_ar1" else 1
+        fitted_scanned_timepoints = fitted_n_timepoints - state_offset
+        fitted_states_a = jnp.arange(
+            n_draws * fitted_scanned_timepoints * n_processes,
+            dtype=jnp.asarray(0.2).dtype,
+        ).reshape(n_draws, fitted_scanned_timepoints, n_processes)
+        fitted_states_a /= 10.0
+        terminal_shift = jnp.array([[0.5, -0.25], [-0.4, 0.75]])
+        fitted_states_b = fitted_states_a.at[:, -1, :].add(terminal_shift)
+        posterior_a, initial_value, _ = _state_forecast_posterior(
+            process_name, fitted_states_a
+        )
+        posterior_b, _, _ = _state_forecast_posterior(process_name, fitted_states_b)
+        process = _make_forecast_process(process_name, "state")
+
+        def model(n_timepoints: int) -> None:
+            """Record the public state-process trajectory."""
+            trajectory = process.sample(
+                n_timepoints=n_timepoints,
+                n_processes=n_processes,
+                initial_value=initial_value,
+                name_prefix="forecast",
+            )
+            numpyro.deterministic("trajectory", trajectory)
+
+        predictive_a = Predictive(
+            model,
+            posterior_samples=posterior_a,
+            return_sites=("trajectory", "forecast_state"),
+        )
+        predictive_b = Predictive(
+            model,
+            posterior_samples=posterior_b,
+            return_sites=("trajectory", "forecast_state"),
+        )
+        key = jax.random.PRNGKey(20)
+        forecast_a = predictive_a(key, forecast_n_timepoints)
+        forecast_b = predictive_b(key, forecast_n_timepoints)
+
+        observed_difference = (
+            forecast_b["forecast_state"][:, fitted_scanned_timepoints, :]
+            - forecast_a["forecast_state"][:, fitted_scanned_timepoints, :]
+        )
+        expected_difference = _transition_mean(
+            process_name, fitted_states_b
+        ) - _transition_mean(process_name, fitted_states_a)
+        assert jnp.allclose(observed_difference, expected_difference)
+
+
+class TestForecastHorizonParity:
+    """State and innovation forms share the public forecasting contract."""
+
+    @pytest.mark.parametrize("parameterization", PARAMETERIZATIONS)
+    @pytest.mark.parametrize(
+        "process_name",
+        ["random_walk", "ar1", "differenced_ar1"],
+    )
+    def test_equal_longer_and_shorter_forecast_horizons(
+        self,
+        process_name: ProcessName,
+        parameterization: Parameterization,
+    ) -> None:
+        """Fitted paths replay or extend, while unsliced shorter paths fail."""
+        n_draws = 2
+        n_processes = 2
+        fitted_n_timepoints = 5
+        forecast_n_timepoints = 8
+        initial_value = jnp.array([0.25, -0.5])
+        process = _make_forecast_process(process_name, parameterization)
+
+        def model(n_timepoints: int) -> None:
+            """Record a temporal-process trajectory for posterior prediction."""
+            trajectory = process.sample(
+                n_timepoints=n_timepoints,
+                n_processes=n_processes,
+                initial_value=initial_value,
+                name_prefix="forecast",
+            )
+            numpyro.deterministic("trajectory", trajectory)
+
+        fitted = Predictive(model, num_samples=n_draws)(
+            jax.random.PRNGKey(30), fitted_n_timepoints
+        )
+        posterior_samples = {
+            site_name: fitted[site_name]
+            for site_name in _forecast_sample_sites(process_name, parameterization)
+        }
+        posterior_predictive = Predictive(
+            model,
+            posterior_samples=posterior_samples,
+            return_sites=("trajectory",),
+        )
+
+        replayed = posterior_predictive(jax.random.PRNGKey(31), fitted_n_timepoints)[
+            "trajectory"
+        ]
+        extended = posterior_predictive(jax.random.PRNGKey(32), forecast_n_timepoints)[
+            "trajectory"
+        ]
+
+        assert replayed.shape == (
+            n_draws,
+            fitted_n_timepoints,
+            n_processes,
+        )
+        assert jnp.allclose(replayed, fitted["trajectory"])
+        assert extended.shape == (
+            n_draws,
+            forecast_n_timepoints,
+            n_processes,
+        )
+        assert jnp.allclose(
+            extended[:, :fitted_n_timepoints, :],
+            fitted["trajectory"],
+        )
+
+        with pytest.raises(RuntimeError, match="requires length"):
+            posterior_predictive(jax.random.PRNGKey(33), fitted_n_timepoints - 1)
+
+
 class TestStepwiseTemporalProcessConstruction:
     """Construction-time validation for StepwiseTemporalProcess."""
 

From 1d1e07760884ae6399da73df062dda4feff17021 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 15 Sep 2026 20:07:04 +0000
Subject: [PATCH 3/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 AGENTS.md | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index 51fe97c3..bdaae94b 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -5,8 +5,11 @@ Avoid overly general and extensible frameworks.
 Always ask if a design or code change requires backwards compatibility.
 In most cases, backwards compatibility is not a concern.
 
-coding style:   always use MyPy type hints.   all functions should have a docstring; public-facing API functions should be user-focussed.  users want to know what the function does, ditto maintainers.  omit design discussion details.
-  inline comments only for critical, non-obvious details.
+coding style: always use MyPy type hints.
+all functions should have a docstring; public-facing API functions should be user-focussed.
+users want to know what the function does, ditto maintainers.
+omit design discussion details.
+inline comments only for critical, non-obvious details.
 
 Design and implementation plans should be written in plain style.
 describe math in terms of python computation.

From f6704da8a61ffe50a97babcd95381a46c9b15c51 Mon Sep 17 00:00:00 2001
From: Mitzi Morris <bf06@cdc.gov>
Date: Tue, 15 Sep 2026 19:45:38 -0400
Subject: [PATCH 4/6] changes per code review

---
 AGENTS.md                                     |   7 +-
 pyrenew/latent/__init__.py                    |   8 -
 .../latent/state_centered_distributions.py    | 420 ------------------
 pyrenew/latent/temporal_processes.py          |  67 ++-
 test/test_temporal_processes.py               | 322 --------------
 5 files changed, 51 insertions(+), 773 deletions(-)
 delete mode 100644 pyrenew/latent/state_centered_distributions.py

diff --git a/AGENTS.md b/AGENTS.md
index 51fe97c3..bdaae94b 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -5,8 +5,11 @@ Avoid overly general and extensible frameworks.
 Always ask if a design or code change requires backwards compatibility.
 In most cases, backwards compatibility is not a concern.
 
-coding style:   always use MyPy type hints.   all functions should have a docstring; public-facing API functions should be user-focussed.  users want to know what the function does, ditto maintainers.  omit design discussion details.
-  inline comments only for critical, non-obvious details.
+coding style: always use MyPy type hints.
+all functions should have a docstring; public-facing API functions should be user-focussed.
+users want to know what the function does, ditto maintainers.
+omit design discussion details.
+inline comments only for critical, non-obvious details.
 
 Design and implementation plans should be written in plain style.
 describe math in terms of python computation.
diff --git a/pyrenew/latent/__init__.py b/pyrenew/latent/__init__.py
index 2fbe4837..b3b04a51 100644
--- a/pyrenew/latent/__init__.py
+++ b/pyrenew/latent/__init__.py
@@ -28,11 +28,6 @@
 from pyrenew.latent.infections import Infections
 from pyrenew.latent.infectionswithfeedback import InfectionsWithFeedback
 from pyrenew.latent.population_infections import PopulationInfections
-from pyrenew.latent.state_centered_distributions import (
-    StateAR1,
-    StateDifferencedAR1,
-    StateRandomWalk,
-)
 from pyrenew.latent.subpopulation_infections import SubpopulationInfections
 from pyrenew.latent.temporal_processes import (
     AR1,
@@ -71,9 +66,6 @@
     "AR1",
     "DifferencedAR1",
     "RandomWalk",
-    "StateAR1",
-    "StateDifferencedAR1",
-    "StateRandomWalk",
     "StepwiseTemporalProcess",
     "WeeklyTemporalProcess",
 ]
diff --git a/pyrenew/latent/state_centered_distributions.py b/pyrenew/latent/state_centered_distributions.py
deleted file mode 100644
index 93c4ed14..00000000
--- a/pyrenew/latent/state_centered_distributions.py
+++ /dev/null
@@ -1,420 +0,0 @@
-"""NumPyro distributions for state-centered temporal-process priors."""
-
-from __future__ import annotations
-
-import jax
-import jax.numpy as jnp
-from jax import lax, random
-from jax.typing import ArrayLike
-from numpyro.distributions import constraints
-from numpyro.distributions.continuous import GaussianRandomWalk, Normal
-from numpyro.distributions.distribution import Distribution
-from numpyro.distributions.util import validate_sample
-from numpyro.util import is_prng_key
-
-
-class StateRandomWalk(Distribution):
-    r"""
-    State-centered random walk on a post-initial state path.
-
-    Given a deterministic initial state $x_0$ = ``initial_loc``:
-
-    $$
-    x_t \sim \mathrm{Normal}(x_{t-1}, \sigma), \quad t = 1, \dots, T
-    $$
-
-    The sampled value is the post-initial path
-    $[x_1, x_2, \ldots, x_{\mathrm{num\_steps}}]$ of length ``num_steps``.
-
-    This is a location shift of [numpyro.distributions.continuous.GaussianRandomWalk][].
-    A zero-mean Gaussian random walk supplies the path, and ``initial_loc``
-    offsets it so the walk is centered on the initial state rather than zero.
-    """
-
-    arg_constraints = {
-        "scale": constraints.positive,
-        "initial_loc": constraints.real,
-    }
-    support = constraints.real_vector
-    reparametrized_params = ["scale", "initial_loc"]
-    pytree_data_fields = ("gaussian_random_walk_", "initial_loc")
-
-    def __init__(
-        self,
-        scale: ArrayLike,
-        initial_loc: ArrayLike = 0.0,
-        num_steps: int = 1,
-        *,
-        validate_args: bool | None = None,
-    ) -> None:
-        """Construct a state-centered random-walk distribution."""
-        if not isinstance(num_steps, int) or num_steps <= 0:
-            raise ValueError(f"num_steps must be a positive integer; got {num_steps!r}")
-        self.initial_loc = initial_loc
-        self.gaussian_random_walk_ = GaussianRandomWalk(scale, num_steps)
-
-        batch_shape = lax.broadcast_shapes(
-            jnp.shape(scale),
-            jnp.shape(initial_loc),
-        )
-        super().__init__(batch_shape, (num_steps,), validate_args=validate_args)
-
-    @property
-    def scale(self) -> ArrayLike:
-        """
-        Innovation standard deviation of the underlying random walk.
-
-        Returns
-        -------
-        ArrayLike
-            The innovation standard deviation $\\sigma$.
-        """
-        return self.gaussian_random_walk_.scale
-
-    @property
-    def num_steps(self) -> int:
-        """
-        Length of the post-initial state path.
-
-        Returns
-        -------
-        int
-            The number of post-initial steps.
-        """
-        return self.gaussian_random_walk_.num_steps
-
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
-        """
-        Forward-sample a post-initial random-walk state path.
-
-        Returns
-        -------
-        ArrayLike
-            Array of shape ``sample_shape + batch_shape + (num_steps,)``.
-        """
-        assert is_prng_key(key)
-        walk = self.gaussian_random_walk_.sample(key, sample_shape)
-        # initial_loc carries batch_shape and walk's last axis is time
-        # (num_steps); add a trailing length-1 time axis so initial_loc aligns
-        # on batch rather than right-aligning against the time axis.
-        return jnp.expand_dims(jnp.asarray(self.initial_loc), -1) + walk
-
-    @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
-        """
-        Compute the log-density of an observed post-initial state path.
-
-        Parameters
-        ----------
-        value
-            Post-initial path of shape
-            ``sample_shape + batch_shape + (num_steps,)``.
-
-        Returns
-        -------
-        ArrayLike
-            Log-density of shape ``sample_shape + batch_shape``.
-        """
-        # initial_loc carries batch_shape and value's last axis is time
-        # (num_steps); add a trailing length-1 time axis so the offset aligns
-        # on batch rather than right-aligning against the time axis.
-        offset = jnp.expand_dims(jnp.asarray(self.initial_loc), -1)
-        return self.gaussian_random_walk_.log_prob(value - offset)
-
-
-class StateAR1(Distribution):
-    r"""
-    State-centered AR(1) process for a post-initial state path.
-
-    Generative form, given a deterministic initial state $x_0$ = ``initial_loc``:
-
-    $$
-    x_t \sim \mathrm{Normal}(\phi \, x_{t-1}, \sigma), \quad t = 1, \dots, T
-    $$
-
-    where $\phi$ is ``autoreg`` and $\sigma$ is ``scale``.
-
-    The sampled value is the post-initial path
-    $[x_1, x_2, \ldots, x_{\mathrm{num\_steps}}]$ of length ``num_steps``.
-    The initial state $x_0$ is not part of the sample; it is supplied as
-    ``initial_loc`` and used to score the first transition. A random initial
-    state drawn from the stationary distribution can be handled by the calling
-    temporal process as a separate sample site, matching the innovation
-    parameterization.
-
-    Parameters
-    ----------
-    autoreg
-        AR(1) coefficient $\phi$. For stationarity, $|\phi| < 1$; this is
-        not enforced.
-    scale
-        Innovation standard deviation $\sigma$. Must be positive.
-    initial_loc
-        Deterministic initial state $x_0$. Used to score the first
-        transition; not itself sampled. Defaults to ``0.0``.
-    num_steps
-        Length of the post-initial path. Must be a positive integer.
-    validate_args
-        Forwarded to the base [numpyro.distributions.distribution.Distribution][].
-    """
-
-    arg_constraints = {
-        "autoreg": constraints.real,
-        "scale": constraints.positive,
-        "initial_loc": constraints.real,
-    }
-    support = constraints.real_vector
-    reparametrized_params = ["autoreg", "scale", "initial_loc"]
-    pytree_aux_fields = ("num_steps",)
-
-    def __init__(
-        self,
-        autoreg: ArrayLike,
-        scale: ArrayLike,
-        initial_loc: ArrayLike = 0.0,
-        num_steps: int = 1,
-        *,
-        validate_args: bool | None = None,
-    ) -> None:
-        """
-        Construct a state-centered AR(1) distribution.
-
-        Raises
-        ------
-        ValueError
-            If ``num_steps`` is not a positive integer.
-        """
-        if not isinstance(num_steps, int) or num_steps <= 0:
-            raise ValueError(f"num_steps must be a positive integer; got {num_steps!r}")
-        self.autoreg = autoreg
-        self.scale = scale
-        self.initial_loc = initial_loc
-        self.num_steps = num_steps
-
-        batch_shape = lax.broadcast_shapes(
-            jnp.shape(autoreg),
-            jnp.shape(scale),
-            jnp.shape(initial_loc),
-        )
-        super().__init__(batch_shape, (num_steps,), validate_args=validate_args)
-
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
-        """
-        Forward-sample a post-initial state path.
-
-        Returns
-        -------
-        ArrayLike
-            Array of shape ``sample_shape + batch_shape + (num_steps,)``.
-        """
-        assert is_prng_key(key)
-
-        per_step_shape = sample_shape + self.batch_shape
-        autoreg = jnp.asarray(self.autoreg)
-        scale = jnp.asarray(self.scale)
-        # lax.scan requires a shape-invariant carry; broadcast the initial
-        # state to per_step_shape so it matches the per-step output shape.
-        initial_loc = jnp.broadcast_to(jnp.asarray(self.initial_loc), per_step_shape)
-
-        noise = random.normal(key, shape=(self.num_steps,) + per_step_shape)
-
-        def step(
-            prev: ArrayLike, z_t: ArrayLike
-        ) -> tuple[ArrayLike, ArrayLike]:  # numpydoc ignore=GL08
-            new = autoreg * prev + scale * z_t
-            return new, new
-
-        _, xs = lax.scan(step, initial_loc, noise)
-        return jnp.moveaxis(
-            xs, 0, -1
-        )  # ensure time is the trailing axis, length num_steps
-
-    @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
-        """
-        Compute the log-density of an observed post-initial state path.
-
-        Parameters
-        ----------
-        value
-            Post-initial path of shape
-            ``sample_shape + batch_shape + (num_steps,)``.
-
-        Returns
-        -------
-        ArrayLike
-            Log-density of shape ``sample_shape + batch_shape``.
-        """
-        scale = jnp.asarray(self.scale)
-        autoreg = jnp.asarray(self.autoreg)
-        initial_loc = jnp.asarray(self.initial_loc)
-
-        per_step_shape = jnp.shape(value)[:-1]
-        initial_loc = jnp.broadcast_to(initial_loc, per_step_shape)
-
-        # add length-1 trailing time dimension to fixed-in-time parameters
-        initial_loc_t = jnp.expand_dims(initial_loc, -1)
-        scale_t = jnp.expand_dims(scale, -1)
-        autoreg_t = jnp.expand_dims(autoreg, -1)
-
-        v = jnp.concatenate([initial_loc_t, value], axis=-1)
-        step_locs = autoreg_t * v[..., :-1]
-        step_probs = Normal(step_locs, scale_t).log_prob(v[..., 1:])
-        return jnp.sum(step_probs, axis=-1)
-
-
-class StateDifferencedAR1(Distribution):
-    r"""
-    State-centered differenced AR(1) prior on a length-``num_steps`` post-initial path.
-
-    Generative form, given a deterministic initial state $x_0$ = ``initial_loc``
-    and initial first difference $d_1$ = ``initial_diff``, so that
-    $x_1 = x_0 + d_1$:
-
-    $$
-    x_t \sim \mathrm{Normal}(x_{t-1} + \phi \, (x_{t-1} - x_{t-2}), \sigma),
-    \quad t \geq 2
-    $$
-
-    where $\phi$ is ``autoreg`` and $\sigma$ is ``scale``.
-
-    The sampled value is the post-initial path
-    $[x_2, x_3, \ldots, x_{\mathrm{num\_steps} + 1}]$ of length ``num_steps``.
-    The initial state $x_0$ and first difference $d_1$ are not part of the
-    sample; they are supplied at construction and used to score the first
-    post-initial transition. A random initial difference drawn from the
-    stationary distribution can be handled by the calling temporal process as a
-    separate sample site, matching the innovation parameterization.
-
-    Parameters
-    ----------
-    autoreg
-        AR(1) coefficient $\phi$ on first differences. For stationarity,
-        $|\phi| < 1$; this is not enforced.
-    scale
-        Innovation standard deviation $\sigma$. Must be positive.
-    initial_loc
-        Deterministic initial state $x_0$. Not itself sampled. Defaults to
-        ``0.0``.
-    initial_diff
-        Deterministic initial first difference $d_1 = x_1 - x_0$. Together with
-        ``initial_loc`` it fixes $x_1 = x_0 + d_1$, used to score the first
-        post-initial transition. Not itself sampled. Defaults to ``0.0``.
-    num_steps
-        Length of the post-initial path. Must be a positive integer.
-    validate_args
-        Forwarded to the base [numpyro.distributions.distribution.Distribution][].
-    """
-
-    arg_constraints = {
-        "autoreg": constraints.real,
-        "scale": constraints.positive,
-        "initial_loc": constraints.real,
-        "initial_diff": constraints.real,
-    }
-    support = constraints.real_vector
-    reparametrized_params = ["autoreg", "scale", "initial_loc", "initial_diff"]
-    pytree_aux_fields = ("num_steps",)
-
-    def __init__(
-        self,
-        autoreg: ArrayLike,
-        scale: ArrayLike,
-        initial_loc: ArrayLike = 0.0,
-        initial_diff: ArrayLike = 0.0,
-        num_steps: int = 1,
-        *,
-        validate_args: bool | None = None,
-    ) -> None:
-        """
-        Construct a state-centered differenced AR(1) distribution.
-
-        Raises
-        ------
-        ValueError
-            If ``num_steps`` is not a positive integer.
-        """
-        if not isinstance(num_steps, int) or num_steps <= 0:
-            raise ValueError(f"num_steps must be a positive integer; got {num_steps!r}")
-        self.autoreg = autoreg
-        self.scale = scale
-        self.initial_loc = initial_loc
-        self.initial_diff = initial_diff
-        self.num_steps = num_steps
-
-        batch_shape = lax.broadcast_shapes(
-            jnp.shape(autoreg),
-            jnp.shape(scale),
-            jnp.shape(initial_loc),
-            jnp.shape(initial_diff),
-        )
-        event_shape = (num_steps,)
-        super().__init__(batch_shape, event_shape, validate_args=validate_args)
-
-    def sample(self, key: jax.Array, sample_shape: tuple[int, ...] = ()) -> ArrayLike:
-        """
-        Forward-sample a post-initial path.
-
-        Returns
-        -------
-        ArrayLike
-            Array of shape ``sample_shape + batch_shape + (num_steps,)``.
-        """
-        assert is_prng_key(key)
-
-        per_step_shape = sample_shape + self.batch_shape
-        autoreg = jnp.asarray(self.autoreg)
-        scale = jnp.asarray(self.scale)
-        # lax.scan requires a shape-invariant carry; broadcast the two initial
-        # states to per_step_shape so they match the per-step output shape.
-        x0 = jnp.broadcast_to(jnp.asarray(self.initial_loc), per_step_shape)
-        x1 = x0 + jnp.broadcast_to(jnp.asarray(self.initial_diff), per_step_shape)
-
-        noise = random.normal(key, shape=(self.num_steps,) + per_step_shape)
-
-        def step(
-            carry: tuple[ArrayLike, ArrayLike], z_t: ArrayLike
-        ) -> tuple[tuple[ArrayLike, ArrayLike], ArrayLike]:  # numpydoc ignore=GL08
-            prev_2, prev_1 = carry
-            new = prev_1 + autoreg * (prev_1 - prev_2) + scale * z_t
-            return (prev_1, new), new
-
-        _, xs = lax.scan(step, (x0, x1), noise)
-        return jnp.moveaxis(xs, 0, -1)
-
-    @validate_sample
-    def log_prob(self, value: ArrayLike) -> ArrayLike:
-        """
-        Compute the log-density of an observed post-initial path.
-
-        Parameters
-        ----------
-        value
-            Post-initial path of shape
-            ``sample_shape + batch_shape + (num_steps,)``.
-
-        Returns
-        -------
-        ArrayLike
-            Log-density of shape ``sample_shape + batch_shape``.
-        """
-        scale = jnp.asarray(self.scale)
-        autoreg = jnp.asarray(self.autoreg)
-        x0 = jnp.asarray(self.initial_loc)
-        x1 = x0 + jnp.asarray(self.initial_diff)
-
-        per_step_shape = jnp.shape(value)[:-1]
-        x0 = jnp.broadcast_to(x0, per_step_shape)
-        x1 = jnp.broadcast_to(x1, per_step_shape)
-
-        # add length-1 trailing time dimension to fixed-in-time parameters
-        x0_t = jnp.expand_dims(x0, -1)
-        x1_t = jnp.expand_dims(x1, -1)
-        scale_t = jnp.expand_dims(scale, -1)
-        autoreg_t = jnp.expand_dims(autoreg, -1)
-
-        v = jnp.concatenate([x0_t, x1_t, value], axis=-1)
-        prev_delta = v[..., 1:-1] - v[..., :-2]
-        means = v[..., 1:-1] + autoreg_t * prev_delta
-        step_probs = Normal(means, scale_t).log_prob(v[..., 2:])
-        return jnp.sum(step_probs, axis=-1)
diff --git a/pyrenew/latent/temporal_processes.py b/pyrenew/latent/temporal_processes.py
index 48f72e08..5603ae8e 100644
--- a/pyrenew/latent/temporal_processes.py
+++ b/pyrenew/latent/temporal_processes.py
@@ -20,6 +20,21 @@
 | Vectorization | Caller manages array shapes | Automatic via ``n_processes`` parameter |
 | Validation | Minimal constraints | Validates positive innovation_sd |
 
+``AR1``, ``DifferencedAR1``, and ``RandomWalk`` expose a common public
+interface for two equivalent parameterizations. The ``"innovation"``
+parameterization is non-centered: it samples independent standardized
+innovations and constructs the state trajectory deterministically using the
+lower-level [pyrenew.process][] implementations. The ``"state"``
+parameterization is centered: it samples each state conditionally on the
+preceding state using a local NumPyro scan.
+
+This implementation asymmetry is intentional. The existing
+[pyrenew.process][] abstractions naturally represent innovation-driven
+processes. State-centered transitions remain local to the public temporal
+processes so that the state itself is the scanned NumPyro sample site, which
+supports replaying a fitted trajectory and extending it during posterior
+prediction.
+
 **When to use which:**
 
 - Use ``pyrenew.process`` classes (``ARProcess``, ``DifferencedProcess``,
@@ -193,13 +208,16 @@ class AR1(TemporalProcess):
     toward a mean level. Keeps Rt bounded near a baseline — values that
     drift away are "pulled back" over time.
 
-    This class wraps [pyrenew.process.ARProcess][] with a simplified,
-    protocol-compliant interface that handles vectorization automatically.
+    The ``parameterization`` argument selects between a non-centered
+    ``"innovation"`` parameterization, which samples independent standardized
+    innovations, and a centered ``"state"`` parameterization, which samples
+    each state conditionally on the preceding state. Both produce the same
+    prior distribution over the state path but can have different posterior
+    sampling geometry.
 
-    The ``parameterization`` argument selects between sampling standardized
-    innovations (``"innovation"``) and sampling the state path directly
-    (``"state"``). Both produce the same prior distribution over the state
-    path; they differ in sampler geometry.
+    The innovation branch uses [pyrenew.process.ARProcess][] to construct the
+    trajectory from sampled innovations. The state branch uses local scanned
+    conditional transitions.
 
     Parameters
     ----------
@@ -358,13 +376,11 @@ class DifferencedAR1(TemporalProcess):
     the rate of change reverting toward a mean. Unlike AR(1), this allows
     Rt to trend persistently upward or downward while the growth rate stabilizes.
 
-    This class wraps [pyrenew.process.DifferencedProcess][] with
-    [pyrenew.process.ARProcess][] as the fundamental process, providing
-    a simplified, protocol-compliant interface.
-
-    The ``parameterization`` argument selects between sampling standardized
-    innovations on the differences (``"innovation"``) and sampling the state
-    path $[x_1, \ldots, x_{T-1}]$ directly (``"state"``) under the priors
+    The ``parameterization`` argument selects between a non-centered
+    ``"innovation"`` parameterization, which samples independent standardized
+    innovations on the differences, and a centered ``"state"``
+    parameterization, which samples the state path
+    $[x_1, \ldots, x_{T-1}]$ conditionally under the priors
 
     $$
     x_1 \sim \mathrm{Normal}(x_0, \sigma / \sqrt{1 - \phi^2})
@@ -377,7 +393,12 @@ class DifferencedAR1(TemporalProcess):
 
     where $\phi$ is ``autoreg`` and $\sigma$ is ``innovation_sd``. The initial
     state $x_0$ is supplied deterministically as ``initial_value``. Both produce
-    the same prior over the state path; they differ in sampler geometry.
+    the same prior over the state path but can have different posterior
+    sampling geometry.
+
+    The innovation branch uses [pyrenew.process.DifferencedProcess][] with
+    [pyrenew.process.ARProcess][] as its fundamental process. The state branch
+    uses local scanned conditional transitions.
 
     Parameters
     ----------
@@ -550,13 +571,17 @@ class RandomWalk(TemporalProcess):
     toward a mean. Allows Rt to drift without bound — suitable when you
     have no prior expectation that Rt will return to a baseline.
 
-    This class wraps [pyrenew.process.RandomWalk][] with a simplified,
-    protocol-compliant interface that handles vectorization automatically.
-
-    The ``parameterization`` argument selects between sampling standardized
-    innovations (``"innovation"``) and sampling the state path directly
-    (``"state"``), with ``x[0] = initial_value`` deterministic. Both produce
-    the same prior over the state path; they differ in sampler geometry.
+    The ``parameterization`` argument selects between a non-centered
+    ``"innovation"`` parameterization, which samples independent standardized
+    innovations, and a centered ``"state"`` parameterization, which samples
+    each state conditionally on the preceding state. The initial state
+    ``x[0] = initial_value`` is deterministic. Both parameterizations produce
+    the same prior over the state path but can have different posterior
+    sampling geometry.
+
+    The innovation branch uses [pyrenew.process.RandomWalk][] to construct the
+    trajectory from sampled innovations. The state branch uses local scanned
+    conditional transitions.
 
     Parameters
     ----------
diff --git a/test/test_temporal_processes.py b/test/test_temporal_processes.py
index 377ad0f2..47b5de53 100644
--- a/test/test_temporal_processes.py
+++ b/test/test_temporal_processes.py
@@ -9,7 +9,6 @@
 import numpyro
 import numpyro.distributions as dist
 import pytest
-from jax.typing import ArrayLike
 from numpyro.infer import Predictive
 
 from pyrenew.deterministic import DeterministicVariable
@@ -21,11 +20,6 @@
     TemporalProcess,
     WeeklyTemporalProcess,
 )
-from pyrenew.latent.state_centered_distributions import (
-    StateAR1,
-    StateDifferencedAR1,
-    StateRandomWalk,
-)
 from pyrenew.randomvariable import DistributionalVariable
 from pyrenew.time import MMWR_WEEK
 
@@ -59,107 +53,6 @@ def fixed_rw_kwargs(innovation_sd=0.05):
     }
 
 
-def _ar1_expected_path(
-    noise: ArrayLike,
-    autoreg: ArrayLike,
-    scale: ArrayLike,
-    initial_loc: ArrayLike,
-) -> ArrayLike:
-    """
-    Replay the StateAR1 state recursion from explicit standard-normal noise.
-
-    Parameters
-    ----------
-    noise
-        Standard-normal innovations with leading axis ``num_steps`` and trailing
-        per-step (sample + batch) shape.
-    autoreg
-        AR(1) coefficient.
-    scale
-        Innovation standard deviation.
-    initial_loc
-        Deterministic initial state.
-
-    Returns
-    -------
-    ArrayLike
-        Post-initial path with time on the trailing axis.
-    """
-    autoreg = jnp.asarray(autoreg)
-    scale = jnp.asarray(scale)
-    state = jnp.broadcast_to(jnp.asarray(initial_loc), noise.shape[1:])
-    states = []
-    for innovation in noise:
-        state = autoreg * state + scale * innovation
-        states.append(state)
-    return jnp.moveaxis(jnp.stack(states, axis=0), 0, -1)
-
-
-def _differenced_ar1_expected_path(
-    noise: ArrayLike,
-    autoreg: ArrayLike,
-    scale: ArrayLike,
-    initial_loc: ArrayLike,
-    initial_diff: ArrayLike,
-) -> ArrayLike:
-    """
-    Replay the StateDifferencedAR1 recursion from explicit standard-normal noise.
-
-    Parameters
-    ----------
-    noise
-        Standard-normal innovations with leading axis ``num_steps`` and trailing
-        per-step (sample + batch) shape.
-    autoreg
-        AR(1) coefficient on first differences.
-    scale
-        Innovation standard deviation.
-    initial_loc
-        Deterministic initial state.
-    initial_diff
-        Deterministic initial first difference.
-
-    Returns
-    -------
-    ArrayLike
-        Post-initial path with time on the trailing axis.
-    """
-    autoreg = jnp.asarray(autoreg)
-    scale = jnp.asarray(scale)
-    per_step_shape = noise.shape[1:]
-    prev_2 = jnp.broadcast_to(jnp.asarray(initial_loc), per_step_shape)
-    prev_1 = prev_2 + jnp.broadcast_to(jnp.asarray(initial_diff), per_step_shape)
-    states = []
-    for innovation in noise:
-        state = prev_1 + autoreg * (prev_1 - prev_2) + scale * innovation
-        states.append(state)
-        prev_2, prev_1 = prev_1, state
-    return jnp.moveaxis(jnp.stack(states, axis=0), 0, -1)
-
-
-AR1_FORWARD_SAMPLING_CASES = [
-    (0.5, 0.2, 1.0, 5, ()),
-    (0.9, 0.05, -0.5, 8, ()),
-    (jnp.array([0.3, 0.8]), jnp.array([0.1, 0.5]), jnp.array([0.0, 1.0]), 6, ()),
-    (0.5, jnp.array([0.1, 0.2, 0.3]), 0.0, 4, ()),
-    (0.7, 0.2, 1.0, 3, (4,)),
-]
-
-DIFFERENCED_AR1_FORWARD_SAMPLING_CASES = [
-    (0.2, 0.5, 1.0, 0.3, 1, ()),
-    (0.6, 0.1, 0.0, 0.05, 7, ()),
-    (
-        jnp.array([0.6, -0.3]),
-        jnp.array([0.2, 0.5]),
-        jnp.array([1.0, -0.5]),
-        jnp.array([0.1, 0.2]),
-        5,
-        (),
-    ),
-    (0.5, 0.2, 0.0, 0.0, 4, (3,)),
-]
-
-
 INNER_PROCESS_PARAMS = [
     (AR1, fixed_ar1_kwargs()),
     (DifferencedAR1, fixed_ar1_kwargs()),
@@ -280,221 +173,6 @@ def _forecast_sample_sites(
 ]
 
 
-class TestStateCenteredDistributionLogProb:
-    """Exact density checks for state-centered temporal-process distributions."""
-
-    def test_state_random_walk_log_prob_matches_manual_transition_sum(self):
-        """Batched StateRandomWalk log_prob equals the explicit RW transition density."""
-        scale = jnp.array([0.3, 0.7])
-        initial_loc = jnp.array([1.0, -0.5])
-        value = jnp.array(
-            [
-                [1.2, 0.6, 0.1, -0.2],
-                [-0.3, 0.4, 0.0, 0.2],
-            ]
-        )
-
-        distribution = StateRandomWalk(
-            scale=scale,
-            initial_loc=initial_loc,
-            num_steps=value.shape[-1],
-        )
-
-        full_path = jnp.concatenate([initial_loc[:, None], value], axis=-1)
-        expected = dist.Normal(full_path[:, :-1], scale[:, None]).log_prob(value)
-        expected = expected.sum(axis=-1)
-
-        assert jnp.allclose(distribution.log_prob(value), expected)
-
-    def test_state_ar1_log_prob_matches_manual_transition_sum(self):
-        """Batched StateAR1 log_prob equals the explicit AR1 transition density."""
-        autoreg = jnp.array([0.4, -0.2])
-        scale = jnp.array([0.3, 0.7])
-        initial_loc = jnp.array([1.0, -0.5])
-        value = jnp.array(
-            [
-                [1.2, 0.6, 0.1, -0.2],
-                [-0.3, 0.4, 0.0, 0.2],
-            ]
-        )
-
-        distribution = StateAR1(
-            autoreg=autoreg,
-            scale=scale,
-            initial_loc=initial_loc,
-            num_steps=value.shape[-1],
-        )
-
-        full_path = jnp.concatenate([initial_loc[:, None], value], axis=-1)
-        transition_locs = autoreg[:, None] * full_path[:, :-1]
-        expected = dist.Normal(transition_locs, scale[:, None]).log_prob(value)
-        expected = expected.sum(axis=-1)
-
-        assert jnp.allclose(distribution.log_prob(value), expected)
-
-    def test_state_differenced_ar1_log_prob_matches_manual_transition_sum(self):
-        """Batched StateDifferencedAR1 log_prob equals the explicit transition density."""
-        autoreg = jnp.array([0.6, -0.3])
-        scale = jnp.array([0.2, 0.5])
-        initial_loc = jnp.array([1.0, -0.5])
-        initial_diff = jnp.array([0.1, 0.2])
-        value = jnp.array(
-            [
-                [1.45, 1.7, 1.6],
-                [-0.1, -0.2, 0.0],
-            ]
-        )
-
-        distribution = StateDifferencedAR1(
-            autoreg=autoreg,
-            scale=scale,
-            initial_loc=initial_loc,
-            initial_diff=initial_diff,
-            num_steps=value.shape[-1],
-        )
-
-        x1 = initial_loc + initial_diff
-        full_path = jnp.concatenate([initial_loc[:, None], x1[:, None], value], axis=-1)
-        previous_delta = full_path[:, 1:-1] - full_path[:, :-2]
-        transition_locs = full_path[:, 1:-1] + autoreg[:, None] * previous_delta
-        transition_probs = dist.Normal(transition_locs, scale[:, None]).log_prob(
-            full_path[:, 2:]
-        )
-        expected = transition_probs.sum(axis=-1)
-
-        assert jnp.allclose(distribution.log_prob(value), expected)
-
-
-class TestStateCenteredDistributionValidation:
-    """Focused coverage for state-centered distribution validation branches."""
-
-    @pytest.mark.parametrize(
-        "distribution_cls,kwargs",
-        [
-            (StateRandomWalk, {"scale": 1.0}),
-            (StateAR1, {"autoreg": 0.5, "scale": 1.0}),
-            (StateDifferencedAR1, {"autoreg": 0.5, "scale": 1.0}),
-        ],
-    )
-    @pytest.mark.parametrize("invalid_num_steps", [0, 1.5])
-    def test_num_steps_must_be_positive_integer(
-        self, distribution_cls, kwargs, invalid_num_steps
-    ):
-        """Constructors reject non-positive and non-integer step counts."""
-        with pytest.raises(ValueError, match="num_steps must be a positive integer"):
-            distribution_cls(**kwargs, num_steps=invalid_num_steps)
-
-
-class TestStateCenteredDistributionForwardSampling:
-    """Forward ``sample`` returns the full post-initial state trajectory."""
-
-    @pytest.mark.parametrize(
-        "autoreg,scale,initial_loc,num_steps,sample_shape",
-        AR1_FORWARD_SAMPLING_CASES,
-    )
-    def test_state_ar1_sample_matches_recursion(
-        self, autoreg, scale, initial_loc, num_steps, sample_shape
-    ):
-        """StateAR1 sampling reproduces the AR(1) state recursion exactly."""
-        key = jax.random.PRNGKey(7)
-        distribution = StateAR1(
-            autoreg=autoreg,
-            scale=scale,
-            initial_loc=initial_loc,
-            num_steps=num_steps,
-        )
-
-        sample = distribution.sample(key, sample_shape=sample_shape)
-        per_step_shape = sample_shape + distribution.batch_shape
-        noise = jax.random.normal(key, shape=(num_steps,) + per_step_shape)
-        expected = _ar1_expected_path(noise, autoreg, scale, initial_loc)
-
-        assert sample.shape == per_step_shape + (num_steps,)
-        assert jnp.allclose(sample, expected)
-
-    @pytest.mark.parametrize(
-        "autoreg,scale,initial_loc,initial_diff,num_steps,sample_shape",
-        DIFFERENCED_AR1_FORWARD_SAMPLING_CASES,
-    )
-    def test_state_differenced_ar1_sample_matches_recursion(
-        self, autoreg, scale, initial_loc, initial_diff, num_steps, sample_shape
-    ):
-        """StateDifferencedAR1 sampling reproduces the differenced recursion exactly."""
-        key = jax.random.PRNGKey(19)
-        distribution = StateDifferencedAR1(
-            autoreg=autoreg,
-            scale=scale,
-            initial_loc=initial_loc,
-            initial_diff=initial_diff,
-            num_steps=num_steps,
-        )
-
-        sample = distribution.sample(key, sample_shape=sample_shape)
-        per_step_shape = sample_shape + distribution.batch_shape
-        noise = jax.random.normal(key, shape=(num_steps,) + per_step_shape)
-        expected = _differenced_ar1_expected_path(
-            noise, autoreg, scale, initial_loc, initial_diff
-        )
-
-        assert sample.shape == per_step_shape + (num_steps,)
-        assert jnp.allclose(sample, expected)
-
-    @pytest.mark.parametrize(
-        "scale,initial_loc",
-        [
-            (0.3, 1.0),
-            (jnp.array([0.2, 0.5]), jnp.array([1.0, -0.5])),
-        ],
-    )
-    def test_state_random_walk_sample_is_location_shift(self, scale, initial_loc):
-        """Shifting initial_loc shifts the whole sampled path by that constant."""
-        key = jax.random.PRNGKey(11)
-        num_steps = 6
-        shift = 4.0
-        base = StateRandomWalk(
-            scale=scale, initial_loc=initial_loc, num_steps=num_steps
-        )
-        shifted = StateRandomWalk(
-            scale=scale,
-            initial_loc=jnp.asarray(initial_loc) + shift,
-            num_steps=num_steps,
-        )
-
-        base_path = base.sample(key)
-        shifted_path = shifted.sample(key)
-
-        assert base_path.shape == base.batch_shape + (num_steps,)
-        assert jnp.allclose(shifted_path - base_path, shift)
-
-    def test_state_random_walk_sample_scales_with_scale(self):
-        """The centered sampled path scales linearly with the innovation scale."""
-        key = jax.random.PRNGKey(13)
-        num_steps = 6
-        initial_loc = 2.0
-        small = StateRandomWalk(scale=0.5, initial_loc=initial_loc, num_steps=num_steps)
-        large = StateRandomWalk(scale=1.5, initial_loc=initial_loc, num_steps=num_steps)
-
-        small_centered = small.sample(key) - initial_loc
-        large_centered = large.sample(key) - initial_loc
-
-        assert jnp.allclose(large_centered, 3.0 * small_centered)
-
-    def test_state_random_walk_sample_variance_grows_linearly(self):
-        """Forward samples accumulate variance linearly in time, as a random walk does."""
-        key = jax.random.PRNGKey(17)
-        num_steps = 10
-        scale = 0.4
-        distribution = StateRandomWalk(
-            scale=scale, initial_loc=0.0, num_steps=num_steps
-        )
-
-        samples = distribution.sample(key, sample_shape=(20000,))
-        empirical_var = samples.var(axis=0)
-        expected_var = scale**2 * jnp.arange(1, num_steps + 1)
-
-        assert jnp.allclose(empirical_var, expected_var, rtol=0.1)
-
-
 class TestTemporalProcessVectorizedSampling:
     """Parameterization-agnostic shape behavior across all temporal process types."""
 

From 759a3e2dcbf9ab8a46c08924dc0b33bd36b72348 Mon Sep 17 00:00:00 2001
From: Mitzi Morris <bf06@cdc.gov>
Date: Fri, 18 Sep 2026 11:04:13 -0400
Subject: [PATCH 5/6] changes per code review

---
 test/test_temporal_processes.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/test/test_temporal_processes.py b/test/test_temporal_processes.py
index 47b5de53..c94de699 100644
--- a/test/test_temporal_processes.py
+++ b/test/test_temporal_processes.py
@@ -996,7 +996,6 @@ def test_fitted_states_are_preserved_as_forecast_prefix(
             n_draws * fitted_scanned_timepoints * n_processes,
             dtype=jnp.asarray(0.2).dtype,
         ).reshape(n_draws, fitted_scanned_timepoints, n_processes)
-        fitted_states /= 10.0
         posterior_samples, initial_value, posterior_state_offset = (
             _state_forecast_posterior(process_name, fitted_states)
         )
@@ -1061,7 +1060,6 @@ def test_first_new_state_uses_fitted_terminal_carry(
             n_draws * fitted_scanned_timepoints * n_processes,
             dtype=jnp.asarray(0.2).dtype,
         ).reshape(n_draws, fitted_scanned_timepoints, n_processes)
-        fitted_states_a /= 10.0
         terminal_shift = jnp.array([[0.5, -0.25], [-0.4, 0.75]])
         fitted_states_b = fitted_states_a.at[:, -1, :].add(terminal_shift)
         posterior_a, initial_value, _ = _state_forecast_posterior(
@@ -1098,9 +1096,11 @@ def model(n_timepoints: int) -> None:
             forecast_b["forecast_state"][:, fitted_scanned_timepoints, :]
             - forecast_a["forecast_state"][:, fitted_scanned_timepoints, :]
         )
-        expected_difference = _transition_mean(
-            process_name, fitted_states_b
-        ) - _transition_mean(process_name, fitted_states_a)
+        transition_mean_a = _transition_mean(process_name, fitted_states_a)
+        transition_mean_b = _transition_mean(process_name, fitted_states_b)
+        assert not jnp.allclose(transition_mean_a, transition_mean_b)
+
+        expected_difference = transition_mean_b - transition_mean_a
         assert jnp.allclose(observed_difference, expected_difference)
 
 

From c477d4f927c55980541d1667db5e560f078fd699 Mon Sep 17 00:00:00 2001
From: Mitzi Morris <bf06@cdc.gov>
Date: Fri, 18 Sep 2026 12:08:39 -0400
Subject: [PATCH 6/6] fix merge conflict

---
 AGENTS.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index bdaae94b..24eb66bb 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -7,8 +7,7 @@ In most cases, backwards compatibility is not a concern.
 
 coding style: always use MyPy type hints.
 all functions should have a docstring; public-facing API functions should be user-focussed.
-users want to know what the function does, ditto maintainers.
-omit design discussion details.
+users and developers want to know what the function does, not design details.
 inline comments only for critical, non-obvious details.
 
 Design and implementation plans should be written in plain style.