Latent

AR1

AR1(
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
)

Bases: TemporalProcess

AR(1) process.

Each value depends on the previous value plus noise, with reversion toward a mean level. Keeps Rt bounded near a baseline — values that drift away are "pulled back" over time.

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 innovation branch uses pyrenew.process.ARProcess to construct the trajectory from sampled innovations. The state branch uses local scanned conditional transitions.

Parameters:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of noise at each time step. Larger values produce more volatile trajectories; smaller values produce smoother ones.	required
parameterization	Parameterization	Which latent object to sample: "innovation" (default) or "state".	'innovation'

Initialize AR(1) process.

Parameters:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of innovations.	required
parameterization	Parameterization	"innovation" (default) or "state". See class docstring.	'innovation'

Raises:

Type	Description
TypeError	If autoreg_rv or innovation_sd_rv are not RandomVariable instances
ValueError	If innovation_sd_rv is a DeterministicVariable with any value <= 0, or if parameterization is not a recognized string

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
) -> None:
    """
    Initialize AR(1) process.

    Parameters
    ----------
    autoreg_rv
        RandomVariable that returns the autoregressive coefficient. For
        stationarity, |autoreg| < 1, but this is not enforced (use priors
        to constrain if needed).
    innovation_sd_rv
        RandomVariable that returns the standard deviation of innovations.
    parameterization
        ``"innovation"`` (default) or ``"state"``. See class docstring.

    Raises
    ------
    TypeError
        If autoreg_rv or innovation_sd_rv are not RandomVariable instances
    ValueError
        If innovation_sd_rv is a DeterministicVariable with any value <= 0,
        or if parameterization is not a recognized string
    """
    if not isinstance(autoreg_rv, RandomVariable):
        raise TypeError("autoreg_rv must be a RandomVariable")
    if not isinstance(innovation_sd_rv, RandomVariable):
        raise TypeError("innovation_sd_rv must be a RandomVariable")
    _validate_deterministic_innovation_sd(innovation_sd_rv)
    _validate_parameterization(parameterization)
    self.autoreg_rv = autoreg_rv
    self.innovation_sd_rv = innovation_sd_rv
    self.parameterization = parameterization
    self.ar_process = ARProcess(name="ar1")

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"AR1(autoreg_rv={self.autoreg_rv}, "
        f"innovation_sd_rv={self.innovation_sd_rv}, "
        f"parameterization={self.parameterization!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample AR(1) trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Mean of the initial-state prior. The first returned value is sampled as Normal(initial_value, innovation_sd / sqrt(1 - autoreg**2)). Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites	'ar1'
first_day_dow	int | None	Unused. See pyrenew.latent.TemporalProcess.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample AR(1) trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Mean of the initial-state prior. The first returned value is sampled
        as ``Normal(initial_value, innovation_sd / sqrt(1 - autoreg**2))``.
        Scalar values are broadcast to all processes; arrays must have
        shape ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites
    first_day_dow
        Unused. See [pyrenew.latent.TemporalProcess][].

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    initial_value = _prepare_initial_value(initial_value, 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):
        init_states = numpyro.sample(
            f"{name_prefix}_init",
            dist.Normal(initial_value, stationary_sd),
        )

    if self.parameterization == "innovation":
        return self.ar_process(
            n=n_timepoints,
            init_vals=init_states[jnp.newaxis, :],
            autoreg=autoreg_broadcast[jnp.newaxis, :],
            noise_sd=innovation_sd,
            noise_name=f"{name_prefix}_noise",
        )

    if n_timepoints == 1:
        return init_states[jnp.newaxis, :]

    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,
    )
    return jnp.concatenate([init_states[jnp.newaxis, :], post_init], axis=0)

BaseLatentInfectionProcess

BaseLatentInfectionProcess(
    *, name: str, gen_int_rv: RandomVariable, n_initialization_points: int
)

Bases: RandomVariable

Base class for latent infection processes with subpopulation structure.

Provides common functionality for hierarchical and partitioned infection models: - Population fraction validation and parsing (at sample time) - Standard output structure via LatentSample

All subclasses return infections as a LatentSample named tuple with fields: (aggregate, all_subpops). Observation processes are responsible for selecting which subpopulations they observe via indexing.

The constructor specifies model structure (generation interval, priors, temporal processes). Population structure (subpop_fractions) is provided at sample time, allowing a single model to be fit to multiple jurisdictions.

Parameters:

Name	Type	Description	Default
gen_int_rv	RandomVariable	Generation interval PMF	required
n_initialization_points	int	Number of initialization days before the first observation day. Latent and observation arrays use a shared padded time axis with element 0 at the start of this initialization period. In observation natural coordinates, day 0 is the first observed data day; on the shared padded axis, that same day is index n_initialization_points. Must be at least len(gen_int_rv()) to provide enough history for the renewal equation convolution.	required

Notes

Population structure (subpop_fractions) is passed to the sample() method, not the constructor. This allows a single model instance to be fit to multiple datasets with different jurisdiction structures.

When using PyrenewBuilder (recommended), n_initialization_points is computed automatically from all observation processes. When constructing latent processes directly, you must specify n_initialization_points explicitly.

Initialize base latent infection process.

Parameters:

Name	Type	Description	Default
name	str	A name for this random variable.	required
gen_int_rv	RandomVariable	Generation interval PMF	required
n_initialization_points	int	Number of initialization days; see the class-level parameter documentation for the shared padded-axis convention.	required

Raises:

Type	Description
ValueError	If gen_int_rv is None or n_initialization_points is insufficient.

Source code in pyrenew/latent/base.py

def __init__(
    self,
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
) -> None:
    """
    Initialize base latent infection process.

    Parameters
    ----------
    name
        A name for this random variable.
    gen_int_rv
        Generation interval PMF
    n_initialization_points
        Number of initialization days; see the class-level parameter
        documentation for the shared padded-axis convention.

    Raises
    ------
    ValueError
        If gen_int_rv is None or n_initialization_points is insufficient.
    """
    super().__init__(name=name)
    if gen_int_rv is None:
        raise ValueError("gen_int_rv is required")
    self.gen_int_rv = gen_int_rv

    gen_int_length = len(self.gen_int_rv())
    if n_initialization_points < gen_int_length:
        raise ValueError(
            f"n_initialization_points must be at least the generation "
            f"interval length ({gen_int_length}), got "
            f"{n_initialization_points}"
        )
    self.n_initialization_points = n_initialization_points

default_subpop_fractions

default_subpop_fractions() -> ArrayLike | None

Return default population fractions, or None if caller must provide them.

Override this function in order to omit specification of subpop_fractions at sample time, in which case it will be called by _parse_and_validate_fractions. Must return a valid array of positive elements that sums to 1.

Returns:

Type	Description
ArrayLike or None	Default fractions array, or None if no default exists.

Source code in pyrenew/latent/base.py

def default_subpop_fractions(self) -> ArrayLike | None:
    """
    Return default population fractions, or None if caller must provide them.

    Override this function in order to omit specification of
    subpop_fractions at sample time, in which case it will be called
    by _parse_and_validate_fractions.
    Must return a valid array of positive elements that sums to 1.

    Returns
    -------
    ArrayLike or None
        Default fractions array, or None if no default exists.
    """
    return None

get_required_lookback

get_required_lookback() -> int

Return the generation interval length for builder pattern support.

This method is used by PyrenewBuilder to compute n_initialization_points from all model components. Returns the generation interval PMF length.

Returns:

Type	Description
int	Length of generation interval PMF

Source code in pyrenew/latent/base.py

def get_required_lookback(self) -> int:
    """
    Return the generation interval length for builder pattern support.

    This method is used by PyrenewBuilder to compute n_initialization_points
    from all model components. Returns the generation interval PMF length.

    Returns
    -------
    int
        Length of generation interval PMF
    """
    return len(self.gen_int_rv())

requires_calendar_anchor

requires_calendar_anchor() -> bool

Report whether this latent process needs a calendar anchor at sample time.

The default implementation inspects this instance's attributes for :class:pyrenew.latent.temporal_processes.TemporalProcess instances and returns True if any of them has requires_calendar_anchor=True. Subclasses may override to add additional calendar-aligned components.

Returns:

Type	Description
bool	True if the caller of :meth:sample must supply a first_day_dow (derived from obs_start_date at the model entry point); False otherwise.

Source code in pyrenew/latent/base.py

def requires_calendar_anchor(self) -> bool:
    """
    Report whether this latent process needs a calendar anchor at sample time.

    The default implementation inspects this instance's attributes
    for :class:`pyrenew.latent.temporal_processes.TemporalProcess`
    instances and returns ``True`` if any of them has
    ``requires_calendar_anchor=True``. Subclasses may override to add
    additional calendar-aligned components.

    Returns
    -------
    bool
        ``True`` if the caller of :meth:`sample` must supply a
        ``first_day_dow`` (derived from ``obs_start_date`` at the
        model entry point); ``False`` otherwise.
    """
    return any(
        isinstance(value, TemporalProcess)
        and getattr(value, "requires_calendar_anchor", False)
        for value in vars(self).values()
    )

sample abstractmethod

sample(
    n_days_post_init: int,
    subpop_fractions: ArrayLike | None = None,
    **kwargs: object,
) -> LatentSample

Sample latent infections for all subpopulations.

Parameters:

Name	Type	Description	Default
n_days_post_init	int	Number of days to simulate after initialization period	required
subpop_fractions	ArrayLike | None	Population fractions for all subpopulations. Shape: (n_subpops,). Must sum to 1.0. If None, falls back to self.default_subpop_fractions().	None
**kwargs	object	Additional parameters required by specific implementations	{}

Returns:

Type	Description
LatentSample	Named tuple with fields: - aggregate: shape (n_total_days,) - all_subpops: shape (n_total_days, n_subpops) where n_total_days = n_initialization_points + n_days_post_init

Source code in pyrenew/latent/base.py

@abstractmethod
def sample(
    self,
    n_days_post_init: int,
    subpop_fractions: ArrayLike | None = None,
    **kwargs: object,
) -> LatentSample:
    """
    Sample latent infections for all subpopulations.

    Parameters
    ----------
    n_days_post_init
        Number of days to simulate after initialization period
    subpop_fractions
        Population fractions for all subpopulations.
        Shape: (n_subpops,). Must sum to 1.0.
        If None, falls back to ``self.default_subpop_fractions()``.
    **kwargs
        Additional parameters required by specific implementations

    Returns
    -------
    LatentSample
        Named tuple with fields:
        - aggregate: shape (n_total_days,)
        - all_subpops: shape (n_total_days, n_subpops)

        where n_total_days = n_initialization_points + n_days_post_init
    """
    pass  # pragma: no cover

BaseSubpopulationInfections

BaseSubpopulationInfections(
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    baseline_rt_process: TemporalProcess,
    subpop_rt_deviation_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
)

Bases: BaseLatentInfectionProcess

Base class for subpopulation renewal models with hierarchical \(\mathcal{R}(t)\).

Subpopulations share a baseline \(\mathcal{R}(t)\). The log \(\mathcal{R}(t)\) of each subpopulation is the log of the baseline plus a deviation, and the deviations sum to zero on each day, so the baseline is the geometric mean of the subpopulation values.

The starting log \(\mathcal{R}(t)\) and the initial infections are each either one value for all subpopulations or one value per subpopulation. Initial infections grow exponentially over the initialization period at the rate implied by each subpopulation's starting \(\mathcal{R}(t)\).

Each subclass defines how infections are generated from \(\mathcal{R}(t)\).

Notes

rt_baseline and rt_subpop are raw reproduction numbers generated by the temporal-process hierarchy. rt_subpop_adjusted is the reproduction number used to generate infections, as each subclass defines it.

Initialize the shared subpopulation structure.

Parameters:

Name	Type	Description	Default
name	str	Name prefix for numpyro sample sites. All deterministic quantities are recorded under this scope (e.g., "{name}::rt_baseline", "{name}::rt_subpop", and "{name}::rt_subpop_adjusted").	required
gen_int_rv	RandomVariable	Generation interval PMF	required
n_initialization_points	int	Number of initialization days before day 0.	required
I0_rv	RandomVariable	New infections per resident on the first day of the initialization period, the earliest day on the time axis. This is not prevalence. Returns either a scalar, one value for all subpopulations, or an array of shape (n_subpops,), one value per subpopulation in the order of subpop_fractions.	required
baseline_rt_process	TemporalProcess	Temporal process for baseline \(\mathcal{R}(t)\) dynamics	required
subpop_rt_deviation_process	TemporalProcess	Temporal process for subpopulation deviations	required
log_rt_time_0_rv	RandomVariable	Starting log(\(\mathcal{R}(t)\)) on the first day of the initialization period. Returns either a scalar, one value for all subpopulations, or an array of shape (n_subpops,), one value per subpopulation in the order of subpop_fractions.	required

Raises:

Type	Description
ValueError	If required parameters are missing or invalid.

Source code in pyrenew/latent/subpopulation_infections.py

def __init__(
    self,
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    baseline_rt_process: TemporalProcess,
    subpop_rt_deviation_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
) -> None:
    """
    Initialize the shared subpopulation structure.

    Parameters
    ----------
    name
        Name prefix for numpyro sample sites. All deterministic
        quantities are recorded under this scope (e.g.,
        ``"{name}::rt_baseline"``, ``"{name}::rt_subpop"``, and
        ``"{name}::rt_subpop_adjusted"``).
    gen_int_rv
        Generation interval PMF
    n_initialization_points
        Number of initialization days before day 0.
    I0_rv
        New infections per resident on the first day of the
        initialization period, the earliest day on the time axis. This
        is not prevalence. Returns either a scalar, one value for all
        subpopulations, or an array of shape ``(n_subpops,)``, one value
        per subpopulation in the order of ``subpop_fractions``.
    baseline_rt_process
        Temporal process for baseline $\\mathcal{R}(t)$ dynamics
    subpop_rt_deviation_process
        Temporal process for subpopulation deviations
    log_rt_time_0_rv
        Starting log($\\mathcal{R}(t)$) on the first day of the
        initialization period. Returns either a scalar, one value for all
        subpopulations, or an array of shape ``(n_subpops,)``, one value per
        subpopulation in the order of ``subpop_fractions``.

    Raises
    ------
    ValueError
        If required parameters are missing or invalid.
    """
    super().__init__(
        name=name,
        gen_int_rv=gen_int_rv,
        n_initialization_points=n_initialization_points,
    )

    if I0_rv is None:
        raise ValueError("I0_rv is required")
    self.I0_rv = I0_rv

    if isinstance(I0_rv, DeterministicVariable):
        self._validate_I0(I0_rv.value)

    if log_rt_time_0_rv is None:
        raise ValueError("log_rt_time_0_rv is required")
    self.log_rt_time_0_rv = log_rt_time_0_rv

    if baseline_rt_process is None:
        raise ValueError("baseline_rt_process is required")
    self.baseline_rt_process = baseline_rt_process

    if subpop_rt_deviation_process is None:
        raise ValueError("subpop_rt_deviation_process is required")
    self.subpop_rt_deviation_process = subpop_rt_deviation_process

sample

sample(
    n_days_post_init: int,
    subpop_fractions: ArrayLike | None = None,
    first_day_dow: int | None = None,
    **kwargs: object,
) -> LatentSample

Sample infections for all subpopulations.

Generates daily baseline \(\mathcal{R}(t)\), daily subpopulation deviations with a sum-to-zero constraint, and initial infections, then computes infections after the initialization period with the subclass's renewal. The temporal processes may sample parameters at any supported cadence (for example daily or weekly/stepwise), but they must return daily-length trajectories before the renewal is evaluated. The raw subpopulation Rt is recorded as rt_subpop; the Rt used to generate infections is recorded as rt_subpop_adjusted.

Parameters:

Name	Type	Description	Default
n_days_post_init	int	Number of days to simulate after initialization period	required
subpop_fractions	ArrayLike | None	Population fractions for all subpopulations. Shape: (n_subpops,). Must sum to 1.0.	None
first_day_dow	int | None	Forwarded to baseline_rt_process and subpop_rt_deviation_process. See pyrenew.latent.TemporalProcess.	None
**kwargs	object	Additional keyword arguments (unused).	{}

Returns:

Type	Description
LatentSample	Named tuple with fields: - aggregate: shape (n_total_days,) - all_subpops: shape (n_total_days, n_subpops)

Raises:

Type	Description
ValueError	If subpop_fractions is invalid, or if log_rt_time_0_rv or I0_rv returns neither a scalar nor an array of shape (n_subpops,). A subclass may raise for its own inputs.

Source code in pyrenew/latent/subpopulation_infections.py

def sample(
    self,
    n_days_post_init: int,
    subpop_fractions: ArrayLike | None = None,
    first_day_dow: int | None = None,
    **kwargs: object,
) -> LatentSample:
    """
    Sample infections for all subpopulations.

    Generates daily baseline $\\mathcal{R}(t)$, daily subpopulation
    deviations with a sum-to-zero constraint, and initial infections,
    then computes infections after the initialization period with the
    subclass's renewal. The temporal processes may sample parameters at
    any supported cadence (for example daily or weekly/stepwise), but
    they must return daily-length trajectories before the renewal is
    evaluated. The raw subpopulation Rt is recorded as ``rt_subpop``;
    the Rt used to generate infections is recorded as
    ``rt_subpop_adjusted``.

    Parameters
    ----------
    n_days_post_init
        Number of days to simulate after initialization period
    subpop_fractions
        Population fractions for all subpopulations. Shape: (n_subpops,).
        Must sum to 1.0.
    first_day_dow
        Forwarded to ``baseline_rt_process`` and
        ``subpop_rt_deviation_process``. See
        [pyrenew.latent.TemporalProcess][].
    **kwargs
        Additional keyword arguments (unused).

    Returns
    -------
    LatentSample
        Named tuple with fields:
        - aggregate: shape (n_total_days,)
        - all_subpops: shape (n_total_days, n_subpops)

    Raises
    ------
    ValueError
        If ``subpop_fractions`` is invalid, or if ``log_rt_time_0_rv``
        or ``I0_rv`` returns neither a scalar nor an array of shape
        ``(n_subpops,)``. A subclass may raise for its own inputs.
    """
    pop = self._parse_and_validate_fractions(
        subpop_fractions=subpop_fractions,
    )

    n_total_days = self.n_initialization_points + n_days_post_init

    baseline_start, deviations_start = self._split_initial_log_rt(
        jnp.asarray(self.log_rt_time_0_rv()), pop
    )

    log_rt_baseline = self.baseline_rt_process.sample(
        n_timepoints=n_total_days,
        initial_value=baseline_start,
        name_prefix="log_rt_baseline",
        first_day_dow=first_day_dow,
    )
    require_shape(log_rt_baseline, (n_total_days, 1), "baseline_rt_process")

    deviations_raw = self.subpop_rt_deviation_process.sample(
        n_timepoints=n_total_days,
        n_processes=pop.n_subpops,
        initial_value=deviations_start,
        name_prefix="subpop_deviations",
        first_day_dow=first_day_dow,
    )
    require_shape(
        deviations_raw,
        (n_total_days, pop.n_subpops),
        "subpop_rt_deviation_process",
    )

    # Sum-to-zero constraint ensures identifiability
    mean_deviation = jnp.mean(deviations_raw, axis=1, keepdims=True)
    deviations = deviations_raw - mean_deviation

    log_rt_subpop = log_rt_baseline + deviations
    rt_subpop = jnp.exp(log_rt_subpop)

    gen_int = self.gen_int_rv()

    I0_subpop = self._validate_and_prepare_I0(jnp.asarray(self.I0_rv()), pop)

    initial_r_subpop = r_approx_from_R(
        R=rt_subpop[0, :], g=gen_int, n_newton_steps=4
    )

    time_indices = jnp.arange(self.n_initialization_points)
    I0_all = I0_subpop[jnp.newaxis, :] * jnp.exp(
        initial_r_subpop[jnp.newaxis, :] * time_indices[:, jnp.newaxis]
    )

    post_init_infections_all, rt_adjusted = self._infections_from_rt(
        rt_raw=rt_subpop[self.n_initialization_points :, :],
        I0=I0_all,
        gen_int=gen_int,
        pop=pop,
    )
    rt_subpop_adjusted = jnp.vstack(
        [rt_subpop[: self.n_initialization_points, :], rt_adjusted]
    )

    infections_all = jnp.vstack([I0_all, post_init_infections_all])

    infections_aggregate = jnp.sum(
        infections_all * pop.fractions[jnp.newaxis, :], axis=1
    )

    self._validate_output_shapes(
        infections_aggregate,
        infections_all,
        n_total_days,
        pop,
    )

    with numpyro.handlers.scope(prefix=self.name, divider="::"):
        numpyro.deterministic("I0_init_all_subpops", I0_all)
        numpyro.deterministic("log_rt_baseline", log_rt_baseline)
        numpyro.deterministic("rt_baseline", jnp.exp(log_rt_baseline))
        numpyro.deterministic("rt_subpop", rt_subpop)
        numpyro.deterministic("rt_subpop_adjusted", rt_subpop_adjusted)
        numpyro.deterministic("subpop_deviations", deviations)
        numpyro.deterministic("infections_aggregate", infections_aggregate)

    return LatentSample(
        aggregate=infections_aggregate,
        all_subpops=infections_all,
    )

validate

validate() -> None

Validate hierarchical infections parameters.

Checks that the generation interval is a valid PMF.

Raises:

Type	Description
ValueError	If gen_int_rv does not return a valid discrete distribution

Source code in pyrenew/latent/subpopulation_infections.py

def validate(self) -> None:
    """
    Validate hierarchical infections parameters.

    Checks that the generation interval is a valid PMF.

    Raises
    ------
    ValueError
        If gen_int_rv does not return a valid discrete distribution
    """
    validate_discrete_dist_vector(self.gen_int_rv())

DifferencedAR1

DifferencedAR1(
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
)

Bases: TemporalProcess

AR(1) process on first differences.

Each change in value depends on the previous change plus noise, with 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.

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}) \]

\[ 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 innovation_sd. The initial state \(x_0\) is supplied deterministically as initial_value. Both produce 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:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient for differences. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of noise added to changes. Larger values produce more erratic growth rates; smaller values produce smoother trends.	required
parameterization	Parameterization	Which latent object to sample: "innovation" (default) or "state".	'innovation'

Initialize differenced AR(1) process.

Parameters:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient for differences. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of innovations.	required
parameterization	Parameterization	"innovation" (default) or "state". See class docstring.	'innovation'

Raises:

Type	Description
TypeError	If autoreg_rv or innovation_sd_rv are not RandomVariable instances
ValueError	If innovation_sd_rv is a DeterministicVariable with any value <= 0, or if parameterization is not a recognized string

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
) -> None:
    """
    Initialize differenced AR(1) process.

    Parameters
    ----------
    autoreg_rv
        RandomVariable that returns the autoregressive coefficient for
        differences. For stationarity, |autoreg| < 1, but this is not
        enforced (use priors to constrain if needed).
    innovation_sd_rv
        RandomVariable that returns the standard deviation of innovations.
    parameterization
        ``"innovation"`` (default) or ``"state"``. See class docstring.

    Raises
    ------
    TypeError
        If autoreg_rv or innovation_sd_rv are not RandomVariable instances
    ValueError
        If innovation_sd_rv is a DeterministicVariable with any value <= 0,
        or if parameterization is not a recognized string
    """
    if not isinstance(autoreg_rv, RandomVariable):
        raise TypeError("autoreg_rv must be a RandomVariable")
    if not isinstance(innovation_sd_rv, RandomVariable):
        raise TypeError("innovation_sd_rv must be a RandomVariable")
    _validate_deterministic_innovation_sd(innovation_sd_rv)
    _validate_parameterization(parameterization)
    self.autoreg_rv = autoreg_rv
    self.innovation_sd_rv = innovation_sd_rv
    self.parameterization = parameterization
    self.process = DifferencedProcess(
        name="diff_ar1",
        fundamental_process=ARProcess(name="diff_ar1_fundamental"),
        differencing_order=1,
    )

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"DifferencedAR1(autoreg_rv={self.autoreg_rv}, "
        f"innovation_sd_rv={self.innovation_sd_rv}, "
        f"parameterization={self.parameterization!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "diff_ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample differenced AR(1) trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Deterministic first state of the trajectory. Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites	'diff_ar1'
first_day_dow	int | None	Unused. See pyrenew.latent.TemporalProcess.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "diff_ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample differenced AR(1) trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Deterministic first state of the trajectory. Scalar values are
        broadcast to all processes; arrays must have shape
        ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites
    first_day_dow
        Unused. See [pyrenew.latent.TemporalProcess][].

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    initial_value = _prepare_initial_value(initial_value, 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):
        init_rates = numpyro.sample(
            f"{name_prefix}_init_rate",
            dist.Normal(0, stationary_sd),
        )

    if self.parameterization == "innovation":
        return self.process(
            n=n_timepoints,
            init_vals=initial_value[jnp.newaxis, :],
            autoreg=autoreg_broadcast[jnp.newaxis, :],
            noise_sd=innovation_sd,
            fundamental_process_init_vals=init_rates[jnp.newaxis, :],
            noise_name=f"{name_prefix}_noise",
        )

    if n_timepoints == 1:
        return initial_value[jnp.newaxis, :]

    x1 = initial_value + init_rates

    if 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,
    )
    return jnp.concatenate(
        [initial_value[jnp.newaxis, :], x1[jnp.newaxis, :], post_init],
        axis=0,
    )

GammaGroupSdPrior

GammaGroupSdPrior(
    name: str,
    sd_mean_rv: RandomVariable,
    sd_concentration_rv: RandomVariable,
    sd_min: float = 0.05,
)

Bases: RandomVariable

Gamma prior for group-level standard deviations, bounded away from zero.

Samples n_groups positive values from Gamma(concentration, rate) + sd_min.

Parameters:

Name	Type	Description	Default
name	str	Unique name for the sampled parameter in numpyro.	required
sd_mean_rv	RandomVariable	RandomVariable returning the mean of the Gamma distribution.	required
sd_concentration_rv	RandomVariable	RandomVariable returning the concentration (shape) parameter of Gamma.	required
sd_min	float	Minimum SD value (lower bound).	0.05

Initialize gamma group SD prior.

Parameters:

Name	Type	Description	Default
name	str	Unique name for the sampled parameter.	required
sd_mean_rv	RandomVariable	RandomVariable returning the mean of the Gamma distribution.	required
sd_concentration_rv	RandomVariable	RandomVariable returning the concentration parameter.	required
sd_min	float	Minimum SD value (lower bound).	0.05

Source code in pyrenew/latent/hierarchical_priors.py

def __init__(
    self,
    name: str,
    sd_mean_rv: RandomVariable,
    sd_concentration_rv: RandomVariable,
    sd_min: float = 0.05,
) -> None:
    """
    Initialize gamma group SD prior.

    Parameters
    ----------
    name
        Unique name for the sampled parameter.
    sd_mean_rv
        RandomVariable returning the mean of the Gamma distribution.
    sd_concentration_rv
        RandomVariable returning the concentration parameter.
    sd_min
        Minimum SD value (lower bound).
    """
    if not isinstance(sd_mean_rv, RandomVariable):
        raise TypeError(
            f"sd_mean_rv must be a RandomVariable, got {type(sd_mean_rv).__name__}. "
            "Use DeterministicVariable(name, value) to wrap a fixed value."
        )
    if not isinstance(sd_concentration_rv, RandomVariable):
        raise TypeError(
            f"sd_concentration_rv must be a RandomVariable, got {type(sd_concentration_rv).__name__}. "
            "Use DeterministicVariable(name, value) to wrap a fixed value."
        )
    if sd_min < 0:
        raise ValueError(f"sd_min must be non-negative, got {sd_min}")

    super().__init__(name=name)
    self.sd_mean_rv = sd_mean_rv
    self.sd_concentration_rv = sd_concentration_rv
    self.sd_min = sd_min

sample

sample(n_groups: int, **kwargs: object) -> ArrayLike

Sample group-level standard deviations.

Parameters:

Name	Type	Description	Default
n_groups	int	Number of groups.	required

Returns:

Type	Description
ArrayLike	Array of shape (n_groups,) with values >= sd_min.

Source code in pyrenew/latent/hierarchical_priors.py

def sample(self, n_groups: int, **kwargs: object) -> ArrayLike:
    """
    Sample group-level standard deviations.

    Parameters
    ----------
    n_groups
        Number of groups.

    Returns
    -------
    ArrayLike
        Array of shape (n_groups,) with values >= sd_min.
    """
    sd_mean = self.sd_mean_rv()
    concentration = self.sd_concentration_rv()
    rate = concentration / sd_mean

    with numpyro.plate(f"n_{self.name}", n_groups):
        raw_sd = numpyro.sample(
            f"{self.name}_raw",
            dist.Gamma(concentration, rate),
        )

        group_sd = numpyro.deterministic(
            self.name,
            jnp.maximum(raw_sd, self.sd_min),
        )
    return group_sd

HierarchicalNormalPrior

HierarchicalNormalPrior(name: str, sd_rv: RandomVariable)

Bases: RandomVariable

Zero-centered Normal prior for group-level effects.

Samples n_groups values from Normal(0, sd).

Parameters:

Name	Type	Description	Default
name	str	Unique name for the sampled parameter in numpyro.	required
sd_rv	RandomVariable	RandomVariable returning the standard deviation.	required

Initialize hierarchical normal prior.

Parameters:

Name	Type	Description	Default
name	str	Unique name for the sampled parameter.	required
sd_rv	RandomVariable	RandomVariable returning the standard deviation.	required

Source code in pyrenew/latent/hierarchical_priors.py

def __init__(
    self,
    name: str,
    sd_rv: RandomVariable,
) -> None:
    """
    Initialize hierarchical normal prior.

    Parameters
    ----------
    name
        Unique name for the sampled parameter.
    sd_rv
        RandomVariable returning the standard deviation.
    """
    if not isinstance(sd_rv, RandomVariable):
        raise TypeError(
            f"sd_rv must be a RandomVariable, got {type(sd_rv).__name__}. "
            "Use DeterministicVariable(name, value) to wrap a fixed value."
        )

    super().__init__(name=name)
    self.sd_rv = sd_rv

sample

sample(n_groups: int, **kwargs: object) -> ArrayLike

Sample group-level effects.

Parameters:

Name	Type	Description	Default
n_groups	int	Number of groups.	required

Returns:

Type	Description
ArrayLike	Array of shape (n_groups,).

Source code in pyrenew/latent/hierarchical_priors.py

def sample(self, n_groups: int, **kwargs: object) -> ArrayLike:
    """
    Sample group-level effects.

    Parameters
    ----------
    n_groups
        Number of groups.

    Returns
    -------
    ArrayLike
        Array of shape (n_groups,).
    """
    sd = self.sd_rv()

    with numpyro.plate(f"n_{self.name}", n_groups):
        effects = numpyro.sample(
            self.name,
            dist.Normal(0.0, sd),
        )
    return effects

InfectionInitializationMethod

InfectionInitializationMethod(n_timepoints: int)

Method for initializing infections in a renewal process.

Default constructor for pyrenew.latent.infection_initialization_method.InfectionInitializationMethod.

Parameters:

Name	Type	Description	Default
n_timepoints	int	the number of time points for which to generate initial infections	required

Returns:

Type	Description
None

Source code in pyrenew/latent/infection_initialization_method.py

def __init__(self, n_timepoints: int) -> None:
    """Default constructor for
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][].

    Parameters
    ----------
    n_timepoints
        the number of time points for which to
        generate initial infections

    Returns
    -------
    None
    """
    self.validate(n_timepoints)
    self.n_timepoints = n_timepoints

initialize_infections abstractmethod

initialize_infections(I_pre_init: ArrayLike) -> ArrayLike

Generate the number of initialized infections at each time point.

Parameters:

Name	Type	Description	Default
I_pre_init	ArrayLike	An array representing some number of latent infections to be used with the specified [pyrenew.latent.infection_initialization_method.InfectionInitializationMethod][].	required

Returns:

Type	Description
ArrayLike	An array of length n_timepoints with the number of initialized infections at each time point.

Source code in pyrenew/latent/infection_initialization_method.py

@abstractmethod
def initialize_infections(self, I_pre_init: ArrayLike) -> ArrayLike:
    """Generate the number of initialized infections at each time point.

    Parameters
    ----------
    I_pre_init
        An array representing some number of latent infections to be used with the specified `[`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][]`.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of initialized infections at each time point.
    """

validate staticmethod

validate(n_timepoints: int) -> None

Validate inputs to the pyrenew.latent.infection_initialization_method.InfectionInitializationMethod constructor.

Parameters:

Name	Type	Description	Default
n_timepoints	int	the number of time points to generate initial infections for	required

Returns:

Type	Description
None

Source code in pyrenew/latent/infection_initialization_method.py

@staticmethod
def validate(n_timepoints: int) -> None:
    """
    Validate inputs to the
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][]
    constructor.

    Parameters
    ----------
    n_timepoints
        the number of time points to generate initial infections for

    Returns
    -------
    None
    """
    if not isinstance(n_timepoints, int):
        raise TypeError(
            f"n_timepoints must be an integer. Got {type(n_timepoints)}"
        )
    if n_timepoints <= 0:
        raise ValueError(f"n_timepoints must be positive. Got {n_timepoints}")

InfectionInitializationProcess

InfectionInitializationProcess(
    name: str,
    I_pre_init_rv: RandomVariable,
    infection_init_method: InfectionInitializationMethod,
)

Bases: RandomVariable

Generate an initial infection history

Default class constructor for InfectionInitializationProcess

Parameters:

Name	Type	Description	Default
name	str	A name to assign to the RandomVariable.	required
I_pre_init_rv	RandomVariable	A RandomVariable representing the number of infections that occur at some time before the renewal process begins. Each infection_init_method uses this random variable in different ways.	required
infection_init_method	InfectionInitializationMethod	An pyrenew.latent.infection_initialization_method.InfectionInitializationMethod that generates the initial infections for the renewal process.	required

Returns:

Type	Description
None

Source code in pyrenew/latent/infection_initialization_process.py

def __init__(
    self,
    name: str,
    I_pre_init_rv: RandomVariable,
    infection_init_method: InfectionInitializationMethod,
) -> None:
    """Default class constructor for InfectionInitializationProcess

    Parameters
    ----------
    name
        A name to assign to the RandomVariable.
    I_pre_init_rv
        A RandomVariable representing the number of infections that occur at some time before the renewal process begins. Each `infection_init_method` uses this random variable in different ways.
    infection_init_method
        An [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][] that generates the initial infections for the renewal process.

    Returns
    -------
    None
    """
    super().__init__(name=name)
    self.I_pre_init_rv = I_pre_init_rv
    self.infection_init_method = infection_init_method

sample

sample() -> ArrayLike

Sample the Infection Initialization Process.

Returns:

Type	Description
ArrayLike	the number of initialized infections at each time point.

Source code in pyrenew/latent/infection_initialization_process.py

def sample(self) -> ArrayLike:
    """Sample the Infection Initialization Process.

    Returns
    -------
    ArrayLike
        the number of initialized infections at each time point.
    """

    I_pre_init = self.I_pre_init_rv()

    infection_initialization = self.infection_init_method(
        I_pre_init,
    )

    return infection_initialization

InfectionProcessSample

Bases: NamedTuple

Infections and the reproduction number used to compute them.

Attributes:

Name	Type	Description
post_initialization_infections	ArrayLike	The estimated latent infections after the initialization period.
rt	ArrayLike	The reproduction number used in the renewal equation. This is the input Rt for plain renewal and the adjusted Rt when feedback is used.

Infections

Infections(name: str)

Bases: RandomVariable

Latent infections

This class samples infections given \(\mathcal{R}(t)\), initial infections, and generation interval.

Parameters:

Name	Type	Description	Default
name	str	A name for this random variable.	required

Notes

The mathematical model is given by:

\[ I(t) = \mathcal{R}(t) \times \sum_{\tau < t} I(\tau) g(t-\tau) \]

where \(I(t)\) is the number of infections at time \(t\), \(\mathcal{R}(t)\) is the reproduction number at time \(t\), and \(g(t-\tau)\) is the generation interval.

Default constructor.

Parameters:

Name	Type	Description	Default
name	str	A name for this random variable.	required

Source code in pyrenew/latent/infections.py

def __init__(self, name: str) -> None:
    """
    Default constructor.

    Parameters
    ----------
    name
        A name for this random variable.
    """
    super().__init__(name=name)

sample

sample(
    Rt: ArrayLike, I0: ArrayLike, gen_int: ArrayLike, **kwargs: object
) -> InfectionProcessSample

Sample infections given \(\mathcal{R}(t)\), initial infections, and generation interval.

Parameters:

Name	Type	Description	Default
Rt	ArrayLike	Reproduction number.	required
I0	ArrayLike	Initial infections vector of the same length as the generation interval.	required
gen_int	ArrayLike	Generation interval pmf vector.	required
**kwargs	object	Additional keyword arguments passed through to internal sample calls, should there be any.	{}

Returns:

Type	Description
InfectionProcessSample	A named tuple with a post_initialization_infections field.

Source code in pyrenew/latent/infections.py

def sample(
    self,
    Rt: ArrayLike,
    I0: ArrayLike,
    gen_int: ArrayLike,
    **kwargs: object,
) -> InfectionProcessSample:
    r"""
    Sample infections given
    $\mathcal{R}(t)$, initial infections,
    and generation interval.

    Parameters
    ----------
    Rt
        Reproduction number.
    I0
        Initial infections vector
        of the same length as the
        generation interval.
    gen_int
        Generation interval pmf vector.
    **kwargs
        Additional keyword arguments passed through to internal
        sample calls, should there be any.

    Returns
    -------
    InfectionProcessSample
        A named tuple with a
        `post_initialization_infections` field.
    """
    if I0.shape[0] < gen_int.size:
        raise ValueError(
            "Initial infections vector must be at least as long as "
            "the generation interval. "
            f"Initial infections vector length: {I0.shape[0]}, "
            f"generation interval length: {gen_int.size}."
        )

    if I0.shape[1:] != Rt.shape[1:]:
        raise ValueError(
            "Initial infections and Rt must have the "
            "same batch shapes. "
            f"Got initial infections of batch shape {I0.shape[1:]} "
            f"and Rt of batch shape {Rt.shape[1:]}."
        )

    gen_int_rev = jnp.flip(gen_int)
    recent_I0 = I0[-gen_int_rev.size :]

    post_initialization_infections = inf.compute_infections_from_rt(
        I0=recent_I0,
        Rt=Rt,
        reversed_generation_interval_pmf=gen_int_rev,
    )

    return InfectionProcessSample(
        post_initialization_infections=post_initialization_infections,
        rt=Rt,
    )

InfectionsWithFeedback

InfectionsWithFeedback(
    name: str,
    infection_feedback_strength: RandomVariable,
    infection_feedback_pmf: RandomVariable,
)

Bases: RandomVariable

Latent infections

This class computes infections, given Rt, initial infections, and generation interval.

Parameters:

Name	Type	Description	Default
infection_feedback_strength	RandomVariable	Infection feedback strength.	required
infection_feedback_pmf	RandomVariable	Infection feedback pmf.	required

Notes

This function implements the following renewal process (reproduced from pyrenew.latent.infection_functions.compute_infections_from_rt_with_feedback):

\[ \begin{aligned} I(t) & = \mathcal{R}(t)\sum_{\tau=1}^{T_g}I(t - \tau)g(\tau) \\ \mathcal{R}(t) & = \mathcal{R}^u(t)\exp\left(\gamma(t)\ \sum_{\tau=1}^{T_f}I(t - \tau)f(\tau)\right) \end{aligned} \]

where \(\mathcal{R}(t)\) is the reproductive number, \(\gamma(t)\) is the infection feedback strength, \(T_g\) is the max-length of the generation interval, \(\mathcal{R}^u(t)\) is the raw reproduction number, \(f(t)\) is the infection feedback pmf, and \(T_f\) is the max-length of the infection feedback pmf.

Default constructor for InfectionsWithFeedback class.

Parameters:

Name	Type	Description	Default
name	str	A name for this random variable.	required
infection_feedback_strength	RandomVariable	Infection feedback strength.	required
infection_feedback_pmf	RandomVariable	Infection feedback pmf.	required

Returns:

Type	Description
None

Source code in pyrenew/latent/infectionswithfeedback.py

def __init__(
    self,
    name: str,
    infection_feedback_strength: RandomVariable,
    infection_feedback_pmf: RandomVariable,
) -> None:
    """
    Default constructor for InfectionsWithFeedback class.

    Parameters
    ----------
    name
        A name for this random variable.
    infection_feedback_strength
        Infection feedback strength.
    infection_feedback_pmf
        Infection feedback pmf.

    Returns
    -------
    None
    """

    super().__init__(name=name)
    self.validate(infection_feedback_strength, infection_feedback_pmf)

    self.infection_feedback_strength = infection_feedback_strength
    self.infection_feedback_pmf = infection_feedback_pmf

    return None

sample

sample(
    Rt: ArrayLike, I0: ArrayLike, gen_int: ArrayLike, **kwargs: object
) -> InfectionProcessSample

Samples infections given Rt, initial infections, and generation interval.

Parameters:

Name	Type	Description	Default
Rt	ArrayLike	Reproduction number.	required
I0	ArrayLike	Initial infections, as an array at least as long as the generation interval PMF.	required
gen_int	ArrayLike	Generation interval PMF.	required
**kwargs	object	Additional keyword arguments passed through to internal sample calls, should there be any.	{}

Returns:

Type	Description
InfectionProcessSample	Named tuple with post-initialization infections and Rt.

Source code in pyrenew/latent/infectionswithfeedback.py

def sample(
    self,
    Rt: ArrayLike,
    I0: ArrayLike,
    gen_int: ArrayLike,
    **kwargs: object,
) -> InfectionProcessSample:
    """
    Samples infections given Rt, initial infections, and generation
    interval.

    Parameters
    ----------
    Rt
        Reproduction number.
    I0
        Initial infections, as an array
        at least as long as the generation
        interval PMF.
    gen_int
        Generation interval PMF.
    **kwargs
        Additional keyword arguments passed through to internal
        sample calls, should there be any.

    Returns
    -------
    InfectionProcessSample
        Named tuple with post-initialization infections and Rt.
    """
    if I0.shape[0] < gen_int.size:
        raise ValueError(
            "Initial infections must be at least as long as the "
            f"generation interval. Got initial infections length {I0.shape[0]}"
            f"and generation interval length {gen_int.size}."
        )

    if I0.shape[1:] != Rt.shape[1:]:
        raise ValueError(
            "Initial infections and Rt must have the same batch shapes. "
            f"Got initial infections of batch shape {I0.shape[1:]} "
            f"and Rt of batch shape {Rt.shape[1:]}."
        )

    gen_int_rev = jnp.flip(gen_int)

    I0 = I0[-gen_int_rev.size :]

    # Sampling inf feedback strength
    inf_feedback_strength = jnp.atleast_1d(
        self.infection_feedback_strength(
            **kwargs,
        )
    )

    try:
        inf_feedback_strength = jnp.broadcast_to(inf_feedback_strength, Rt.shape)
    except Exception as e:
        raise ValueError(
            "Could not broadcast inf_feedback_strength "
            f"(shape {inf_feedback_strength.shape}) "
            "to the shape of Rt"
            f"{Rt.shape}"
        ) from e

    # Sampling inf feedback pmf
    inf_feedback_pmf = self.infection_feedback_pmf(**kwargs)

    inf_fb_pmf_rev = jnp.flip(inf_feedback_pmf)

    (
        post_initialization_infections,
        Rt_adj,
    ) = inf.compute_infections_from_rt_with_feedback(
        I0=I0,
        Rt_raw=Rt,
        infection_feedback_strength=inf_feedback_strength,
        reversed_generation_interval_pmf=gen_int_rev,
        reversed_infection_feedback_pmf=inf_fb_pmf_rev,
    )

    return InfectionProcessSample(
        post_initialization_infections=post_initialization_infections,
        rt=Rt_adj,
    )

validate staticmethod

validate(inf_feedback_strength: any, inf_feedback_pmf: any) -> None

Validates the input parameters.

Parameters:

Name	Type	Description	Default
inf_feedback_strength	any	Infection feedback strength.	required
inf_feedback_pmf	any	Infection feedback pmf.	required

Returns:

Type	Description
None

Source code in pyrenew/latent/infectionswithfeedback.py

@staticmethod
def validate(
    inf_feedback_strength: any,
    inf_feedback_pmf: any,
) -> None:  # numpydoc ignore=GL08
    """
    Validates the input parameters.

    Parameters
    ----------
    inf_feedback_strength
        Infection feedback strength.
    inf_feedback_pmf
        Infection feedback pmf.

    Returns
    -------
    None
    """
    assert isinstance(inf_feedback_strength, RandomVariable)
    assert isinstance(inf_feedback_pmf, RandomVariable)

    return None

InitializeInfectionsExponentialGrowth

InitializeInfectionsExponentialGrowth(
    n_timepoints: int, rate_rv: RandomVariable, t_pre_init: int | None = None
)

Bases: InfectionInitializationMethod

Generate initial infections according to exponential growth.

Notes

The number of incident infections at time t is given by:

\[ I(t) = I_p \exp \left( r (t - t_p) \right) \]

Where \(I_p\) is I_pre_init, \(r\) is rate, and \(t_p\) is t_pre_init. This ensures that \(I(t_p) = I_p\). We default to t_pre_init = n_timepoints - 1, so that I_pre_init represents the number of incident infections immediately before the renewal process begins.

Default constructor for the pyrenew.latent.infection_initialization_method.InitializeInfectionsExponentialGrowth class.

Parameters:

Name	Type	Description	Default
n_timepoints	int	the number of time points to generate initial infections for	required
rate_rv	RandomVariable	A random variable representing the rate of exponential growth	required
t_pre_init	int | None	The time point whose number of infections is described by I_pre_init. Defaults to n_timepoints - 1.	None

Source code in pyrenew/latent/infection_initialization_method.py

def __init__(
    self,
    n_timepoints: int,
    rate_rv: RandomVariable,
    t_pre_init: int | None = None,
) -> None:
    """Default constructor for the [`pyrenew.latent.infection_initialization_method.InitializeInfectionsExponentialGrowth`][] class.

    Parameters
    ----------
    n_timepoints
        the number of time points to generate initial infections for
    rate_rv
        A random variable representing the rate of exponential growth
    t_pre_init
         The time point whose number of infections is described by ``I_pre_init``. Defaults to ``n_timepoints - 1``.
    """
    super().__init__(n_timepoints)
    self.rate_rv = rate_rv
    if t_pre_init is None:
        t_pre_init = n_timepoints - 1
    self.t_pre_init = t_pre_init

initialize_infections

initialize_infections(I_pre_init: ArrayLike) -> ArrayLike

Generate initial infections according to exponential growth.

Parameters:

Name	Type	Description	Default
I_pre_init	ArrayLike	An array of size 1 representing the number of infections at time t_pre_init.	required

Returns:

Type	Description
ArrayLike	An array of length n_timepoints with the number of initialized infections at each time point.

Source code in pyrenew/latent/infection_initialization_method.py

def initialize_infections(self, I_pre_init: ArrayLike) -> ArrayLike:
    """Generate initial infections according to exponential growth.

    Parameters
    ----------
    I_pre_init
        An array of size 1 representing the number of infections at time ``t_pre_init``.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of initialized infections at each time point.
    """
    I_pre_init = jnp.array(I_pre_init)
    rate = jnp.array(self.rate_rv())
    initial_infections = I_pre_init * jnp.exp(
        rate * (jnp.arange(self.n_timepoints)[:, jnp.newaxis] - self.t_pre_init)
    )
    return jnp.squeeze(initial_infections)

InitializeInfectionsFromVec

InitializeInfectionsFromVec(n_timepoints: int)

Bases: InfectionInitializationMethod

Create initial infections from a vector of infections.

Source code in pyrenew/latent/infection_initialization_method.py

def __init__(self, n_timepoints: int) -> None:
    """Default constructor for
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][].

    Parameters
    ----------
    n_timepoints
        the number of time points for which to
        generate initial infections

    Returns
    -------
    None
    """
    self.validate(n_timepoints)
    self.n_timepoints = n_timepoints

initialize_infections

initialize_infections(I_pre_init: ArrayLike) -> ArrayLike

Create initial infections from a vector of infections.

Parameters:

Name	Type	Description	Default
I_pre_init	ArrayLike	An array with the same length as n_timepoints to be used as the initial infections.	required

Returns:

Type	Description
ArrayLike	An array of length n_timepoints with the number of initialized infections at each time point.

Source code in pyrenew/latent/infection_initialization_method.py

def initialize_infections(self, I_pre_init: ArrayLike) -> ArrayLike:
    """Create initial infections from a vector of infections.

    Parameters
    ----------
    I_pre_init
        An array with the same length as ``n_timepoints`` to be
        used as the initial infections.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of
        initialized infections at each time point.
    """
    I_pre_init = jnp.array(I_pre_init)
    if I_pre_init.size != self.n_timepoints:
        raise ValueError(
            "I_pre_init must have the same size as n_timepoints. "
            f"Got I_pre_init of size {I_pre_init.size} "
            f"and n_timepoints of size {self.n_timepoints}."
        )
    return I_pre_init

InitializeInfectionsZeroPad

InitializeInfectionsZeroPad(n_timepoints: int)

Bases: InfectionInitializationMethod

Create an initial infection vector of specified length by padding a shorter vector with an appropriate number of zeros at the beginning of the time series.

Source code in pyrenew/latent/infection_initialization_method.py

def __init__(self, n_timepoints: int) -> None:
    """Default constructor for
    [`pyrenew.latent.infection_initialization_method.InfectionInitializationMethod`][].

    Parameters
    ----------
    n_timepoints
        the number of time points for which to
        generate initial infections

    Returns
    -------
    None
    """
    self.validate(n_timepoints)
    self.n_timepoints = n_timepoints

initialize_infections

initialize_infections(I_pre_init: ArrayLike) -> ArrayLike

Pad the initial infections with zeros at the beginning of the time series.

Parameters:

Name	Type	Description	Default
I_pre_init	ArrayLike	An array with initialized infections to be padded with zeros.	required

Returns:

Type	Description
ArrayLike	An array of length n_timepoints with the number of initialized infections at each time point.

Source code in pyrenew/latent/infection_initialization_method.py

def initialize_infections(self, I_pre_init: ArrayLike) -> ArrayLike:
    """Pad the initial infections with zeros at the beginning of the time series.

    Parameters
    ----------
    I_pre_init
        An array with initialized infections to be padded with zeros.

    Returns
    -------
    ArrayLike
        An array of length ``n_timepoints`` with the number of initialized infections at each time point.
    """
    I_pre_init = jnp.atleast_1d(I_pre_init)
    if self.n_timepoints < I_pre_init.size:
        raise ValueError(
            "I_pre_init must be no longer than n_timepoints. "
            f"Got I_pre_init of size {I_pre_init.size} and "
            f" n_timepoints of size {self.n_timepoints}."
        )
    return jnp.pad(I_pre_init, (self.n_timepoints - I_pre_init.size, 0))

LatentSample

Bases: NamedTuple

Output from latent infection process sampling.

Attributes:

Name	Type	Description
aggregate	ArrayLike	Total infections aggregated across all subpopulations. Shape: (n_total_days,)
all_subpops	ArrayLike	Infections for all subpopulations. Shape: (n_total_days, n_subpops)

MobilitySubpopulationInfections

MobilitySubpopulationInfections(
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    baseline_rt_process: TemporalProcess,
    subpop_rt_deviation_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
    mobility_matrix_rv: DeterministicVariable,
    susceptible_depletion: bool = True,
)

Bases: BaseSubpopulationInfections

Subpopulation renewal model coupled by a fixed mobility matrix.

Subpopulations are home regions whose residents spend part of each day at other locations. Transmission happens at each location, at that location's reproduction number, among the people present there, and each infection is assigned to the home region of the person infected. This implements Equation 2 of Roy, Clapham, and Mishra (2025), Incorporating human mobility to enhance epidemic response and estimate real-time reproduction numbers. With depletion enabled, PyRenew additionally caps each day's new infections at the remaining susceptible population.

Infections are per resident: infections[t, k] is new infections among residents of home region k on day t, divided by the number of residents of k. With M the mobility matrix and p = subpop_fractions, one day with depletion enabled is:

present = M @ p
share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None]
infectiousness = reversed_gen_int @ history
susceptible = 1.0 - cumulative
new_infections_raw = M.T @ (rt * (share_present @ infectiousness))
new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible)

share_present[j, k] is the share of the people present at location j who live in region k, and cumulative counts every infection since the first day of the initialization period. Initial infections must be finite and non-negative. With depletion enabled, their sum over all initialization days must be at most 1 per resident in each region. Concrete invalid values raise ValueError; during JAX tracing, callers must ensure sampled initialization meets these bounds.

M[j, k] is the fraction of residents of home region k present at location j: columns are home regions, rows are locations, and each column sums to 1. Regions must be in the same order in the matrix, subpop_fractions, a vector-valued I0_rv, a vector-valued log_rt_time_0_rv, and every observation's subpop_indices.

The recorded values are those of pyrenew.latent.SubpopulationInfections, plus susceptible_fraction. Under mobility:

• rt_subpop[t, j] is the reproduction number at location j, where transmission happens, when everyone is susceptible.
• rt_subpop_adjusted is rt_subpop * (share_present @ susceptible_fraction) after the initialization period, with depletion enabled. This adjusts for susceptibility but does not account for the incidence ceiling. It equals new infections per unit of infectiousness present only when the ceiling does not bind. It is 0 at a location with no one present.
• susceptible_fraction[t, i] is the share of residents of home region i not infected before day t. It is recorded as 1 minus cumulative infections and can become negative with depletion off.

Susceptible depletion is on by default, as in Eq 2. With it off, neither susceptibility scaling nor the incidence ceiling is applied, and inferred raw values of rt_subpop will need to directly absorb any reductions in transmission due to susceptible depletion.

The matrix is fixed. Fitting and prediction use the matrix the model was built with; it is available as mobility_matrix_rv.value.

Initialize the mobility-coupled infections process.

Parameters:

Name	Type	Description	Default
name	str	Name prefix for numpyro sample sites. All deterministic quantities are recorded under this scope (e.g., "{name}::rt_subpop", "{name}::rt_subpop_adjusted", and "{name}::susceptible_fraction").	required
gen_int_rv	RandomVariable	Generation interval PMF	required
n_initialization_points	int	Number of initialization days before day 0.	required
I0_rv	RandomVariable	New infections per resident on the first day of the initialization period, the earliest day on the time axis. This is not prevalence. Returns either a scalar, one value for all regions, or an array of shape (n_subpops,), one value per region in the order of subpop_fractions. Every initialization day counts toward susceptible depletion. With depletion enabled, the sum over all initialization days must be at most 1 per resident in each region.	required
baseline_rt_process	TemporalProcess	Temporal process for baseline \(\mathcal{R}(t)\) dynamics	required
subpop_rt_deviation_process	TemporalProcess	Temporal process for subpopulation deviations	required
log_rt_time_0_rv	RandomVariable	Starting log(\(\mathcal{R}(t)\)) on the first day of the initialization period. Returns either a scalar, one value for all regions, or an array of shape (n_subpops,), one value per region in the order of subpop_fractions, which gives each region its own starting \(\mathcal{R}(t)\), as in Roy et al.	required
mobility_matrix_rv	DeterministicVariable	The mobility matrix used on every day, as a DeterministicVariable whose value has shape (n_subpops, n_subpops). mobility_matrix[j, k] is the fraction of residents of home region k present at location j; each column sums to 1. Rows and columns are in the order of subpop_fractions. Convert a data frame with .to_numpy(). Transpose a matrix whose rows sum to 1.	required
susceptible_depletion	bool	If True (default), new infections among residents of each region are scaled by the fraction of those residents not yet infected, as in Roy et al. Eq 2, and capped at that fraction. If False, neither scaling nor the ceiling is applied, and rt_subpop absorbs depletion.	True

Raises:

Type	Description
TypeError	If mobility_matrix_rv is not a DeterministicVariable.
ValueError	If required parameters are missing or invalid, or if the mobility matrix is not a valid column-stochastic matrix.

Source code in pyrenew/latent/mobility_subpopulation_infections.py

def __init__(
    self,
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    baseline_rt_process: TemporalProcess,
    subpop_rt_deviation_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
    mobility_matrix_rv: DeterministicVariable,
    susceptible_depletion: bool = True,
) -> None:
    """
    Initialize the mobility-coupled infections process.

    Parameters
    ----------
    name
        Name prefix for numpyro sample sites. All deterministic
        quantities are recorded under this scope (e.g.,
        ``"{name}::rt_subpop"``, ``"{name}::rt_subpop_adjusted"``, and
        ``"{name}::susceptible_fraction"``).
    gen_int_rv
        Generation interval PMF
    n_initialization_points
        Number of initialization days before day 0.
    I0_rv
        New infections per resident on the first day of the
        initialization period, the earliest day on the time axis. This
        is not prevalence. Returns either a scalar, one value for all
        regions, or an array of shape ``(n_subpops,)``, one value per
        region in the order of ``subpop_fractions``. Every
        initialization day counts toward susceptible depletion. With
        depletion enabled, the sum over all initialization days must
        be at most 1 per resident in each region.
    baseline_rt_process
        Temporal process for baseline $\\mathcal{R}(t)$ dynamics
    subpop_rt_deviation_process
        Temporal process for subpopulation deviations
    log_rt_time_0_rv
        Starting log($\\mathcal{R}(t)$) on the first day of the
        initialization period. Returns either a scalar, one value for
        all regions, or an array of shape ``(n_subpops,)``, one value per
        region in the order of ``subpop_fractions``, which gives each
        region its own starting $\\mathcal{R}(t)$, as in Roy et al.
    mobility_matrix_rv
        The mobility matrix used on every day, as a
        ``DeterministicVariable`` whose value has shape
        ``(n_subpops, n_subpops)``. ``mobility_matrix[j, k]`` is the
        fraction of residents of home region ``k`` present at location
        ``j``; each column sums to 1. Rows and columns are in the order
        of ``subpop_fractions``. Convert a data frame with
        ``.to_numpy()``. Transpose a matrix whose rows sum to 1.
    susceptible_depletion
        If True (default), new infections among residents of each region
        are scaled by the fraction of those residents not yet infected,
        as in Roy et al. Eq 2, and capped at that fraction. If False,
        neither scaling nor the ceiling is applied, and ``rt_subpop``
        absorbs depletion.

    Raises
    ------
    TypeError
        If mobility_matrix_rv is not a DeterministicVariable.
    ValueError
        If required parameters are missing or invalid, or if the
        mobility matrix is not a valid column-stochastic matrix.
    """
    super().__init__(
        name=name,
        gen_int_rv=gen_int_rv,
        n_initialization_points=n_initialization_points,
        I0_rv=I0_rv,
        baseline_rt_process=baseline_rt_process,
        subpop_rt_deviation_process=subpop_rt_deviation_process,
        log_rt_time_0_rv=log_rt_time_0_rv,
    )
    if not isinstance(mobility_matrix_rv, DeterministicVariable):
        raise TypeError(
            "mobility_matrix_rv must be a DeterministicVariable; the "
            "mobility matrix is data and cannot be random. Got "
            f"{type(mobility_matrix_rv).__name__}."
        )
    self._validate_mobility_matrix(jnp.asarray(mobility_matrix_rv.value))
    self.mobility_matrix_rv = mobility_matrix_rv
    self.susceptible_depletion = susceptible_depletion

PopulationInfections

PopulationInfections(
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    single_rt_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
    infection_feedback_strength_rv: RandomVariable | None = None,
    infection_feedback_pmf_rv: RandomVariable | None = None,
)

Bases: BaseLatentInfectionProcess

A single \(\mathcal{R}(t)\) trajectory drives one renewal equation.

The constructor specifies model structure (priors, temporal processes). rt_single is the raw reproduction number generated by single_rt_process. rt_single_adjusted is the reproduction number used to generate infections. The two are equal unless infection feedback is on.

Infection feedback is optional. It is on when both infection_feedback_strength_rv and infection_feedback_pmf_rv are given. On each day after the initialization period, it multiplies the raw \(\mathcal{R}(t)\) by a factor that depends on recent infections:

\[ \mathcal{R}^{\mathrm{adj}}(t) = \mathcal{R}(t) \exp\left(\gamma(t) \sum_{\tau=1}^{T_g} I(t - \tau) f(\tau)\right) \]

where \(I(t)\) is new infections per resident, \(\gamma(t)\) is the feedback strength, \(f\) is the feedback PMF, and \(T_g\) is the length of the generation interval. The renewal equation then uses the adjusted \(\mathcal{R}^{\mathrm{adj}}(t)\) in place of \(\mathcal{R}(t)\). Negative strength lowers \(\mathcal{R}(t)\) when recent infections are high, and positive strength raises it. Because \(I(t)\) is per resident, the effect of a given strength depends on the level of infections per resident. With both arguments left out, infections follow the plain renewal equation.

Initialize population-level infections process.

Parameters:

Name	Type	Description	Default
name	str	Name prefix for numpyro sample sites. All deterministic quantities are recorded under this scope (e.g., "{name}::rt_single" and "{name}::rt_single_adjusted").	required
gen_int_rv	RandomVariable	Generation interval PMF	required
n_initialization_points	int	Number of initialization days before day 0.	required
I0_rv	RandomVariable	Initial infection prevalence (proportion of population)	required
single_rt_process	TemporalProcess	Temporal process for single \(\mathcal{R}(t)\) dynamics	required
log_rt_time_0_rv	RandomVariable	Initial value for log(\(\mathcal{R}(t)\)) at time 0.	required
infection_feedback_strength_rv	RandomVariable | None	Infection feedback strength \(\gamma\). Returns a scalar, or an array of shape (n_days_post_init,) for a strength that varies by day. Give together with infection_feedback_pmf_rv, or leave both out for plain renewal.	None
infection_feedback_pmf_rv	RandomVariable | None	Feedback PMF \(f\), which weights recent infections in the feedback factor. Element 0 weights infections on the previous day. Must have the same length as the generation interval. Give together with infection_feedback_strength_rv.	None

Raises:

Type	Description
ValueError	If required parameters are missing or invalid, or if only one of the two infection feedback arguments is given.

Source code in pyrenew/latent/population_infections.py

def __init__(
    self,
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    single_rt_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
    infection_feedback_strength_rv: RandomVariable | None = None,
    infection_feedback_pmf_rv: RandomVariable | None = None,
) -> None:
    """
    Initialize population-level infections process.

    Parameters
    ----------
    name
        Name prefix for numpyro sample sites. All deterministic
        quantities are recorded under this scope (e.g.,
        ``"{name}::rt_single"`` and
        ``"{name}::rt_single_adjusted"``).
    gen_int_rv
        Generation interval PMF
    n_initialization_points
        Number of initialization days before day 0.
    I0_rv
        Initial infection prevalence (proportion of population)
    single_rt_process
        Temporal process for single $\\mathcal{R}(t)$ dynamics
    log_rt_time_0_rv
        Initial value for log($\\mathcal{R}(t)$) at time 0.
    infection_feedback_strength_rv
        Infection feedback strength $\\gamma$. Returns a scalar, or an
        array of shape ``(n_days_post_init,)`` for a strength that varies
        by day. Give together with ``infection_feedback_pmf_rv``, or leave
        both out for plain renewal.
    infection_feedback_pmf_rv
        Feedback PMF $f$, which weights recent infections in the feedback
        factor. Element 0 weights infections on the previous day. Must have
        the same length as the generation interval. Give together with
        ``infection_feedback_strength_rv``.

    Raises
    ------
    ValueError
        If required parameters are missing or invalid, or if only one of
        the two infection feedback arguments is given.
    """
    super().__init__(
        name=name,
        gen_int_rv=gen_int_rv,
        n_initialization_points=n_initialization_points,
    )

    if I0_rv is None:
        raise ValueError("I0_rv is required")
    self.I0_rv = I0_rv

    if isinstance(I0_rv, DeterministicVariable):
        self._validate_I0(I0_rv.value)

    if log_rt_time_0_rv is None:
        raise ValueError("log_rt_time_0_rv is required")
    self.log_rt_time_0_rv = log_rt_time_0_rv

    if single_rt_process is None:
        raise ValueError("single_rt_process is required")
    self.single_rt_process = single_rt_process

    _validate_feedback_options(
        infection_feedback_strength_rv, infection_feedback_pmf_rv
    )
    self.infection_feedback_strength_rv = infection_feedback_strength_rv
    self.infection_feedback_pmf_rv = infection_feedback_pmf_rv

default_subpop_fractions

default_subpop_fractions() -> ArrayLike

Return default population fractions for a single-population model.

Returns:

Type	Description
ArrayLike	jnp.array([1.0])

Source code in pyrenew/latent/population_infections.py

def default_subpop_fractions(self) -> ArrayLike:
    """
    Return default population fractions for a single-population model.

    Returns
    -------
    ArrayLike
        ``jnp.array([1.0])``
    """
    return jnp.array([1.0])

sample

sample(
    n_days_post_init: int,
    subpop_fractions: ArrayLike | None = None,
    first_day_dow: int | None = None,
    **kwargs: object,
) -> LatentSample

Sample population infections using a single renewal process.

Generates a single daily \(\mathcal{R}(t)\) trajectory, computes initial infections via exponential backprojection, and computes infections with the renewal equation, with infection feedback if it is on. The temporal process may sample parameters at any supported cadence (for example daily or weekly/stepwise), but it must return a daily-length trajectory before the renewal equation is evaluated. The raw Rt is recorded as rt_single; the adjusted Rt used to generate infections is recorded as rt_single_adjusted.

Parameters:

Name	Type	Description	Default
n_days_post_init	int	Number of days to simulate after initialization period	required
subpop_fractions	ArrayLike | None	Population fractions. Defaults to [1.0] (single population). Must be [1.0] if provided.	None
first_day_dow	int | None	Day of week for element 0 of the full latent infection time axis, including initialization days. When this latent process is used inside MultiSignalModel, the caller passes obs_start_date to the model and it converts the first observation day to this axis-origin day of week by subtracting n_initialization_points. Forwarded to single_rt_process. See pyrenew.latent.TemporalProcess.	None
**kwargs	object	Additional keyword arguments (unused).	{}

Returns:

Type	Description
LatentSample	Named tuple with fields: - aggregate: shape (n_total_days,) - all_subpops: shape (n_total_days, 1)

Raises:

Type	Description
ValueError	If subpop_fractions does not represent a single population with fraction [1.0] or if I0_rv does not return a scalar, or if the infection feedback strength does not broadcast to shape (n_days_post_init,), or the feedback PMF length differs from the generation interval length.

Source code in pyrenew/latent/population_infections.py

def sample(
    self,
    n_days_post_init: int,
    subpop_fractions: ArrayLike | None = None,
    first_day_dow: int | None = None,
    **kwargs: object,
) -> LatentSample:
    """
    Sample population infections using a single renewal process.

    Generates a single daily $\\mathcal{R}(t)$ trajectory, computes initial
    infections via exponential backprojection, and computes infections
    with the renewal equation, with infection feedback if it is on. The
    temporal process may sample parameters at any supported cadence (for
    example daily or weekly/stepwise), but it must return a daily-length
    trajectory before the renewal equation is evaluated. The raw Rt is
    recorded as ``rt_single``; the adjusted Rt used to generate infections
    is recorded as ``rt_single_adjusted``.

    Parameters
    ----------
    n_days_post_init
        Number of days to simulate after initialization period
    subpop_fractions
        Population fractions. Defaults to ``[1.0]`` (single population).
        Must be ``[1.0]`` if provided.
    first_day_dow
        Day of week for element 0 of the full latent infection time axis,
        including initialization days. When this latent process is used
        inside MultiSignalModel, the caller passes obs_start_date to the
        model and it converts the first observation day to this axis-origin
        day of week by subtracting n_initialization_points. Forwarded to
        ``single_rt_process``. See [pyrenew.latent.TemporalProcess][].
    **kwargs
        Additional keyword arguments (unused).

    Returns
    -------
    LatentSample
        Named tuple with fields:
        - aggregate: shape (n_total_days,)
        - all_subpops: shape (n_total_days, 1)

    Raises
    ------
    ValueError
        If ``subpop_fractions`` does not represent a single population
        with fraction ``[1.0]`` or if ``I0_rv`` does not return a scalar,
        or if the infection feedback strength does not broadcast to shape
        ``(n_days_post_init,)``, or the feedback PMF length differs from
        the generation interval length.
    """
    pop = self._parse_and_validate_fractions(
        subpop_fractions=subpop_fractions,
    )
    frac_check = jnp.isclose(pop.fractions[0], 1.0, atol=1e-6)
    if pop.n_subpops != 1 or (not_jax_tracer(frac_check) and not frac_check):
        raise ValueError(
            "PopulationInfections requires exactly one subpopulation "
            "with fraction [1.0]"
        )

    n_total_days = self.n_initialization_points + n_days_post_init

    initial_log_rt = self.log_rt_time_0_rv()

    log_rt_single = self.single_rt_process.sample(
        n_timepoints=n_total_days,
        initial_value=initial_log_rt,
        name_prefix="log_rt_single",
        first_day_dow=first_day_dow,
    )
    require_shape(log_rt_single, (n_total_days, 1), "single_rt_process")

    rt_single = jnp.exp(log_rt_single)

    gen_int = self.gen_int_rv()

    I0 = self._validate_and_prepare_I0(jnp.asarray(self.I0_rv()), pop)

    initial_r = r_approx_from_R(
        R=rt_single[0, 0],
        g=gen_int,
        n_newton_steps=4,
    )

    time_indices = jnp.arange(self.n_initialization_points)
    I0_init = I0 * jnp.exp(initial_r * time_indices)

    post_init_infections, rt_adjusted = _compute_renewal(
        I0=I0_init,
        rt_raw=rt_single[self.n_initialization_points :, 0],
        gen_int=gen_int,
        infection_feedback_strength_rv=self.infection_feedback_strength_rv,
        infection_feedback_pmf_rv=self.infection_feedback_pmf_rv,
    )
    rt_single_adjusted = jnp.concatenate(
        [rt_single[: self.n_initialization_points, 0], rt_adjusted]
    )[:, jnp.newaxis]

    infections_1d = jnp.concatenate([I0_init, post_init_infections])
    infections_all = infections_1d[:, jnp.newaxis]
    infections_aggregate = infections_1d

    self._validate_output_shapes(
        infections_aggregate,
        infections_all,
        n_total_days,
        pop,
    )

    with numpyro.handlers.scope(prefix=self.name, divider="::"):
        numpyro.deterministic("I0_init", I0_init)
        numpyro.deterministic("log_rt_single", log_rt_single)
        numpyro.deterministic("rt_single", rt_single)
        numpyro.deterministic("rt_single_adjusted", rt_single_adjusted)
        numpyro.deterministic("infections_aggregate", infections_aggregate)

    return LatentSample(
        aggregate=infections_aggregate,
        all_subpops=infections_all,
    )

validate

validate() -> None

Validate population infections parameters.

Checks that the generation interval is a valid PMF.

Raises:

Type	Description
ValueError	If gen_int_rv does not return a valid discrete distribution

Source code in pyrenew/latent/population_infections.py

def validate(self) -> None:
    """
    Validate population infections parameters.

    Checks that the generation interval is a valid PMF.

    Raises
    ------
    ValueError
        If gen_int_rv does not return a valid discrete distribution
    """
    validate_discrete_dist_vector(self.gen_int_rv())

PopulationStructure dataclass

PopulationStructure(fractions: ArrayLike)

Parsed and validated population structure for a jurisdiction.

Attributes:

Name	Type	Description
fractions	ArrayLike	Population fractions for all subpopulations. Shape: (n_subpops,)

n_subpops property

n_subpops: int

Total number of subpopulations.

Returns:

Type	Description
int	The number of subpopulations.

RandomWalk

RandomWalk(
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
)

Bases: TemporalProcess

Random walk process for log(Rt).

Each value equals the previous value plus noise, with no reversion toward a mean. Allows Rt to drift without bound — suitable when you have no prior expectation that Rt will return to a baseline.

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:

Name	Type	Description	Default
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of noise at each time step. Larger values produce faster drift; smaller values produce more gradual changes.	required
parameterization	Parameterization	Which latent object to sample: "innovation" (default) or "state".	'innovation'

Notes

Unlike AR(1), variance grows over time — the process can wander arbitrarily far from its starting point. For long time horizons, consider AR(1) if you want Rt to stay bounded near a baseline.

Initialize random walk process.

Parameters:

Name	Type	Description	Default
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of innovations.	required
parameterization	Parameterization	"innovation" (default) or "state". See class docstring.	'innovation'

Raises:

Type	Description
TypeError	If innovation_sd_rv is not a RandomVariable instance
ValueError	If innovation_sd_rv is a DeterministicVariable with any value <= 0, or if parameterization is not a recognized string

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
) -> None:
    """
    Initialize random walk process.

    Parameters
    ----------
    innovation_sd_rv
        RandomVariable that returns the standard deviation of innovations.
    parameterization
        ``"innovation"`` (default) or ``"state"``. See class docstring.

    Raises
    ------
    TypeError
        If innovation_sd_rv is not a RandomVariable instance
    ValueError
        If innovation_sd_rv is a DeterministicVariable with any value <= 0,
        or if parameterization is not a recognized string
    """
    if not isinstance(innovation_sd_rv, RandomVariable):
        raise TypeError("innovation_sd_rv must be a RandomVariable")
    _validate_deterministic_innovation_sd(innovation_sd_rv)
    _validate_parameterization(parameterization)
    self.innovation_sd_rv = innovation_sd_rv
    self.parameterization = parameterization

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"RandomWalk(innovation_sd_rv={self.innovation_sd_rv}, "
        f"parameterization={self.parameterization!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "rw",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample random walk trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Deterministic first state of the trajectory. Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites	'rw'
first_day_dow	int | None	Unused. See pyrenew.latent.TemporalProcess.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "rw",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample random walk trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Deterministic first state of the trajectory. Scalar values are
        broadcast to all processes; arrays must have shape
        ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites
    first_day_dow
        Unused. See [pyrenew.latent.TemporalProcess][].

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    initial_value = _prepare_initial_value(initial_value, n_processes)

    innovation_sd = jnp.asarray(self.innovation_sd_rv())

    if self.parameterization == "innovation":
        rw = ProcessRandomWalk(
            name=f"{name_prefix}_random_walk",
            step_rv=DistributionalVariable(
                name=f"{name_prefix}_step",
                distribution=dist.Normal(
                    jnp.zeros(n_processes),
                    innovation_sd,
                ),
            ),
        )

        return rw.sample(
            init_vals=initial_value[jnp.newaxis, :],
            n=n_timepoints,
        )

    if n_timepoints == 1:
        return initial_value[jnp.newaxis, :]

    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,
    )
    return jnp.concatenate([initial_value[jnp.newaxis, :], post_init], axis=0)

StepwiseTemporalProcess

StepwiseTemporalProcess(inner: TemporalProcess, step_size: int)

Bases: TemporalProcess

Parameterize an inner temporal process at a coarser cadence and broadcast to the per-timepoint scale by model-index repetition.

Each step_size consecutive output timepoints share a single sampled value from the inner process. Use when a parameter should be estimated at a coarser cadence while downstream model components still need one value per evaluation timepoint. Blocks always start at output index 0.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the coarse-scale trajectory. Must satisfy the TemporalProcess protocol.	required
step_size	int	Number of per-timepoint units that share each inner sample. Must be a positive integer.	required

Raises:

Type	Description
ValueError	If step_size is not a positive integer.

Initialize stepwise temporal process.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the coarse trajectory.	required
step_size	int	Number of per-timepoint units that share each inner sample.	required

Raises:

Type	Description
ValueError	If step_size is not a positive integer.

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    inner: TemporalProcess,
    step_size: int,
) -> None:
    """
    Initialize stepwise temporal process.

    Parameters
    ----------
    inner
        Inner ``TemporalProcess`` that generates the coarse trajectory.
    step_size
        Number of per-timepoint units that share each inner sample.

    Raises
    ------
    ValueError
        If ``step_size`` is not a positive integer.
    """
    if not isinstance(step_size, int) or step_size < 1:
        raise ValueError(f"step_size must be a positive integer, got {step_size!r}")
    self.inner = inner
    self.step_size = step_size

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"StepwiseTemporalProcess(inner={self.inner!r}, step_size={self.step_size})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "stepwise",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample coarse trajectory from inner process and broadcast.

Computes the number of coarse time steps needed for the requested length, samples the inner process at that cadence, then broadcasts each coarse value to the per-timepoint axis and trims to n_timepoints. The returned value always has one row per evaluation timepoint, regardless of the inner parameter cadence. The coarse trajectory is recorded as a NumPyro deterministic site named "{name_prefix}_coarse".

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of per-timepoint outputs to produce.	required
initial_value	float | ArrayLike | None	Initial value(s) for the inner process. Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites; forwarded to the inner process.	'stepwise'
first_day_dow	int | None	Ignored. Accepted for compatibility with the TemporalProcess protocol.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes), constant within each block of step_size consecutive rows.

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "stepwise",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample coarse trajectory from inner process and broadcast.

    Computes the number of coarse time steps needed for the requested
    length, samples the inner process at that cadence, then broadcasts each
    coarse value to the per-timepoint axis and trims to ``n_timepoints``.
    The returned value always has one row per evaluation timepoint,
    regardless of the inner parameter cadence. The coarse
    trajectory is recorded as a NumPyro deterministic site named
    ``"{name_prefix}_coarse"``.

    Parameters
    ----------
    n_timepoints
        Number of per-timepoint outputs to produce.
    initial_value
        Initial value(s) for the inner process. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites; forwarded to the inner process.
    first_day_dow
        Ignored. Accepted for compatibility with the
        ``TemporalProcess`` protocol.

    Returns
    -------
    ArrayLike
        Trajectories of shape ``(n_timepoints, n_processes)``, constant
        within each block of ``step_size`` consecutive rows.
    """
    n_steps = self._resolve_n_coarse(n_timepoints)
    coarse = self.inner.sample(
        n_timepoints=n_steps,
        initial_value=initial_value,
        n_processes=n_processes,
        name_prefix=name_prefix,
    )
    numpyro.deterministic(f"{name_prefix}_coarse", coarse)
    return jnp.repeat(coarse, repeats=self.step_size, axis=0)[:n_timepoints]

StudentTGroupModePrior

StudentTGroupModePrior(name: str, sd_rv: RandomVariable, df_rv: RandomVariable)

Bases: RandomVariable

Zero-centered Student-t prior for group-level modes (robust alternative to Normal).

Samples n_groups values from StudentT(df, 0, sd).

Parameters:

Name	Type	Description	Default
name	str	Unique name for the sampled parameter in numpyro.	required
sd_rv	RandomVariable	RandomVariable returning the scale parameter.	required
df_rv	RandomVariable	RandomVariable returning the degrees of freedom.	required

Initialize Student-t group mode prior.

Parameters:

Name	Type	Description	Default
name	str	Unique name for the sampled parameter.	required
sd_rv	RandomVariable	RandomVariable returning the scale parameter.	required
df_rv	RandomVariable	RandomVariable returning the degrees of freedom.	required

Source code in pyrenew/latent/hierarchical_priors.py

def __init__(
    self,
    name: str,
    sd_rv: RandomVariable,
    df_rv: RandomVariable,
) -> None:
    """
    Initialize Student-t group mode prior.

    Parameters
    ----------
    name
        Unique name for the sampled parameter.
    sd_rv
        RandomVariable returning the scale parameter.
    df_rv
        RandomVariable returning the degrees of freedom.
    """
    if not isinstance(sd_rv, RandomVariable):
        raise TypeError(
            f"sd_rv must be a RandomVariable, got {type(sd_rv).__name__}. "
            "Use DeterministicVariable(name, value) to wrap a fixed value."
        )
    if not isinstance(df_rv, RandomVariable):
        raise TypeError(
            f"df_rv must be a RandomVariable, got {type(df_rv).__name__}. "
            "Use DeterministicVariable(name, value) to wrap a fixed value."
        )

    super().__init__(name=name)
    self.sd_rv = sd_rv
    self.df_rv = df_rv

sample

sample(n_groups: int, **kwargs: object) -> ArrayLike

Sample group-level modes.

Parameters:

Name	Type	Description	Default
n_groups	int	Number of groups.	required

Returns:

Type	Description
ArrayLike	Array of shape (n_groups,).

Source code in pyrenew/latent/hierarchical_priors.py

def sample(self, n_groups: int, **kwargs: object) -> ArrayLike:
    """
    Sample group-level modes.

    Parameters
    ----------
    n_groups
        Number of groups.

    Returns
    -------
    ArrayLike
        Array of shape (n_groups,).
    """
    sd = self.sd_rv()
    df = self.df_rv()

    with numpyro.plate(f"n_{self.name}", n_groups):
        effects = numpyro.sample(
            self.name,
            dist.StudentT(df=df, loc=0.0, scale=sd),
        )
    return effects

SubpopulationInfections

SubpopulationInfections(
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    baseline_rt_process: TemporalProcess,
    subpop_rt_deviation_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
    infection_feedback_strength_rv: RandomVariable | None = None,
    infection_feedback_pmf_rv: RandomVariable | None = None,
)

Bases: BaseSubpopulationInfections

Multi-subpopulation renewal model with hierarchical \(\mathcal{R}(t)\) structure.

Each subpopulation has its own renewal equation with \(\mathcal{R}(t)\) deviating from a shared baseline. Suitable when transmission dynamics vary substantially across subpopulations. Each subpopulation renews independently: infections in one subpopulation do not depend on infections in any other.

Mathematical form: - Baseline \(\mathcal{R}(t)\): log[R_baseline(t)] ~ TemporalProcess - Subpopulation \(\mathcal{R}(t)\): log R_k(t) = log[R_baseline(t)] + delta_k(t) - Deviations: delta_k(t) ~ TemporalProcess with sum-to-zero constraint - Starting values on the first day of the initialization period, from log_rt_time_0_rv: a scalar s gives log[R_baseline(0)] = s and delta_k(0) = 0; one value v_k per subpopulation gives log[R_baseline(0)] = mean(v) and delta_k(0) = v_k - mean(v), so log R_k(0) = v_k. - Renewal per subpop: I_k(t) = R_k(t) * sum_tau I_k(t-tau) * g(tau) - Aggregate total: I_aggregate(t) = sum_k p_k * I_k(t)

Infection feedback is optional. It is on when both infection_feedback_strength_rv and infection_feedback_pmf_rv are given. On each day after the initialization period, it multiplies the raw \(\mathcal{R}(t)\) of each subpopulation by a factor that depends on recent infections in that subpopulation:

\[ \mathcal{R}^{\mathrm{adj}}_k(t) = \mathcal{R}_k(t) \exp\left(\gamma_k(t) \sum_{\tau=1}^{T_g} I_k(t - \tau) f(\tau)\right) \]

where \(I_k(t)\) is new infections per resident of subpopulation \(k\), \(\gamma_k(t)\) is the feedback strength, \(f\) is the feedback PMF, and \(T_g\) is the length of the generation interval. The renewal equation of subpopulation \(k\) then uses the adjusted \(\mathcal{R}^{\mathrm{adj}}_k(t)\) in place of \(\mathcal{R}_k(t)\). Negative strength lowers \(\mathcal{R}_k(t)\) when recent infections are high, and positive strength raises it. With both arguments left out, infections follow the plain renewal equation. Sampling raises ValueError if the feedback strength does not broadcast to shape (n_days_post_init, n_subpops), or if the feedback PMF length differs from the generation interval length.

The constructor specifies model structure (priors, temporal processes). Population structure (subpop_fractions) is provided at sample time, allowing a single model to be fit to multiple jurisdictions.

Notes

Sum-to-zero constraint on deviations ensures R_baseline(t) is the geometric mean of subpopulation \(\mathcal{R}(t)\) values, providing identifiability. rt_baseline and rt_subpop are raw reproduction numbers generated by the temporal-process hierarchy. rt_subpop_adjusted is the reproduction number used to generate infections. It equals rt_subpop unless infection feedback is on.

Initialize hierarchical infections process.

Parameters:

Name	Type	Description	Default
name	str	Name prefix for numpyro sample sites. All deterministic quantities are recorded under this scope (e.g., "{name}::rt_baseline", "{name}::rt_subpop", and "{name}::rt_subpop_adjusted").	required
gen_int_rv	RandomVariable	Generation interval PMF	required
n_initialization_points	int	Number of initialization days before day 0.	required
I0_rv	RandomVariable	New infections per resident on the first day of the initialization period, the earliest day on the time axis. This is not prevalence. Returns either a scalar, one value for all subpopulations, or an array of shape (n_subpops,), one value per subpopulation in the order of subpop_fractions.	required
baseline_rt_process	TemporalProcess	Temporal process for baseline \(\mathcal{R}(t)\) dynamics	required
subpop_rt_deviation_process	TemporalProcess	Temporal process for subpopulation deviations	required
log_rt_time_0_rv	RandomVariable	Starting log(\(\mathcal{R}(t)\)) on the first day of the initialization period. Returns either a scalar, one value for all subpopulations, or an array of shape (n_subpops,), one value per subpopulation in the order of subpop_fractions. Per-subpopulation values also give each subpopulation its own initial growth rate. The baseline starts at the mean of the per-subpopulation values, so giving every subpopulation the same prior narrows the prior on the baseline's starting value: with dist.Normal(0.0, s).expand([n_subpops]), its standard deviation is s / sqrt(n_subpops), not s.	required
infection_feedback_strength_rv	RandomVariable | None	Infection feedback strength \(\gamma\). Returns an array that broadcasts to (n_days_post_init, n_subpops): a scalar: one strength for all subpopulations and days; shape (n_subpops,): one strength per subpopulation; shape (n_days_post_init, 1): one strength per day; shape (n_days_post_init, n_subpops): one strength per day and subpopulation. A 1-D array is read as one strength per subpopulation, so a strength that varies by day must have shape (n_days_post_init, 1). Give together with infection_feedback_pmf_rv, or leave both out for plain renewal.	None
infection_feedback_pmf_rv	RandomVariable | None	Feedback PMF \(f\), which weights recent infections in the feedback factor. Element 0 weights infections on the previous day. Must have the same length as the generation interval. Give together with infection_feedback_strength_rv.	None

Raises:

Type	Description
ValueError	If required parameters are missing or invalid, or if only one of the two infection feedback arguments is given.

Source code in pyrenew/latent/subpopulation_infections.py

def __init__(
    self,
    *,
    name: str,
    gen_int_rv: RandomVariable,
    n_initialization_points: int,
    I0_rv: RandomVariable,
    baseline_rt_process: TemporalProcess,
    subpop_rt_deviation_process: TemporalProcess,
    log_rt_time_0_rv: RandomVariable,
    infection_feedback_strength_rv: RandomVariable | None = None,
    infection_feedback_pmf_rv: RandomVariable | None = None,
) -> None:
    """
    Initialize hierarchical infections process.

    Parameters
    ----------
    name
        Name prefix for numpyro sample sites. All deterministic
        quantities are recorded under this scope (e.g.,
        ``"{name}::rt_baseline"``, ``"{name}::rt_subpop"``, and
        ``"{name}::rt_subpop_adjusted"``).
    gen_int_rv
        Generation interval PMF
    n_initialization_points
        Number of initialization days before day 0.
    I0_rv
        New infections per resident on the first day of the
        initialization period, the earliest day on the time axis. This
        is not prevalence. Returns either a scalar, one value for all
        subpopulations, or an array of shape ``(n_subpops,)``, one value
        per subpopulation in the order of ``subpop_fractions``.
    baseline_rt_process
        Temporal process for baseline $\\mathcal{R}(t)$ dynamics
    subpop_rt_deviation_process
        Temporal process for subpopulation deviations
    log_rt_time_0_rv
        Starting log($\\mathcal{R}(t)$) on the first day of the
        initialization period. Returns either a scalar, one value for all
        subpopulations, or an array of shape ``(n_subpops,)``, one value per
        subpopulation in the order of ``subpop_fractions``. Per-subpopulation
        values also give each subpopulation its own initial growth rate.
        The baseline starts at the mean of the per-subpopulation values, so
        giving every subpopulation the same prior narrows the prior on the
        baseline's starting value: with
        ``dist.Normal(0.0, s).expand([n_subpops])``, its standard deviation
        is ``s / sqrt(n_subpops)``, not ``s``.
    infection_feedback_strength_rv
        Infection feedback strength $\\gamma$. Returns an array that
        broadcasts to ``(n_days_post_init, n_subpops)``:

        - a scalar: one strength for all subpopulations and days;
        - shape ``(n_subpops,)``: one strength per subpopulation;
        - shape ``(n_days_post_init, 1)``: one strength per day;
        - shape ``(n_days_post_init, n_subpops)``: one strength per day
          and subpopulation.

        A 1-D array is read as one strength per subpopulation, so a
        strength that varies by day must have shape
        ``(n_days_post_init, 1)``. Give together with
        ``infection_feedback_pmf_rv``, or leave both out for plain renewal.
    infection_feedback_pmf_rv
        Feedback PMF $f$, which weights recent infections in the feedback
        factor. Element 0 weights infections on the previous day. Must have
        the same length as the generation interval. Give together with
        ``infection_feedback_strength_rv``.

    Raises
    ------
    ValueError
        If required parameters are missing or invalid, or if only one of
        the two infection feedback arguments is given.
    """
    super().__init__(
        name=name,
        gen_int_rv=gen_int_rv,
        n_initialization_points=n_initialization_points,
        I0_rv=I0_rv,
        baseline_rt_process=baseline_rt_process,
        subpop_rt_deviation_process=subpop_rt_deviation_process,
        log_rt_time_0_rv=log_rt_time_0_rv,
    )

    _validate_feedback_options(
        infection_feedback_strength_rv, infection_feedback_pmf_rv
    )
    self.infection_feedback_strength_rv = infection_feedback_strength_rv
    self.infection_feedback_pmf_rv = infection_feedback_pmf_rv

TemporalProcess

Bases: Protocol

Protocol for temporal processes generating time-varying parameters.

Used for jurisdiction-level Rt dynamics, subpopulation deviations, or allocation trajectories. All processes return 2D arrays of shape (n_timepoints, n_processes) for consistent handling.

Attributes:

Name	Type	Description
step_size	int	Number of consecutive timepoints that share the same sampled value. Defaults to 1 for the standard processes (one independent sample per timepoint). Wrapper processes like StepwiseTemporalProcess and WeeklyTemporalProcess expose a larger value so that model builders can inspect R(t) parametrization cadence.

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "temporal",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample temporal trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Per-process starting value or initial-location parameter. Processes with a deterministic initial state return this value at the first timepoint. AR1 uses it as the mean of the initial-state prior. Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample site names to avoid collisions	'temporal'
first_day_dow	int | None	Day of week for element 0 of the shared model time axis (0=Monday, ..., 6=Sunday). Standard temporal processes ignore this value; calendar-aligned wrappers may use it.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "temporal",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample temporal trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Per-process starting value or initial-location parameter. Processes
        with a deterministic initial state return this value at the first
        timepoint. ``AR1`` uses it as the mean of the initial-state prior.
        Scalar values are broadcast to all processes; arrays must have
        shape ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample site names to avoid collisions
    first_day_dow
        Day of week for element 0 of the shared model time axis
        (0=Monday, ..., 6=Sunday). Standard temporal processes ignore
        this value; calendar-aligned wrappers may use it.

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    ...

WeeklyTemporalProcess

WeeklyTemporalProcess(inner: TemporalProcess, start_dow: int)

Bases: TemporalProcess

Parameterize an inner temporal process at a calendar-week cadence.

Each output timepoint in the same calendar week shares a single sampled value from the inner process. Use when a parameter should be estimated once per week while downstream model components still need one value per evaluation timepoint. For example, a weekly-parameterized R(t) process can return daily R(t) values for a daily deterministic renewal equation.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the weekly trajectory. Must satisfy the TemporalProcess protocol.	required
start_dow	int	Day-of-week on which the calendar-week cycle begins (0=Monday, 6=Sunday, ISO convention). Use 6 for MMWR Sunday-Saturday epiweeks or 0 for ISO Monday-Sunday weeks.	required

Initialize weekly temporal process.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the weekly trajectory.	required
start_dow	int	Day-of-week on which the calendar-week cycle begins (0=Monday, 6=Sunday, ISO convention).	required

Source code in pyrenew/latent/temporal_processes.py

def __init__(self, inner: TemporalProcess, start_dow: int) -> None:
    """
    Initialize weekly temporal process.

    Parameters
    ----------
    inner
        Inner ``TemporalProcess`` that generates the weekly trajectory.
    start_dow
        Day-of-week on which the calendar-week cycle begins
        (0=Monday, 6=Sunday, ISO convention).
    """
    validate_dow(start_dow, "start_dow")
    self.inner = inner
    self.start_dow = start_dow

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"WeeklyTemporalProcess(inner={self.inner!r}, start_dow={self.start_dow!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "weekly",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample weekly trajectory from inner process and broadcast.

Computes the number of weekly time steps needed for the configured calendar-week cycle, samples the inner process at that cadence, then broadcasts each weekly value to the per-timepoint axis and trims to n_timepoints. The returned value always has one row per evaluation timepoint, regardless of the inner parameter cadence. The weekly trajectory is recorded as a NumPyro deterministic site named "{name_prefix}_weekly".

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of per-timepoint outputs to produce.	required
initial_value	float | ArrayLike | None	Initial value(s) for the inner process. Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites; forwarded to the inner process.	'weekly'
first_day_dow	int | None	Day of week for element 0 of the shared model time axis (0=Monday, ..., 6=Sunday). Required.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes), constant within each calendar week.

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "weekly",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample weekly trajectory from inner process and broadcast.

    Computes the number of weekly time steps needed for the configured
    calendar-week cycle, samples the inner process at that cadence, then
    broadcasts each weekly value to the per-timepoint axis and trims to
    ``n_timepoints``. The returned value always has one row per evaluation
    timepoint, regardless of the inner parameter cadence. The weekly
    trajectory is recorded as a NumPyro deterministic site named
    ``"{name_prefix}_weekly"``.

    Parameters
    ----------
    n_timepoints
        Number of per-timepoint outputs to produce.
    initial_value
        Initial value(s) for the inner process. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites; forwarded to the inner process.
    first_day_dow
        Day of week for element 0 of the shared model time axis
        (0=Monday, ..., 6=Sunday). Required.

    Returns
    -------
    ArrayLike
        Trajectories of shape ``(n_timepoints, n_processes)``, constant
        within each calendar week.
    """
    n_steps = self._resolve_n_weekly(n_timepoints, first_day_dow)
    # first_day_dow intentionally not forwarded: inner operates on the
    # weekly axis; the outer's axis-origin day-of-week does not apply.
    weekly = self.inner.sample(
        n_timepoints=n_steps,
        initial_value=initial_value,
        n_processes=n_processes,
        name_prefix=name_prefix,
    )
    numpyro.deterministic(f"{name_prefix}_weekly", weekly)
    return weekly_to_daily(
        weekly,
        week_start_dow=self.start_dow,
        output_data_first_dow=first_day_dow,
    )[:n_timepoints]

compute_infections_from_rt

compute_infections_from_rt(
    I0: ArrayLike, Rt: ArrayLike, reversed_generation_interval_pmf: ArrayLike
) -> ndarray

Generate infections according to a renewal process with a time-varying reproduction number \(\mathcal{R}(t)\)

Parameters:

Name	Type	Description	Default
I0	ArrayLike	Array of initial infections of the same length as the generation interval pmf vector.	required
Rt	ArrayLike	Timeseries of \(\mathcal{R}(t)\) values	required
reversed_generation_interval_pmf	ArrayLike	discrete probability mass vector representing the generation interval of the infection process, where the final entry represents an infection 1 time unit in the past, the second-to-last entry represents an infection two time units in the past, etc.	required

Returns:

Type	Description
ndarray	The timeseries of infections.

Source code in pyrenew/latent/infection_functions.py

def compute_infections_from_rt(
    I0: ArrayLike,
    Rt: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
) -> jnp.ndarray:
    """
    Generate infections according to a
    renewal process with a time-varying
    reproduction number $\\mathcal{R}(t)$

    Parameters
    ----------
    I0
        Array of initial infections of the
        same length as the generation interval
        pmf vector.
    Rt
        Timeseries of $\\mathcal{R}(t)$ values
    reversed_generation_interval_pmf
        discrete probability mass vector
        representing the generation interval
        of the infection process, where the final
        entry represents an infection 1 time unit in the
        past, the second-to-last entry represents
        an infection two time units in the past, etc.

    Returns
    -------
    jnp.ndarray
        The timeseries of infections.
    """
    incidence_func = new_convolve_scanner(
        reversed_generation_interval_pmf, IdentityTransform()
    )

    latest, all_infections = jax.lax.scan(f=incidence_func, init=I0, xs=Rt)

    return all_infections

compute_infections_from_rt_with_feedback

compute_infections_from_rt_with_feedback(
    I0: ArrayLike,
    Rt_raw: ArrayLike,
    infection_feedback_strength: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
    reversed_infection_feedback_pmf: ArrayLike,
) -> tuple

Generate infections according to a renewal process with infection feedback (generalizing Asher 2018 <https://doi.org/10.1016/j.epidem.2017.02.009>_).

Parameters:

Name	Type	Description	Default
I0	ArrayLike	Array of initial infections of the same length as the generation interval pmf vector.	required
Rt_raw	ArrayLike	Timeseries of raw \(\mathcal{R}(t)\) values not adjusted by infection feedback	required
infection_feedback_strength	ArrayLike	Strength of the infection feedback. Either a scalar (constant feedback strength in time) or a vector representing the infection feedback strength at a given point in time.	required
reversed_generation_interval_pmf	ArrayLike	discrete probability mass vector representing the generation interval of the infection process, where the final entry represents an infection 1 time unit in the past, the second-to-last entry represents an infection two time units in the past, etc.	required
reversed_infection_feedback_pmf	ArrayLike	discrete probability mass vector representing the infection feedback process, where the final entry represents the relative contribution to infection feedback from infections that occurred 1 time unit in the past, the second-to-last entry represents the contribution from infections that occurred 2 time units in the past, etc.	required

Returns:

Type	Description
tuple	A tuple (infections, Rt_adjusted), where Rt_adjusted is the infection-feedback-adjusted timeseries of the reproduction number \(\mathcal{R}(t)\) and infections is the incident infection timeseries.

Notes

This function implements the following renewal process:

\[ \begin{aligned} I(t) & = \mathcal{R}(t)\sum_{\tau=1}^{T_g}I(t - \tau)g(\tau) \\ \mathcal{R}(t) & = \mathcal{R}^u(t)\exp\left(\gamma(t)\ \sum_{\tau=1}^{T_f}I(t - \tau)f(\tau)\right) \end{aligned} \]

where \(\mathcal{R}(t)\) is the reproductive number, \(\gamma(t)\) is the infection feedback strength, \(T_g\) is the max-length of the generation interval, \(\mathcal{R}^u(t)\) is the raw reproduction number, \(f(t)\) is the infection feedback pmf, and \(T_f\) is the max-length of the infection feedback pmf.

Note that negative \(\gamma(t)\) implies that recent incident infections reduce \(\mathcal{R}(t)\) below its raw value in the absence of feedback, while positive \(\gamma\) implies that recent incident infections increase \(\mathcal{R}(t)\) above its raw value, and \(\gamma(t)=0\) implies no feedback.

In general, negative \(\gamma\) is the more common modeling choice, as it can be used to model susceptible depletion, reductions in contact rate due to awareness of high incidence, et cetera.

Source code in pyrenew/latent/infection_functions.py

def compute_infections_from_rt_with_feedback(
    I0: ArrayLike,
    Rt_raw: ArrayLike,
    infection_feedback_strength: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
    reversed_infection_feedback_pmf: ArrayLike,
) -> tuple:
    r"""
    Generate infections according to
    a renewal process with infection
    feedback (generalizing `Asher 2018
    <https://doi.org/10.1016/j.epidem.2017.02.009>`_).

    Parameters
    ----------
    I0
        Array of initial infections of the
        same length as the generation interval
        pmf vector.
    Rt_raw
        Timeseries of raw $\mathcal{R}(t)$ values not
        adjusted by infection feedback
    infection_feedback_strength
        Strength of the infection feedback.
        Either a scalar (constant feedback
        strength in time) or a vector representing
        the infection feedback strength at a
        given point in time.
    reversed_generation_interval_pmf
        discrete probability mass vector
        representing the generation interval
        of the infection process, where the final
        entry represents an infection 1 time unit in the
        past, the second-to-last entry represents
        an infection two time units in the past, etc.
    reversed_infection_feedback_pmf
        discrete probability mass vector
        representing the infection feedback
        process, where the final entry represents
        the relative contribution to infection
        feedback from infections that occurred
        1 time unit in the past, the second-to-last
        entry represents the contribution from infections
        that occurred 2 time units in the past, etc.

    Returns
    -------
    tuple
        A tuple ``(infections, Rt_adjusted)``,
        where `Rt_adjusted` is the infection-feedback-adjusted
        timeseries of the reproduction number $\mathcal{R}(t)$
        and `infections` is the incident infection timeseries.

    Notes
    -----
    This function implements the following renewal process:

    ```math
    \begin{aligned}
    I(t) & = \mathcal{R}(t)\sum_{\tau=1}^{T_g}I(t - \tau)g(\tau) \\
    \mathcal{R}(t) & = \mathcal{R}^u(t)\exp\left(\gamma(t)\
        \sum_{\tau=1}^{T_f}I(t - \tau)f(\tau)\right)
    \end{aligned}
    ```

    where $\mathcal{R}(t)$ is the reproductive number,
    $\gamma(t)$ is the infection feedback strength,
    $T_g$ is the max-length of the
    generation interval, $\mathcal{R}^u(t)$ is the raw reproduction
    number, $f(t)$ is the infection feedback pmf, and $T_f$
    is the max-length of the infection feedback pmf.

    Note that negative $\gamma(t)$ implies
    that recent incident infections reduce $\mathcal{R}(t)$
    below its raw value in the absence of feedback, while
    positive $\gamma$ implies that recent incident infections
    *increase* $\mathcal{R}(t)$ above its raw value, and
    $\gamma(t)=0$ implies no feedback.

    In general, negative $\gamma$ is the more common modeling
    choice, as it can be used to model susceptible depletion,
    reductions in contact rate due to awareness of high incidence,
    et cetera.
    """
    feedback_scanner = new_double_convolve_scanner(
        arrays_to_convolve=(
            reversed_infection_feedback_pmf,
            reversed_generation_interval_pmf,
        ),
        transforms=(ExpTransform(), IdentityTransform()),
    )
    latest, infs_and_R_adj = jax.lax.scan(
        f=feedback_scanner,
        init=I0,
        xs=(infection_feedback_strength, Rt_raw),
    )

    infections, R_adjustment = infs_and_R_adj
    Rt_adjusted = R_adjustment * Rt_raw
    return infections, Rt_adjusted

compute_infections_from_rt_with_mobility

compute_infections_from_rt_with_mobility(
    I0: ArrayLike,
    Rt_raw: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
    mobility_matrix: ArrayLike,
    subpop_fractions: ArrayLike,
    susceptible_depletion: bool = True,
) -> tuple[ndarray, ndarray]

Generate infections from a renewal process coupled by a fixed mobility matrix.

Residents of each home region spend part of each day at other locations. Transmission happens at each location, at that location's reproduction number, among the people present there. Each infection is assigned to the home region of the person infected.

Parameters:

Name	Type	Description	Default
I0	ArrayLike	New infections per resident on each day of the initialization period, shape (n_init, n), with n_init at least the length of the generation interval. Every row counts toward susceptible depletion. Only the last len(reversed_generation_interval_pmf) rows enter the generation-interval sum. Entries must be finite and non-negative. With depletion enabled, I0.sum(axis=0) must be at most 1 in every region.	required
Rt_raw	ArrayLike	Finite, non-negative reproduction number at each location on each day after the initialization period, shape (n_days, n).	required
reversed_generation_interval_pmf	ArrayLike	Generation interval PMF in reverse order: the last entry is one day back, the second-to-last two days back, and so on.	required
mobility_matrix	ArrayLike	Matrix used on every day, shape (n, n). mobility_matrix[j, k] is the fraction of residents of home region k present at location j. Each column sums to 1.	required
subpop_fractions	ArrayLike	Each home region's share of the total population, shape (n,). Sums to 1.	required
susceptible_depletion	bool	If True (default), new infections among residents of each region are scaled by, and capped at, the fraction of those residents not yet infected. If False, neither adjustment is applied.	True

Returns:

Type	Description
tuple[ndarray, ndarray]	New infections per resident, shape (n_days, n), and the reproduction number at each location after susceptibility adjustment, the same shape. The second array does not account for the incidence ceiling. Without depletion, it is Rt_raw.

Raises:

Type	Description
ValueError	If concrete I0 values are non-finite or negative, or if their sum exceeds 1 in any region with depletion enabled. Value checks are skipped during JAX tracing; callers must ensure traced inputs meet the same requirements.

Notes

This implements Equation 2 of Roy, Clapham, and Mishra (2025), https://doi.org/10.1371/journal.pcbi.1013642. In infection counts (before converting to infections per resident and applying the incidence ceiling), the equation is:

\[ I_{t,i} = \sum_{j=1}^{n} \frac{m_{ji} S_{t,i}}{\sum_{l=1}^{n} m_{jl} N_l} R_{t,j} \left[ \sum_{k=1}^{n} m_{jk} \sum_{\tau} I_{\tau,k}\, g_{t-\tau} \right] \]

where \(m_{jk}\) is mobility_matrix[j, k], \(N_k\) is the number of residents of region \(k\), \(R_{t,j}\) is the reproduction number at location \(j\), \(g\) is the generation interval, and \(S_{t,i}\) is the number of residents of region \(i\) not infected before day \(t\).

With M = mobility_matrix and p = subpop_fractions, one day with depletion enabled is:

present = M @ p
share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None]
infectiousness = reversed_generation_interval_pmf @ history
susceptible = 1.0 - cumulative
new_infections_raw = M.T @ (rt * (share_present @ infectiousness))
new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible)

share_present[j, k] is the share of the people present at location j who live in region k. cumulative counts every day of I0 and every earlier day of the result.

A location with no one present contributes no infections. With depletion enabled, its adjusted reproduction number is 0.

The incidence ceiling is an addition to Roy et al.'s Equation 2. With valid initialization, it keeps cumulative infections at most 1 per resident, so the susceptible fraction needs no floor. Without depletion, cumulative infections may exceed 1 per resident.

Source code in pyrenew/latent/infection_functions.py

def compute_infections_from_rt_with_mobility(
    I0: ArrayLike,
    Rt_raw: ArrayLike,
    reversed_generation_interval_pmf: ArrayLike,
    mobility_matrix: ArrayLike,
    subpop_fractions: ArrayLike,
    susceptible_depletion: bool = True,
) -> tuple[jnp.ndarray, jnp.ndarray]:
    r"""
    Generate infections from a renewal process coupled by a fixed mobility matrix.

    Residents of each home region spend part of each day at other locations.
    Transmission happens at each location, at that location's reproduction
    number, among the people present there. Each infection is assigned to
    the home region of the person infected.

    Parameters
    ----------
    I0
        New infections per resident on each day of the initialization
        period, shape ``(n_init, n)``, with ``n_init`` at least the length of
        the generation interval. Every row counts toward susceptible
        depletion. Only the last ``len(reversed_generation_interval_pmf)``
        rows enter the generation-interval sum. Entries must be finite and
        non-negative. With depletion enabled, ``I0.sum(axis=0)`` must be
        at most 1 in every region.
    Rt_raw
        Finite, non-negative reproduction number at each location on each
        day after the initialization period, shape ``(n_days, n)``.
    reversed_generation_interval_pmf
        Generation interval PMF in reverse order: the last entry is one day
        back, the second-to-last two days back, and so on.
    mobility_matrix
        Matrix used on every day, shape ``(n, n)``.
        ``mobility_matrix[j, k]`` is the fraction of residents of home
        region ``k`` present at location ``j``. Each column sums to 1.
    subpop_fractions
        Each home region's share of the total population, shape ``(n,)``.
        Sums to 1.
    susceptible_depletion
        If True (default), new infections among residents of each region
        are scaled by, and capped at, the fraction of those residents not
        yet infected. If False, neither adjustment is applied.

    Returns
    -------
    tuple[jnp.ndarray, jnp.ndarray]
        New infections per resident, shape ``(n_days, n)``, and the
        reproduction number at each location after susceptibility
        adjustment, the same shape. The second array does not account for
        the incidence ceiling. Without depletion, it is ``Rt_raw``.

    Raises
    ------
    ValueError
        If concrete ``I0`` values are non-finite or negative, or if their
        sum exceeds 1 in any region with depletion enabled. Value checks
        are skipped during JAX tracing; callers must ensure traced inputs
        meet the same requirements.

    Notes
    -----
    This implements Equation 2 of Roy, Clapham, and Mishra (2025),
    <https://doi.org/10.1371/journal.pcbi.1013642>. In infection counts (before
    converting to infections per resident and applying the incidence
    ceiling), the equation is:

    ```math
    I_{t,i} = \sum_{j=1}^{n} \frac{m_{ji} S_{t,i}}{\sum_{l=1}^{n} m_{jl} N_l}
    R_{t,j} \left[ \sum_{k=1}^{n} m_{jk} \sum_{\tau} I_{\tau,k}\, g_{t-\tau} \right]
    ```

    where $m_{jk}$ is ``mobility_matrix[j, k]``, $N_k$ is the number of
    residents of region $k$, $R_{t,j}$ is the reproduction number at
    location $j$, $g$ is the generation interval, and $S_{t,i}$ is the
    number of residents of region $i$ not infected before day $t$.

    With ``M = mobility_matrix`` and ``p = subpop_fractions``, one day with
    depletion enabled is:

    ```python
    present = M @ p
    share_present = M * p[None, :] / jnp.where(present > 0, present, 1.0)[:, None]
    infectiousness = reversed_generation_interval_pmf @ history
    susceptible = 1.0 - cumulative
    new_infections_raw = M.T @ (rt * (share_present @ infectiousness))
    new_infections = jnp.minimum(susceptible * new_infections_raw, susceptible)
    ```

    ``share_present[j, k]`` is the share of the people present at location
    ``j`` who live in region ``k``. ``cumulative`` counts every day of
    ``I0`` and every earlier day of the result.

    A location with no one present contributes no infections. With
    depletion enabled, its adjusted reproduction number is 0.

    The incidence ceiling is an addition to Roy et al.'s Equation 2.
    With valid initialization, it keeps cumulative infections at most 1
    per resident, so the susceptible fraction needs no floor. Without
    depletion, cumulative infections may exceed 1 per resident.
    """
    I0 = jnp.asarray(I0)
    invalid_initial = jnp.any(~jnp.isfinite(I0) | (I0 < 0))
    if not_jax_tracer(invalid_initial) and invalid_initial:
        raise ValueError("I0 must contain only finite, non-negative infections.")
    cumulative = jnp.sum(I0, axis=0)
    if susceptible_depletion:
        exceeds_population = jnp.any(cumulative > 1.0)
        if not_jax_tracer(exceeds_population) and exceeds_population:
            raise ValueError(
                "I0.sum(axis=0) must be <= 1 in every region when "
                "susceptible depletion is enabled."
            )
    present = mobility_matrix @ subpop_fractions
    share_present = (
        mobility_matrix
        * subpop_fractions[jnp.newaxis, :]
        / jnp.where(present > 0, present, 1.0)[:, jnp.newaxis]
    )
    n_gen = reversed_generation_interval_pmf.size
    initial_state = (I0[-n_gen:], cumulative)
    step = functools.partial(
        _mobility_renewal_step,
        mobility_matrix=mobility_matrix,
        share_present=share_present,
        reversed_generation_interval_pmf=reversed_generation_interval_pmf,
        susceptible_depletion=susceptible_depletion,
    )
    _, (infections, Rt_adjusted) = jax.lax.scan(step, initial_state, Rt_raw)
    return infections, Rt_adjusted

logistic_susceptibility_adjustment

logistic_susceptibility_adjustment(
    I_raw_t: float, frac_susceptible: float, n_population: float
) -> float

Apply the logistic susceptibility adjustment to a potential new incidence I_raw_t proposed in equation 6 of Bhatt et al 2023 <https://doi.org/10.1093/jrsssa/qnad030>_.

Parameters:

Name	Type	Description	Default
I_raw_t	float	The "unadjusted" incidence at time t, i.e. the incidence given an infinite number of available susceptible individuals.	required
frac_susceptible	float	fraction of remaining susceptible individuals in the population	required
n_population	float	Total size of the population.	required

Returns:

Type	Description
float	The adjusted value of \(I(t)\).

Source code in pyrenew/latent/infection_functions.py

def logistic_susceptibility_adjustment(
    I_raw_t: float,
    frac_susceptible: float,
    n_population: float,
) -> float:
    """
    Apply the logistic susceptibility
    adjustment to a potential new
    incidence `I_raw_t` proposed in
    equation 6 of `Bhatt et al 2023
    <https://doi.org/10.1093/jrsssa/qnad030>`_.

    Parameters
    ----------
    I_raw_t
        The "unadjusted" incidence at time t,
        i.e. the incidence given an infinite
        number of available susceptible individuals.
    frac_susceptible
        fraction of remaining susceptible individuals
        in the population
    n_population
        Total size of the population.

    Returns
    -------
    float
        The adjusted value of $I(t)$.
    """
    approx_frac_infected = 1 - jnp.exp(-I_raw_t / n_population)
    return n_population * frac_susceptible * approx_frac_infected

Temporal Processes

Temporal processes for latent infection models.

Provides time-series processes for modeling Rt dynamics and subpopulation deviations in hierarchical infection models. All processes return 2D arrays of shape (n_timepoints, n_processes) through a unified TemporalProcess protocol.

Relationship to pyrenew.process:

This module provides high-level, domain-specific wrappers around the low-level building blocks in pyrenew.process. The key differences:

Aspect	pyrenew.process	pyrenew.latent.temporal_processes
Abstraction level	Low-level composable primitives	High-level domain-specific API
Interface	Varied signatures per class	Unified TemporalProcess protocol
Target use	General time-series modeling	Rt dynamics, hierarchical infections
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, RandomWalk) when building novel statistical models or when you need fine-grained control over array shapes and numpyro sampling semantics.

• Use this module's classes (AR1, DifferencedAR1, RandomWalk) when modeling Rt trajectories in hierarchical infection models. These provide a consistent interface, automatic vectorization, and enforce epidemiologically-sensible constraints.

Temporal processes provided:

• AR1: Autoregressive process with mean reversion. Keeps Rt bounded near a baseline. Wraps pyrenew.process.ARProcess.
• DifferencedAR1: AR(1) on first differences. Allows persistent trends while stabilizing the growth rate. Wraps pyrenew.process.DifferencedProcess.
• RandomWalk: No mean reversion. Rt can drift without bound. Wraps pyrenew.process.RandomWalk.
• StepwiseTemporalProcess: Wrapper that parameterizes any inner TemporalProcess at a coarser model-index cadence and broadcasts to the per-timepoint scale by repetition.
• WeeklyTemporalProcess: Wrapper that parameterizes any inner TemporalProcess at a calendar-week cadence and broadcasts to the per-timepoint scale by repetition.

All implementations satisfy the TemporalProcess protocol and can be used interchangeably in hierarchical infection models.

AR1

AR1(
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
)

Bases: TemporalProcess

AR(1) process.

Each value depends on the previous value plus noise, with reversion toward a mean level. Keeps Rt bounded near a baseline — values that drift away are "pulled back" over time.

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 innovation branch uses pyrenew.process.ARProcess to construct the trajectory from sampled innovations. The state branch uses local scanned conditional transitions.

Parameters:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of noise at each time step. Larger values produce more volatile trajectories; smaller values produce smoother ones.	required
parameterization	Parameterization	Which latent object to sample: "innovation" (default) or "state".	'innovation'

Initialize AR(1) process.

Parameters:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of innovations.	required
parameterization	Parameterization	"innovation" (default) or "state". See class docstring.	'innovation'

Raises:

Type	Description
TypeError	If autoreg_rv or innovation_sd_rv are not RandomVariable instances
ValueError	If innovation_sd_rv is a DeterministicVariable with any value <= 0, or if parameterization is not a recognized string

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
) -> None:
    """
    Initialize AR(1) process.

    Parameters
    ----------
    autoreg_rv
        RandomVariable that returns the autoregressive coefficient. For
        stationarity, |autoreg| < 1, but this is not enforced (use priors
        to constrain if needed).
    innovation_sd_rv
        RandomVariable that returns the standard deviation of innovations.
    parameterization
        ``"innovation"`` (default) or ``"state"``. See class docstring.

    Raises
    ------
    TypeError
        If autoreg_rv or innovation_sd_rv are not RandomVariable instances
    ValueError
        If innovation_sd_rv is a DeterministicVariable with any value <= 0,
        or if parameterization is not a recognized string
    """
    if not isinstance(autoreg_rv, RandomVariable):
        raise TypeError("autoreg_rv must be a RandomVariable")
    if not isinstance(innovation_sd_rv, RandomVariable):
        raise TypeError("innovation_sd_rv must be a RandomVariable")
    _validate_deterministic_innovation_sd(innovation_sd_rv)
    _validate_parameterization(parameterization)
    self.autoreg_rv = autoreg_rv
    self.innovation_sd_rv = innovation_sd_rv
    self.parameterization = parameterization
    self.ar_process = ARProcess(name="ar1")

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"AR1(autoreg_rv={self.autoreg_rv}, "
        f"innovation_sd_rv={self.innovation_sd_rv}, "
        f"parameterization={self.parameterization!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample AR(1) trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Mean of the initial-state prior. The first returned value is sampled as Normal(initial_value, innovation_sd / sqrt(1 - autoreg**2)). Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites	'ar1'
first_day_dow	int | None	Unused. See pyrenew.latent.TemporalProcess.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample AR(1) trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Mean of the initial-state prior. The first returned value is sampled
        as ``Normal(initial_value, innovation_sd / sqrt(1 - autoreg**2))``.
        Scalar values are broadcast to all processes; arrays must have
        shape ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites
    first_day_dow
        Unused. See [pyrenew.latent.TemporalProcess][].

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    initial_value = _prepare_initial_value(initial_value, 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):
        init_states = numpyro.sample(
            f"{name_prefix}_init",
            dist.Normal(initial_value, stationary_sd),
        )

    if self.parameterization == "innovation":
        return self.ar_process(
            n=n_timepoints,
            init_vals=init_states[jnp.newaxis, :],
            autoreg=autoreg_broadcast[jnp.newaxis, :],
            noise_sd=innovation_sd,
            noise_name=f"{name_prefix}_noise",
        )

    if n_timepoints == 1:
        return init_states[jnp.newaxis, :]

    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,
    )
    return jnp.concatenate([init_states[jnp.newaxis, :], post_init], axis=0)

DifferencedAR1

DifferencedAR1(
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
)

Bases: TemporalProcess

AR(1) process on first differences.

Each change in value depends on the previous change plus noise, with 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.

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}) \]

\[ 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 innovation_sd. The initial state \(x_0\) is supplied deterministically as initial_value. Both produce 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:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient for differences. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of noise added to changes. Larger values produce more erratic growth rates; smaller values produce smoother trends.	required
parameterization	Parameterization	Which latent object to sample: "innovation" (default) or "state".	'innovation'

Initialize differenced AR(1) process.

Parameters:

Name	Type	Description	Default
autoreg_rv	RandomVariable	RandomVariable that returns the autoregressive coefficient for differences. For stationarity, |autoreg| < 1, but this is not enforced (use priors to constrain if needed).	required
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of innovations.	required
parameterization	Parameterization	"innovation" (default) or "state". See class docstring.	'innovation'

Raises:

Type	Description
TypeError	If autoreg_rv or innovation_sd_rv are not RandomVariable instances
ValueError	If innovation_sd_rv is a DeterministicVariable with any value <= 0, or if parameterization is not a recognized string

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    autoreg_rv: RandomVariable,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
) -> None:
    """
    Initialize differenced AR(1) process.

    Parameters
    ----------
    autoreg_rv
        RandomVariable that returns the autoregressive coefficient for
        differences. For stationarity, |autoreg| < 1, but this is not
        enforced (use priors to constrain if needed).
    innovation_sd_rv
        RandomVariable that returns the standard deviation of innovations.
    parameterization
        ``"innovation"`` (default) or ``"state"``. See class docstring.

    Raises
    ------
    TypeError
        If autoreg_rv or innovation_sd_rv are not RandomVariable instances
    ValueError
        If innovation_sd_rv is a DeterministicVariable with any value <= 0,
        or if parameterization is not a recognized string
    """
    if not isinstance(autoreg_rv, RandomVariable):
        raise TypeError("autoreg_rv must be a RandomVariable")
    if not isinstance(innovation_sd_rv, RandomVariable):
        raise TypeError("innovation_sd_rv must be a RandomVariable")
    _validate_deterministic_innovation_sd(innovation_sd_rv)
    _validate_parameterization(parameterization)
    self.autoreg_rv = autoreg_rv
    self.innovation_sd_rv = innovation_sd_rv
    self.parameterization = parameterization
    self.process = DifferencedProcess(
        name="diff_ar1",
        fundamental_process=ARProcess(name="diff_ar1_fundamental"),
        differencing_order=1,
    )

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"DifferencedAR1(autoreg_rv={self.autoreg_rv}, "
        f"innovation_sd_rv={self.innovation_sd_rv}, "
        f"parameterization={self.parameterization!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "diff_ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample differenced AR(1) trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Deterministic first state of the trajectory. Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites	'diff_ar1'
first_day_dow	int | None	Unused. See pyrenew.latent.TemporalProcess.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "diff_ar1",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample differenced AR(1) trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Deterministic first state of the trajectory. Scalar values are
        broadcast to all processes; arrays must have shape
        ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites
    first_day_dow
        Unused. See [pyrenew.latent.TemporalProcess][].

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    initial_value = _prepare_initial_value(initial_value, 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):
        init_rates = numpyro.sample(
            f"{name_prefix}_init_rate",
            dist.Normal(0, stationary_sd),
        )

    if self.parameterization == "innovation":
        return self.process(
            n=n_timepoints,
            init_vals=initial_value[jnp.newaxis, :],
            autoreg=autoreg_broadcast[jnp.newaxis, :],
            noise_sd=innovation_sd,
            fundamental_process_init_vals=init_rates[jnp.newaxis, :],
            noise_name=f"{name_prefix}_noise",
        )

    if n_timepoints == 1:
        return initial_value[jnp.newaxis, :]

    x1 = initial_value + init_rates

    if 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,
    )
    return jnp.concatenate(
        [initial_value[jnp.newaxis, :], x1[jnp.newaxis, :], post_init],
        axis=0,
    )

RandomWalk

RandomWalk(
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
)

Bases: TemporalProcess

Random walk process for log(Rt).

Each value equals the previous value plus noise, with no reversion toward a mean. Allows Rt to drift without bound — suitable when you have no prior expectation that Rt will return to a baseline.

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:

Name	Type	Description	Default
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of noise at each time step. Larger values produce faster drift; smaller values produce more gradual changes.	required
parameterization	Parameterization	Which latent object to sample: "innovation" (default) or "state".	'innovation'

Notes

Unlike AR(1), variance grows over time — the process can wander arbitrarily far from its starting point. For long time horizons, consider AR(1) if you want Rt to stay bounded near a baseline.

Initialize random walk process.

Parameters:

Name	Type	Description	Default
innovation_sd_rv	RandomVariable	RandomVariable that returns the standard deviation of innovations.	required
parameterization	Parameterization	"innovation" (default) or "state". See class docstring.	'innovation'

Raises:

Type	Description
TypeError	If innovation_sd_rv is not a RandomVariable instance
ValueError	If innovation_sd_rv is a DeterministicVariable with any value <= 0, or if parameterization is not a recognized string

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    innovation_sd_rv: RandomVariable,
    parameterization: Parameterization = "innovation",
) -> None:
    """
    Initialize random walk process.

    Parameters
    ----------
    innovation_sd_rv
        RandomVariable that returns the standard deviation of innovations.
    parameterization
        ``"innovation"`` (default) or ``"state"``. See class docstring.

    Raises
    ------
    TypeError
        If innovation_sd_rv is not a RandomVariable instance
    ValueError
        If innovation_sd_rv is a DeterministicVariable with any value <= 0,
        or if parameterization is not a recognized string
    """
    if not isinstance(innovation_sd_rv, RandomVariable):
        raise TypeError("innovation_sd_rv must be a RandomVariable")
    _validate_deterministic_innovation_sd(innovation_sd_rv)
    _validate_parameterization(parameterization)
    self.innovation_sd_rv = innovation_sd_rv
    self.parameterization = parameterization

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"RandomWalk(innovation_sd_rv={self.innovation_sd_rv}, "
        f"parameterization={self.parameterization!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "rw",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample random walk trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Deterministic first state of the trajectory. Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites	'rw'
first_day_dow	int | None	Unused. See pyrenew.latent.TemporalProcess.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "rw",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample random walk trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Deterministic first state of the trajectory. Scalar values are
        broadcast to all processes; arrays must have shape
        ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites
    first_day_dow
        Unused. See [pyrenew.latent.TemporalProcess][].

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    initial_value = _prepare_initial_value(initial_value, n_processes)

    innovation_sd = jnp.asarray(self.innovation_sd_rv())

    if self.parameterization == "innovation":
        rw = ProcessRandomWalk(
            name=f"{name_prefix}_random_walk",
            step_rv=DistributionalVariable(
                name=f"{name_prefix}_step",
                distribution=dist.Normal(
                    jnp.zeros(n_processes),
                    innovation_sd,
                ),
            ),
        )

        return rw.sample(
            init_vals=initial_value[jnp.newaxis, :],
            n=n_timepoints,
        )

    if n_timepoints == 1:
        return initial_value[jnp.newaxis, :]

    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,
    )
    return jnp.concatenate([initial_value[jnp.newaxis, :], post_init], axis=0)

StepwiseTemporalProcess

StepwiseTemporalProcess(inner: TemporalProcess, step_size: int)

Bases: TemporalProcess

Parameterize an inner temporal process at a coarser cadence and broadcast to the per-timepoint scale by model-index repetition.

Each step_size consecutive output timepoints share a single sampled value from the inner process. Use when a parameter should be estimated at a coarser cadence while downstream model components still need one value per evaluation timepoint. Blocks always start at output index 0.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the coarse-scale trajectory. Must satisfy the TemporalProcess protocol.	required
step_size	int	Number of per-timepoint units that share each inner sample. Must be a positive integer.	required

Raises:

Type	Description
ValueError	If step_size is not a positive integer.

Initialize stepwise temporal process.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the coarse trajectory.	required
step_size	int	Number of per-timepoint units that share each inner sample.	required

Raises:

Type	Description
ValueError	If step_size is not a positive integer.

Source code in pyrenew/latent/temporal_processes.py

def __init__(
    self,
    inner: TemporalProcess,
    step_size: int,
) -> None:
    """
    Initialize stepwise temporal process.

    Parameters
    ----------
    inner
        Inner ``TemporalProcess`` that generates the coarse trajectory.
    step_size
        Number of per-timepoint units that share each inner sample.

    Raises
    ------
    ValueError
        If ``step_size`` is not a positive integer.
    """
    if not isinstance(step_size, int) or step_size < 1:
        raise ValueError(f"step_size must be a positive integer, got {step_size!r}")
    self.inner = inner
    self.step_size = step_size

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"StepwiseTemporalProcess(inner={self.inner!r}, step_size={self.step_size})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "stepwise",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample coarse trajectory from inner process and broadcast.

Computes the number of coarse time steps needed for the requested length, samples the inner process at that cadence, then broadcasts each coarse value to the per-timepoint axis and trims to n_timepoints. The returned value always has one row per evaluation timepoint, regardless of the inner parameter cadence. The coarse trajectory is recorded as a NumPyro deterministic site named "{name_prefix}_coarse".

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of per-timepoint outputs to produce.	required
initial_value	float | ArrayLike | None	Initial value(s) for the inner process. Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites; forwarded to the inner process.	'stepwise'
first_day_dow	int | None	Ignored. Accepted for compatibility with the TemporalProcess protocol.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes), constant within each block of step_size consecutive rows.

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "stepwise",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample coarse trajectory from inner process and broadcast.

    Computes the number of coarse time steps needed for the requested
    length, samples the inner process at that cadence, then broadcasts each
    coarse value to the per-timepoint axis and trims to ``n_timepoints``.
    The returned value always has one row per evaluation timepoint,
    regardless of the inner parameter cadence. The coarse
    trajectory is recorded as a NumPyro deterministic site named
    ``"{name_prefix}_coarse"``.

    Parameters
    ----------
    n_timepoints
        Number of per-timepoint outputs to produce.
    initial_value
        Initial value(s) for the inner process. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites; forwarded to the inner process.
    first_day_dow
        Ignored. Accepted for compatibility with the
        ``TemporalProcess`` protocol.

    Returns
    -------
    ArrayLike
        Trajectories of shape ``(n_timepoints, n_processes)``, constant
        within each block of ``step_size`` consecutive rows.
    """
    n_steps = self._resolve_n_coarse(n_timepoints)
    coarse = self.inner.sample(
        n_timepoints=n_steps,
        initial_value=initial_value,
        n_processes=n_processes,
        name_prefix=name_prefix,
    )
    numpyro.deterministic(f"{name_prefix}_coarse", coarse)
    return jnp.repeat(coarse, repeats=self.step_size, axis=0)[:n_timepoints]

TemporalProcess

Bases: Protocol

Protocol for temporal processes generating time-varying parameters.

Used for jurisdiction-level Rt dynamics, subpopulation deviations, or allocation trajectories. All processes return 2D arrays of shape (n_timepoints, n_processes) for consistent handling.

Attributes:

Name	Type	Description
step_size	int	Number of consecutive timepoints that share the same sampled value. Defaults to 1 for the standard processes (one independent sample per timepoint). Wrapper processes like StepwiseTemporalProcess and WeeklyTemporalProcess expose a larger value so that model builders can inspect R(t) parametrization cadence.

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "temporal",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample temporal trajectory or trajectories.

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of time points to generate	required
initial_value	float | ArrayLike | None	Per-process starting value or initial-location parameter. Processes with a deterministic initial state return this value at the first timepoint. AR1 uses it as the mean of the initial-state prior. Scalar values are broadcast to all processes; arrays must have shape (n_processes,). Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample site names to avoid collisions	'temporal'
first_day_dow	int | None	Day of week for element 0 of the shared model time axis (0=Monday, ..., 6=Sunday). Standard temporal processes ignore this value; calendar-aligned wrappers may use it.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes)

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "temporal",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample temporal trajectory or trajectories.

    Parameters
    ----------
    n_timepoints
        Number of time points to generate
    initial_value
        Per-process starting value or initial-location parameter. Processes
        with a deterministic initial state return this value at the first
        timepoint. ``AR1`` uses it as the mean of the initial-state prior.
        Scalar values are broadcast to all processes; arrays must have
        shape ``(n_processes,)``. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample site names to avoid collisions
    first_day_dow
        Day of week for element 0 of the shared model time axis
        (0=Monday, ..., 6=Sunday). Standard temporal processes ignore
        this value; calendar-aligned wrappers may use it.

    Returns
    -------
    ArrayLike
        Trajectories of shape (n_timepoints, n_processes)
    """
    ...

WeeklyTemporalProcess

WeeklyTemporalProcess(inner: TemporalProcess, start_dow: int)

Bases: TemporalProcess

Parameterize an inner temporal process at a calendar-week cadence.

Each output timepoint in the same calendar week shares a single sampled value from the inner process. Use when a parameter should be estimated once per week while downstream model components still need one value per evaluation timepoint. For example, a weekly-parameterized R(t) process can return daily R(t) values for a daily deterministic renewal equation.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the weekly trajectory. Must satisfy the TemporalProcess protocol.	required
start_dow	int	Day-of-week on which the calendar-week cycle begins (0=Monday, 6=Sunday, ISO convention). Use 6 for MMWR Sunday-Saturday epiweeks or 0 for ISO Monday-Sunday weeks.	required

Initialize weekly temporal process.

Parameters:

Name	Type	Description	Default
inner	TemporalProcess	Inner TemporalProcess that generates the weekly trajectory.	required
start_dow	int	Day-of-week on which the calendar-week cycle begins (0=Monday, 6=Sunday, ISO convention).	required

Source code in pyrenew/latent/temporal_processes.py

def __init__(self, inner: TemporalProcess, start_dow: int) -> None:
    """
    Initialize weekly temporal process.

    Parameters
    ----------
    inner
        Inner ``TemporalProcess`` that generates the weekly trajectory.
    start_dow
        Day-of-week on which the calendar-week cycle begins
        (0=Monday, 6=Sunday, ISO convention).
    """
    validate_dow(start_dow, "start_dow")
    self.inner = inner
    self.start_dow = start_dow

__repr__

__repr__() -> str

Return string representation.

Source code in pyrenew/latent/temporal_processes.py

def __repr__(self) -> str:
    """Return string representation."""
    return (
        f"WeeklyTemporalProcess(inner={self.inner!r}, start_dow={self.start_dow!r})"
    )

sample

sample(
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "weekly",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike

Sample weekly trajectory from inner process and broadcast.

Computes the number of weekly time steps needed for the configured calendar-week cycle, samples the inner process at that cadence, then broadcasts each weekly value to the per-timepoint axis and trims to n_timepoints. The returned value always has one row per evaluation timepoint, regardless of the inner parameter cadence. The weekly trajectory is recorded as a NumPyro deterministic site named "{name_prefix}_weekly".

Parameters:

Name	Type	Description	Default
n_timepoints	int	Number of per-timepoint outputs to produce.	required
initial_value	float | ArrayLike | None	Initial value(s) for the inner process. Defaults to 0.0.	None
n_processes	int	Number of parallel processes.	1
name_prefix	str	Prefix for numpyro sample sites; forwarded to the inner process.	'weekly'
first_day_dow	int | None	Day of week for element 0 of the shared model time axis (0=Monday, ..., 6=Sunday). Required.	None

Returns:

Type	Description
ArrayLike	Trajectories of shape (n_timepoints, n_processes), constant within each calendar week.

Source code in pyrenew/latent/temporal_processes.py

def sample(
    self,
    n_timepoints: int,
    initial_value: float | ArrayLike | None = None,
    n_processes: int = 1,
    name_prefix: str = "weekly",
    *,
    first_day_dow: int | None = None,
) -> ArrayLike:
    """
    Sample weekly trajectory from inner process and broadcast.

    Computes the number of weekly time steps needed for the configured
    calendar-week cycle, samples the inner process at that cadence, then
    broadcasts each weekly value to the per-timepoint axis and trims to
    ``n_timepoints``. The returned value always has one row per evaluation
    timepoint, regardless of the inner parameter cadence. The weekly
    trajectory is recorded as a NumPyro deterministic site named
    ``"{name_prefix}_weekly"``.

    Parameters
    ----------
    n_timepoints
        Number of per-timepoint outputs to produce.
    initial_value
        Initial value(s) for the inner process. Defaults to 0.0.
    n_processes
        Number of parallel processes.
    name_prefix
        Prefix for numpyro sample sites; forwarded to the inner process.
    first_day_dow
        Day of week for element 0 of the shared model time axis
        (0=Monday, ..., 6=Sunday). Required.

    Returns
    -------
    ArrayLike
        Trajectories of shape ``(n_timepoints, n_processes)``, constant
        within each calendar week.
    """
    n_steps = self._resolve_n_weekly(n_timepoints, first_day_dow)
    # first_day_dow intentionally not forwarded: inner operates on the
    # weekly axis; the outer's axis-origin day-of-week does not apply.
    weekly = self.inner.sample(
        n_timepoints=n_steps,
        initial_value=initial_value,
        n_processes=n_processes,
        name_prefix=name_prefix,
    )
    numpyro.deterministic(f"{name_prefix}_weekly", weekly)
    return weekly_to_daily(
        weekly,
        week_start_dow=self.start_dow,
        output_data_first_dow=first_day_dow,
    )[:n_timepoints]