Let `log_rt_time_0_rv` give each subpopulation its own starting Rt

Start with pyrenew/latent/subpopulation_infections.py, especially sample and _validate_and_prepare_I0, then review test/test_subpopulation_infections.py and docs/tutorials/latent_subpopulation_infections.qmd. Implement the stated scalar or per-subpopulation shape behavior for initial Rt and I0, preserving scalar draw order. Done means the listed tests pass and the tutorial describes the vector option.

Two constructor arguments of SubpopulationInfections set each subpopulation's state on the first day of the initialization period. I0_rv sets the infections per resident. log_rt_time_0_rv sets the log Rt, which also sets how fast infections grow during the initialization period. I0_rv can return one value for all subpopulations or one value per subpopulation. log_rt_time_0_rv can return only one value, so every subpopulation starts with the same Rt and the same growth rate. This issue lets log_rt_time_0_rv return one value per subpopulation as well.

Background

SubpopulationInfections.sample builds subpopulation log Rt from a baseline and deviations that sum to zero:

log_rt_baseline = baseline_rt_process.sample(..., initial_value=log_rt_time_0_rv())
deviations_raw = subpop_rt_deviation_process.sample(..., initial_value=jnp.zeros(n_subpops))
deviations = deviations_raw - jnp.mean(deviations_raw, axis=1, keepdims=True)
log_rt_subpop = log_rt_baseline + deviations

• The deviations start at zero. RandomWalk and DifferencedAR1 start at initial_value, so with either as the deviation process, every subpopulation starts the initialization period with the same Rt. AR1 draws its first value around initial_value, here zero.
• Every subpopulation also starts with the same growth rate. Initial infections are projected backward with a growth rate computed from rt_subpop[0, :].
• A vector log_rt_time_0_rv fails. It is passed to the single baseline process and raises ValueError: Incompatible types for broadcasting.

Motivation

• Fits that start mid-epidemic. When subpopulations are at different stages of an epidemic on the first day, the deviation process has to move each one away from the shared starting Rt, and the backward projection gives all of them the same initial growth.
• Roy et al. (2025). Their model gives each region its own starting reproduction number, $\log R_{0,i} \sim \mathcal{N}(-0.1, 0.5)$, independently for each region. #901 implements their coupled renewal.

What the user writes

One starting value for all subpopulations, unchanged:

builder.configure_latent(
    SubpopulationInfections,
    ...,
    log_rt_time_0_rv=DistributionalVariable(
        "log_rt_time_0", dist.Normal(0.0, 0.5)
    ),
)

One starting value per subpopulation, new:

builder.configure_latent(
    SubpopulationInfections,
    ...,
    log_rt_time_0_rv=DistributionalVariable(
        "log_rt_time_0", dist.Normal(0.0, 0.5).expand([n_subpops])
    ),
)

As with a vector I0_rv, a vector log_rt_time_0_rv fixes the number of subpopulations when the model is built. Its length must equal the length of subpop_fractions at sample time, and its entries are in the same order.

Behavior

• Scalar log_rt_time_0_rv: unchanged. The baseline starts at the scalar, and the deviations start at zero.
• Vector v of shape (n_subpops,): the baseline starts at jnp.mean(v), and the deviations start at v - jnp.mean(v). Those starting deviations already sum to zero, so centering leaves them unchanged, and log_rt_subpop starts at v.
  • With RandomWalk or DifferencedAR1 for both temporal processes, rt_subpop[0] == jnp.exp(v) and rt_baseline[0] == jnp.exp(jnp.mean(v)).
  • With AR1, the first value is drawn around the starting value, as it is today for a scalar.
• Initial growth: the backward projection already uses rt_subpop[0, :], so each subpopulation gets its own initial growth rate with no further change.
• Any other shape: ValueError when the latent process is sampled. This covers a vector of the wrong length, including a shape (1,) array when there is more than one subpopulation, and a 2-D array. The message names the shape received and n_subpops. Shapes are static, so the check works during JAX tracing.
• I0_rv: follows the same shape rule: a scalar or an array of shape (n_subpops,). Any other shape raises ValueError when the latent process is sampled, with a message of the same form. Today a wrong length fails with a broadcasting error, and a shape (1,) array is used for all subpopulations.
• Recorded values: names and shapes are unchanged. With a vector, rt_baseline[0] is the geometric mean of jnp.exp(v), consistent with the class docstring's statement that the baseline is the geometric mean of the subpopulation Rt.
• PopulationInfections: unchanged. Its log_rt_time_0_rv stays a scalar.

Implementation

pyrenew/latent/subpopulation_infections.py

Add a private method that splits the starting log Rt into the starting values of the two temporal processes:

def _split_initial_log_rt(
    self,
    initial_log_rt: ArrayLike,
    pop: PopulationStructure,
) -> tuple[ArrayLike, ArrayLike]:
    """
    Split the starting log Rt into baseline and deviation starting values.
    """
    if initial_log_rt.ndim == 0:
        return initial_log_rt, jnp.zeros(pop.n_subpops)
    if initial_log_rt.shape == (pop.n_subpops,):
        baseline_start = jnp.mean(initial_log_rt)
        return baseline_start, initial_log_rt - baseline_start
    raise ValueError(
        "log_rt_time_0_rv must return a scalar or an array of shape "
        f"({pop.n_subpops},), one value per subpopulation; "
        f"got shape {initial_log_rt.shape}."
    )

In sample, call it on jnp.asarray(self.log_rt_time_0_rv()), and pass its two results as initial_value to baseline_rt_process.sample and subpop_rt_deviation_process.sample. The rest of sample is unchanged.

In _validate_and_prepare_I0, after the scalar case, raise ValueError if I0.shape != (pop.n_subpops,), with a message that names I0_rv, (n_subpops,), and the shape received.

• Docstrings. The log_rt_time_0_rv parameter: the starting log Rt on the first day of the initialization period, either one value for all subpopulations or one value per subpopulation. The class docstring's mathematical form gains the starting values.
• Draw order. log_rt_time_0_rv is called once, in the same place as now, so seeded runs with a scalar give identical draws.

Unchanged

• PopulationInfections, PyrenewBuilder, MultiSignalModel, and the temporal processes.

Tests

In test/test_subpopulation_infections.py, using RandomWalk processes unless noted:

All existing tests must pass unchanged; they cover the scalar case. Construction repeated in more than two tests goes in a conftest.py fixture.

Documentation

docs/tutorials/latent_subpopulation_infections.qmd: the log_rt_time_0_rv bullet currently says it is the initial value of the baseline. It changes to say it can be one value per subpopulation, with a short example.

Breaking change

A scalar log_rt_time_0_rv behaves as now, and a vector of length n_subpops currently raises an error. A shape (1,) array, for log_rt_time_0_rv or I0_rv, which today is used as one value for all subpopulations, raises an error unless there is one subpopulation.

Out of scope

• PopulationInfections, which has one population.
• Per-subpopulation step sizes for the deviation process. RandomWalk passes innovation_sd_rv() to Normal(jnp.zeros(n_processes), innovation_sd), so an innovation_sd_rv of length n_subpops already gives one step size per subpopulation.
• Seeding dates that differ by subpopulation, as in Roy et al.'s England fits.
• Changes to how initial infections are projected backward.