Add Hilbert Space Gaussian Processes Approximation as a Temporal Process

[damonbayer]

https://num.pyro.ai/en/stable/contrib.html#hilbert-space-gaussian-processes-approximation

[cdc-mitzimorris]

I have reviewed the numpyro HSGP documentation and tutorial, and would like to propose the following:

Implement one new class, HSGP, in pyrenew/latent/temporal_processes.py. This class, unlike the other temporal processes, does not have a corresponding pyrenew/process class. The existing pyrenew/process/ classes (ARProcess, DifferencedProcess, RandomWalk) are RandomVariable subclasses implementing stateful, scan-based sequential processes that are composable.

HSGP is fundamentally different:

• It is a batch function. It takes all time indices at once and returns the full trajectory in one call.
• It has no scan state to compose.
• The NumPyro helpers are already clean primitives. A RandomVariable wrapper adds a layer with no benefit.

Review: why use HSGP (instead of a regular GP)

HSGP decomposes into two parts with different computational character:

Part 1 — pure deterministic JAX math:

1. Time indices x = jnp.arange(n_timepoints, dtype=float).

2. Laplacian eigenfunctions phi of shape (n_timepoints, m), computed from x and boundary parameter ell. Formula for the j-th eigenfunction at position x:

   $$\phi_j(x) = \frac{1}{\sqrt{\ell}} \sin!\left(\frac{j \pi (x + \ell)}{2\ell}\right)$$

3. Square root of the spectral density spd of shape (m,), computed from the kernel hyperparameters and the j-th eigenvalue
   $\lambda_j = \left(\frac{j \pi}{2\ell}\right)^2$.

Part 2 — stochastic:

4. Spectral coefficients beta ~ Normal(0, 1) of shape (m,) (non-centered parameterization).

Output: phi @ (spd * beta), shape (n_timepoints,).

The only step that involves numpyro.sample is step 4.

Computational concern: K independent subpopulation trajectories (HierarchicalInfections)

The convenience wrappers hsgp_squared_exponential / hsgp_matern hardcode sample site names "basis" and "beta" inside numpyro.plate. To obtain independent trajectories for n_processes > 1 while using these wrappers, each call must be wrapped in numpyro.handlers.scope — a Python loop that is not JAX-compiled.

The vectorized alternative uses the lower-level helpers directly and samples beta as a single (n_processes, m) array:

from numpyro.contrib.hsgp.laplacian import eigenfunctions
from numpyro.contrib.hsgp.spectral_densities import (
    diag_spectral_density_squared_exponential,
    diag_spectral_density_matern,
)

phi = eigenfunctions(x, ell, m)          # (n_timepoints, m), pure JAX
spd = jnp.sqrt(
    diag_spectral_density_squared_exponential(alpha, length, ell, m, dim=1)
)                                         # (m,), pure JAX

with numpyro.plate(f"{name_prefix}_processes", n_processes):
    beta = numpyro.sample(
        f"{name_prefix}_beta",
        dist.Normal(jnp.zeros(m), jnp.ones(m)),
    )                                     # (n_processes, m)

trajectories = phi @ (spd * beta).T      # (n_timepoints, n_processes)

This is a single numpyro.sample call and one matrix multiply regardless of n_processes. The cost scales with n_processes * m inside JAX, not with n_processes Python iterations.

For n_processes=1 the plate collapses, beta has shape (m,), and trajectories has shape (n_timepoints,). Reshape to (n_timepoints, 1) for protocol compliance.

Design questions/decisions

Kernel Support

Support two kernels via a kernel string parameter:

kernel value	Spectral density helper	Extra parameter
"squared_exponential"	diag_spectral_density_squared_exponential	none
"matern"	diag_spectral_density_matern	nu (smoothness)

Rational quadratic is excluded from the initial implementation (its spectral density has a different functional form requiring additional care). Raise ValueError for unknown kernel strings in __init__.

Kernel choice for epidemiological use:

The squared exponential kernel produces infinitely smooth trajectories. This may over-smooth Rt dynamics where rapid changes occur (e.g., intervention effects). Matern $\nu = 3/2$ or $\nu = 5/2$ produces smoother-than-random-walk but finitely differentiable trajectories and is a better default for Rt modeling. (Note: Matern $\nu = 1/2$ (exponential kernel) is the continuous analogue of AR(1).)

Boundary Parameter ell

ell must satisfy ell > max(x) = n_timepoints - 1. Since n_timepoints is only known at sample() time, accept ell_factor: float in __init__ and compute ell = ell_factor * max(1, n_timepoints - 1) at sample time.

Validate ell_factor > 1.0 in __init__. Default: 1.5.

Hyperparameter Priors

alpha (amplitude) and length (length scale) are sampled inside sample().
Accept numpyro.distributions.Distribution objects as constructor parameters.

Defaults:

• alpha_prior = dist.HalfNormal(1.0) — amplitude on the log(Rt) scale
• length_prior = dist.HalfNormal(1.0) — length scale in days

For Matern, nu is a fixed float (not sampled). Common values: 0.5, 1.5, 2.5. Default: 1.5.

Hyperparameters are shared across all processes within one sample() call. This implements partial pooling: all subpopulations are assumed to have deviation trajectories with the same smoothness and amplitude structure, differing only in their specific realization. If independent hyperparameters per process are needed, instantiate and call HSGP separately per process.

initial_value as Additive Offset

HSGP has no sequential starting point. initial_value is interpreted as an additive offset (intercept) applied to the entire trajectory:

trajectories = phi @ (spd * beta).T + initial_value

initial_value=None defaults to 0.0. A scalar is broadcast over all processes. This is consistent with how HierarchicalInfections uses initial_log_rt for the baseline process and jnp.zeros(n_subpops) for deviations.

Non-Centered Parameterization

The non-centered parameterization (beta ~ Normal(0, 1), scaled by spd) is used unconditionally. It is strongly recommended for HMC and avoids funnel geometry. No centered option is provided in the initial implementation.

Class Interface

class HSGP(TemporalProcess):
    def __init__(
        self,
        m: int,
        ell_factor: float = 1.5,
        kernel: str = "squared_exponential",
        nu: float = 1.5,
        alpha_prior: dist.Distribution = dist.HalfNormal(1.0),
        length_prior: dist.Distribution = dist.HalfNormal(1.0),
    ) -> None: ...

    def __repr__(self) -> str: ...

    def sample(
        self,
        n_timepoints: int,
        initial_value: float | ArrayLike | None = None,
        n_processes: int = 1,
        name_prefix: str = "hsgp",
    ) -> ArrayLike: ...
        # Returns shape (n_timepoints, n_processes)

Validation in __init__:

• m >= 1
• ell_factor > 1.0
• kernel in {"squared_exponential", "matern"}
• nu > 0 (only meaningful for Matern, but validate unconditionally)

Validation in sample:

• n_timepoints >= 1
• n_processes >= 1

NumPyro Sample Site Names

For name_prefix="hsgp", n_processes=3, m=10:

Site	Shape	Description
hsgp_alpha	scalar	GP amplitude, sampled from alpha_prior
hsgp_length	scalar	GP length scale, sampled from length_prior
hsgp_processes	plate of size n_processes	process plate
hsgp_beta	(n_processes, m)	spectral coefficients, all processes

For n_processes=1, hsgp_beta has shape (m,) (plate collapses to scalar batch dimension).

Contrast with the loop-based approach in hsgp.md, which produced n_processes separate sites (hsgp_proc0::beta, hsgp_proc1::beta, ...) using numpyro.handlers.scope(prefix=..., divider="::"). The vectorized approach produces one site, which simplifies posterior analysis.

When used inside HierarchicalInfections, all HSGP sites are further prefixed by numpyro.handlers.scope(prefix=self.name, divider="::") (see hierarchical_infections.py line 262). With name="latent_infections" and name_prefix="subpop_deviations", the fully-qualified site names become:

Fully-qualified site	Description
latent_infections::subpop_deviations_alpha	GP amplitude
latent_infections::subpop_deviations_length	GP length scale
latent_infections::subpop_deviations_beta	spectral coefficients (n_subpops, m)

---

Use in HierarchicalInfections

HierarchicalInfections calls its temporal processes as:

log_rt_baseline = self.baseline_rt_process.sample(
    n_timepoints=n_total_days,
    initial_value=initial_log_rt,
    name_prefix="log_rt_baseline",
)

deviations_raw = self.subpop_rt_deviation_process.sample(
    n_timepoints=n_total_days,
    n_processes=pop.n_subpops,
    initial_value=jnp.zeros(pop.n_subpops),
    name_prefix="subpop_deviations",
)

HSGP satisfies the TemporalProcess protocol and can be passed as either argument without changes to HierarchicalInfections.

Baseline process (n_processes=1): initial_log_rt becomes the GP intercept rather than a sequential starting value. The interpretation is: the GP trajectory fluctuates around initial_log_rt. This is epidemiologically coherent and requires no special handling.

Deviation process (n_processes=n_subpops): All subpopulations share alpha and length but have independent coefficient vectors. The sum-to-zero constraint in HierarchicalInfections (line 207) is applied after sampling and works regardless of temporal process type.

Note on n_total_days: n_total_days includes the initialization period. HSGP trajectories during initialization are GP draws over the full time window, not extrapolations from a starting point. This is more principled than backward extrapolation from an AR initial condition but means the trajectory during initialization is not constrained to any particular shape.

[damonbayer]

This looks good @cdc-mitzimorris! However, in reviewing this I still found it hard to understand the distinctions between latent, temporal_processes, and process. I think the names and organization could be improved.

I would also be curious for @dylanhmorris's review.

[cdc-mitzimorris]

This looks good @cdc-mitzimorris! However, in reviewing this I still found it hard to understand the distinctions between latent, temporal_processes, and process. I think the names and organization could be improved.

happy to discuss further. temporal_process is only used by class latent/base.py and everything that inherits from it, therefore to my mind, it should be part of the latent submodule. Is there a description/rationale behind the current set of submodules and what belongs where? Why so many different submodules?