"""Hyperpriors and prior distributions for observation model parameters."""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from latents.observation.realizations import ObsParamsRealization
[docs]
@dataclass(frozen=True, slots=True, kw_only=True)
class ObsParamsHyperPrior:
"""Homogeneous hyperpriors for observation model parameters.
Scalar values broadcast to all groups/latents. Typical for inference
with uninformative priors.
Parameters
----------
a_alpha : float, default 1e-12
Shape parameter of the ARD prior (Gamma). Must be > 0.
b_alpha : float, default 1e-12
Rate parameter of the ARD prior (Gamma). Must be > 0.
a_phi : float, default 1e-12
Shape parameter of the observation precision prior (Gamma). Must be > 0.
b_phi : float, default 1e-12
Rate parameter of the observation precision prior (Gamma). Must be > 0.
beta_d : float, default 1e-12
Precision of the observation mean prior (Gaussian). Must be > 0.
Examples
--------
>>> priors = ObsParamsHyperPrior() # Use defaults (uninformative)
>>> priors = ObsParamsHyperPrior(a_alpha=1e-6, b_alpha=1e-6)
"""
a_alpha: float = 1e-12
b_alpha: float = 1e-12
a_phi: float = 1e-12
b_phi: float = 1e-12
beta_d: float = 1e-12
def __post_init__(self) -> None:
"""Validate all parameters are positive."""
for name in ("a_alpha", "b_alpha", "a_phi", "b_phi", "beta_d"):
value = getattr(self, name)
if not isinstance(value, (int, float)) or value <= 0:
msg = f"{name} must be a positive number, got {value!r}"
raise ValueError(msg)
[docs]
@dataclass(frozen=True, slots=True, kw_only=True)
class ObsParamsHyperPriorStructured:
"""Structured hyperpriors with per-group, per-latent control.
Enables sparsity constraints (np.inf in a_alpha forces zero loadings)
and incorporation of prior knowledge.
Parameters
----------
a_alpha : ndarray of float, shape (n_groups, x_dim)
Shape parameters for ARD priors.
Use np.inf to force zero loadings (sparsity pattern).
b_alpha : ndarray of float, shape (n_groups, x_dim)
Rate parameters for ARD priors.
Typically ones or matched to a_alpha.
a_phi : float, default 1.0
Shape parameter of observation precision prior. Must be > 0.
b_phi : float, default 1.0
Rate parameter of observation precision prior. Must be > 0.
beta_d : float, default 1.0
Precision of observation mean prior. Must be > 0.
Examples
--------
>>> # 3 groups, 4 latents, with sparsity pattern
>>> sparsity = np.array([
... [1, 1, np.inf, 1], # Group 0: latents 0,1,3
... [1, np.inf, 1, 1], # Group 1: latents 0,2,3
... [np.inf, 1, 1, 1], # Group 2: latents 1,2,3
... ])
>>> priors = ObsParamsHyperPriorStructured(
... a_alpha=100 * sparsity,
... b_alpha=100 * np.ones((3, 4)),
... )
"""
a_alpha: np.ndarray
b_alpha: np.ndarray
a_phi: float = 1.0
b_phi: float = 1.0
beta_d: float = 1.0
def __post_init__(self) -> None:
"""Validate array shapes and scalar positivity."""
# Validate a_alpha is ndarray
if not isinstance(self.a_alpha, np.ndarray):
msg = f"a_alpha must be a numpy array, got {type(self.a_alpha).__name__}"
raise TypeError(msg)
# Validate b_alpha is ndarray
if not isinstance(self.b_alpha, np.ndarray):
msg = f"b_alpha must be a numpy array, got {type(self.b_alpha).__name__}"
raise TypeError(msg)
# Validate shapes match
if self.a_alpha.shape != self.b_alpha.shape:
msg = (
f"a_alpha shape {self.a_alpha.shape} "
f"must match b_alpha shape {self.b_alpha.shape}"
)
raise ValueError(msg)
# Validate 2D
if self.a_alpha.ndim != 2:
msg = f"a_alpha must be 2D (n_groups, x_dim), got {self.a_alpha.ndim}D"
raise ValueError(msg)
# Validate b_alpha values are positive (a_alpha can have np.inf)
if np.any(self.b_alpha <= 0):
msg = "b_alpha values must all be > 0"
raise ValueError(msg)
# Validate finite a_alpha values are positive
finite_mask = np.isfinite(self.a_alpha)
if np.any(self.a_alpha[finite_mask] <= 0):
msg = "Finite a_alpha values must be > 0 (use np.inf for sparsity)"
raise ValueError(msg)
# Validate scalars
for name in ("a_phi", "b_phi", "beta_d"):
value = getattr(self, name)
if value <= 0:
msg = f"{name} must be > 0, got {value}"
raise ValueError(msg)
@property
def n_groups(self) -> int:
"""Number of observed groups."""
return self.a_alpha.shape[0]
@property
def x_dim(self) -> int:
"""Number of latent dimensions."""
return self.a_alpha.shape[1]
[docs]
@dataclass
class ObsParamsPrior:
"""Prior distributions over observation model parameters.
Encapsulates p(C, d, phi, alpha) and handles correct sampling order
(alpha must be sampled before C, since C|alpha ~ N(0, alpha^-1)).
Parameters
----------
hyperprior : ObsParamsHyperPrior or ObsParamsHyperPriorStructured
Hyperprior parameters controlling the prior distributions.
"""
hyperprior: ObsParamsHyperPrior | ObsParamsHyperPriorStructured
[docs]
def sample(
self,
y_dims: np.ndarray,
x_dim: int,
rng: np.random.Generator,
) -> ObsParamsRealization:
"""Sample from joint prior p(alpha)p(C|alpha)p(d)p(phi).
Parameters
----------
y_dims : ndarray of int, shape (n_groups,)
Dimensionalities of each observed group.
x_dim : int
Number of latent dimensions.
rng : numpy.random.Generator
Random number generator.
Returns
-------
ObsParamsRealization
Sampled parameter values.
"""
n_groups = len(y_dims)
y_dim = int(y_dims.sum())
# Get hyperprior values, expanding scalars to arrays if needed
if isinstance(self.hyperprior, ObsParamsHyperPriorStructured):
a_alpha = self.hyperprior.a_alpha
b_alpha = self.hyperprior.b_alpha
else:
a_alpha = np.full((n_groups, x_dim), self.hyperprior.a_alpha)
b_alpha = np.full((n_groups, x_dim), self.hyperprior.b_alpha)
# Sample observation mean: d ~ N(0, 1/beta_d)
d = rng.normal(0, 1 / np.sqrt(self.hyperprior.beta_d), size=y_dim)
# Sample observation precision: phi ~ Gamma(a_phi, b_phi)
phi = rng.gamma(
shape=self.hyperprior.a_phi,
scale=1 / self.hyperprior.b_phi,
size=y_dim,
)
# Sample ARD parameters and loadings group by group
alpha = np.zeros((n_groups, x_dim))
C = np.zeros((y_dim, x_dim))
# Split C by group for in-place assignment
y_boundaries = np.cumsum(y_dims)[:-1]
C_split = np.split(C, y_boundaries, axis=0)
for group_idx in range(n_groups):
for x_idx in range(x_dim):
a = a_alpha[group_idx, x_idx]
b = b_alpha[group_idx, x_idx]
if np.isinf(a):
# Infinite shape parameter forces zero loadings
alpha[group_idx, x_idx] = np.inf
C_split[group_idx][:, x_idx] = 0.0
else:
# Sample alpha ~ Gamma(a, b)
alpha[group_idx, x_idx] = rng.gamma(shape=a, scale=1 / b)
# Sample C|alpha ~ N(0, 1/alpha)
C_split[group_idx][:, x_idx] = rng.normal(
0,
1 / np.sqrt(alpha[group_idx, x_idx]),
size=y_dims[group_idx],
)
return ObsParamsRealization(
C=C,
d=d,
phi=phi,
alpha=alpha,
y_dims=y_dims.copy(),
x_dim=x_dim,
)
