"""Posterior distributions for observation model parameters."""
from __future__ import annotations
import sys
from itertools import combinations
import numpy as np
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
from latents._internal.numerics import stability_floor
from latents.base import ArrayContainer
from latents.observation.realizations import ObsParamsRealization
[docs]
class LoadingPosterior(ArrayContainer):
"""Posterior distribution over loading matrices.
Parameters
----------
mean : ndarray of float, shape (y_dim, x_dim) or None, default None
Posterior mean.
cov : ndarray of float, shape (y_dim, x_dim, x_dim) or None, default None
Posterior covariances.
moment : ndarray of float, shape (y_dim, x_dim, x_dim) or None, default None
Posterior second moments.
Attributes
----------
mean : ndarray of float, shape (y_dim, x_dim) or None
Posterior mean.
cov : ndarray of float, shape (y_dim, x_dim, x_dim) or None
Posterior covariances.
moment : ndarray of float, shape (y_dim, x_dim, x_dim) or None
Posterior second moments.
"""
def __init__(
self,
mean: np.ndarray | None = None,
cov: np.ndarray | None = None,
moment: np.ndarray | None = None,
):
if mean is not None and not isinstance(mean, np.ndarray):
msg = "mean must be a numpy.ndarray."
raise TypeError(msg)
self.mean = mean
if cov is not None and not isinstance(cov, np.ndarray):
msg = "cov must be a numpy.ndarray."
raise TypeError(msg)
self.cov = cov
if moment is not None and not isinstance(moment, np.ndarray):
msg = "moment must be a numpy.ndarray."
raise TypeError(msg)
self.moment = moment
[docs]
def get_groups(
self,
y_dims: np.ndarray,
) -> tuple[list[np.ndarray], list[np.ndarray], list[np.ndarray]]:
"""Get list of views into posterior parameters, one per group.
Parameters
----------
y_dims : ndarray of int, shape (n_groups,)
Dimensionalities of each observed group.
Returns
-------
group_means : list of ndarray or None
List of views into mean, one per group. None if mean is None.
group_covs : list of ndarray or None
List of views into cov, one per group. None if cov is None.
group_moments : list of ndarray or None
List of views into moment, one per group. None if moment is None.
"""
group_means = None
if self.mean is not None:
if np.sum(y_dims) != self.mean.shape[0]:
msg = "The sum of y_dims must equal the number of rows in mean."
raise ValueError(msg)
group_means = np.split(self.mean, np.cumsum(y_dims)[:-1], axis=0)
group_covs = None
if self.cov is not None:
if np.sum(y_dims) != self.cov.shape[0]:
msg = "The sum of y_dims must equal the size of the first axis of cov."
raise ValueError(msg)
group_covs = np.split(self.cov, np.cumsum(y_dims)[:-1], axis=0)
group_moments = None
if self.moment is not None:
if np.sum(y_dims) != self.moment.shape[0]:
msg = (
"The sum of y_dims must equal the size of the first axis of moment."
)
raise ValueError(msg)
group_moments = np.split(self.moment, np.cumsum(y_dims)[:-1], axis=0)
return group_means, group_covs, group_moments
[docs]
def compute_moment(self, in_place: bool = True) -> np.ndarray:
"""Compute posterior second moments E[C_i C_i^T] for each row i.
Parameters
----------
in_place : bool, default True
If True, store in self.moment and return reference.
If False, return new array.
Returns
-------
ndarray
Second moments, shape (y_dim, x_dim, x_dim).
"""
if in_place:
if self.moment is None:
self.moment = np.zeros_like(self.cov)
# Outer product per row:
# (y_dim, x_dim), (y_dim, x_dim) -> (y_dim, x_dim, x_dim)
np.einsum("ij,ik->ijk", self.mean, self.mean, out=self.moment)
self.moment += self.cov
return self.moment
return np.einsum("ij,ik->ijk", self.mean, self.mean) + self.cov
[docs]
def compute_squared_norms(self, y_dims: np.ndarray) -> np.ndarray:
"""Compute expected squared norm of each column per group.
Parameters
----------
y_dims : ndarray of int, shape (n_groups,)
Dimensionalities of each observed group.
Returns
-------
ndarray
Squared norms, shape (n_groups, x_dim).
"""
n_groups = len(y_dims)
x_dim = self.moment.shape[1]
_, _, moments = self.get_groups(y_dims)
squared_norms = np.zeros((n_groups, x_dim))
for group_idx in range(n_groups):
squared_norms[group_idx, :] = np.sum(
moments[group_idx].diagonal(offset=0, axis1=1, axis2=2), axis=0
)
return squared_norms
[docs]
def get_subset_dims(
self,
x_indices: np.ndarray,
in_place: bool = True,
) -> Self:
"""Keep only specified latent dimensions.
Parameters
----------
x_indices : ndarray of int
Indices of latent dimensions to keep.
in_place : bool, default True
If True, modify self. If False, return new instance.
Returns
-------
Self
Modified or new instance.
"""
new_mean = self.mean[:, x_indices] if self.mean is not None else None
new_cov = (
self.cov[np.ix_(np.arange(self.cov.shape[0]), x_indices, x_indices)]
if self.cov is not None
else None
)
new_moment = (
self.moment[np.ix_(np.arange(self.moment.shape[0]), x_indices, x_indices)]
if self.moment is not None
else None
)
if in_place:
self.mean = new_mean
self.cov = new_cov
self.moment = new_moment
return self
return self.__class__(mean=new_mean, cov=new_cov, moment=new_moment)
[docs]
class ARDPosterior(ArrayContainer):
"""Posterior distribution over ARD parameters.
Parameters
----------
a : ndarray of float, shape (n_groups,) or None, default None
Shape parameters.
b : ndarray of float, shape (n_groups, x_dim) or None, default None
Rate parameters.
mean : ndarray of float, shape (n_groups, x_dim) or None, default None
Posterior mean a/b.
Attributes
----------
a : ndarray of float, shape (n_groups,) or None
Shape parameters.
b : ndarray of float, shape (n_groups, x_dim) or None
Rate parameters.
mean : ndarray of float, shape (n_groups, x_dim) or None
Posterior mean a/b.
"""
def __init__(
self,
a: np.ndarray | None = None,
b: np.ndarray | None = None,
mean: np.ndarray | None = None,
):
if a is not None and not isinstance(a, np.ndarray):
msg = "a must be a numpy.ndarray."
raise TypeError(msg)
self.a = a
if b is not None and not isinstance(b, np.ndarray):
msg = "b must be a numpy.ndarray."
raise TypeError(msg)
self.b = b
if mean is not None and not isinstance(mean, np.ndarray):
msg = "mean must be a numpy.ndarray."
raise TypeError(msg)
self.mean = mean
[docs]
def compute_mean(self, in_place: bool = True) -> np.ndarray:
"""Compute posterior mean a/b.
Parameters
----------
in_place : bool, default True
If True, store in self.mean and return reference.
If False, return new array.
Returns
-------
ndarray
Posterior mean, shape (n_groups, x_dim).
"""
floor = stability_floor(self.b.dtype)
if in_place:
if self.mean is None:
self.mean = np.zeros_like(self.b)
# a: (n_groups,) -> (n_groups, 1) for broadcast with b: (n_groups, x_dim)
self.mean[:] = self.a[:, np.newaxis] / np.maximum(self.b, floor)
return self.mean
return self.a[:, np.newaxis] / np.maximum(self.b, floor)
[docs]
def get_subset_dims(
self,
x_indices: np.ndarray,
in_place: bool = True,
) -> Self:
"""Keep only specified latent dimensions.
Parameters
----------
x_indices : ndarray of int
Indices of latent dimensions to keep.
in_place : bool, default True
If True, modify self. If False, return new instance.
Returns
-------
Self
Modified or new instance.
"""
new_b = self.b[:, x_indices] if self.b is not None else None
new_mean = self.mean[:, x_indices] if self.mean is not None else None
if in_place:
self.b = new_b
self.mean = new_mean
return self
new_a = self.a.copy() if self.a is not None else None
return self.__class__(a=new_a, b=new_b, mean=new_mean)
[docs]
class ObsMeanPosterior(ArrayContainer):
"""Posterior distribution over observation mean parameters.
Parameters
----------
mean : ndarray of float, shape (y_dim,) or None, default None
Posterior mean.
cov : ndarray of float, shape (y_dim,) or None, default None
Posterior variance (diagonal).
Attributes
----------
mean : ndarray of float, shape (y_dim,) or None
Posterior mean.
cov : ndarray of float, shape (y_dim,) or None
Posterior variance (diagonal).
"""
def __init__(
self,
mean: np.ndarray | None = None,
cov: np.ndarray | None = None,
):
if mean is not None and not isinstance(mean, np.ndarray):
msg = "mean must be a numpy.ndarray."
raise TypeError(msg)
self.mean = mean
if cov is not None and not isinstance(cov, np.ndarray):
msg = "cov must be a numpy.ndarray."
raise TypeError(msg)
self.cov = cov
[docs]
def get_groups(
self,
y_dims: np.ndarray,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
"""Get list of views into posterior parameters, one per group.
Parameters
----------
y_dims : ndarray of int, shape (n_groups,)
Dimensionalities of each observed group.
Returns
-------
group_means : list of ndarray or None
List of views into mean, one per group. None if mean is None.
group_covs : list of ndarray or None
List of views into cov, one per group. None if cov is None.
"""
group_means = None
if self.mean is not None:
if np.sum(y_dims) != len(self.mean):
msg = "The sum of y_dims must equal the length of mean."
raise ValueError(msg)
group_means = np.split(self.mean, np.cumsum(y_dims)[:-1], axis=0)
group_covs = None
if self.cov is not None:
if np.sum(y_dims) != len(self.cov):
msg = "The sum of y_dims must equal the length of cov."
raise ValueError(msg)
group_covs = np.split(self.cov, np.cumsum(y_dims)[:-1], axis=0)
return group_means, group_covs
[docs]
class ObsPrecPosterior(ArrayContainer):
"""Posterior distribution over observation precision parameters.
Parameters
----------
a : float or None, default None
Shape parameter (scalar, shared across dimensions).
b : ndarray of float, shape (y_dim,) or None, default None
Rate parameters.
mean : ndarray of float, shape (y_dim,) or None, default None
Posterior mean a/b.
Attributes
----------
a : float or None
Shape parameter (scalar, shared across dimensions).
b : ndarray of float, shape (y_dim,) or None
Rate parameters.
mean : ndarray of float, shape (y_dim,) or None
Posterior mean a/b.
"""
def __init__(
self,
a: float | None = None,
b: np.ndarray | None = None,
mean: np.ndarray | None = None,
):
if a is not None and not isinstance(a, float):
msg = "a must be a float."
raise TypeError(msg)
self.a = a
if b is not None and not isinstance(b, np.ndarray):
msg = "b must be a numpy.ndarray."
raise TypeError(msg)
self.b = b
if mean is not None and not isinstance(mean, np.ndarray):
msg = "mean must be a numpy.ndarray."
raise TypeError(msg)
self.mean = mean
[docs]
def get_groups(
self,
y_dims: np.ndarray,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
"""Get list of views into posterior parameters, one per group.
Parameters
----------
y_dims : ndarray of int, shape (n_groups,)
Dimensionalities of each observed group.
Returns
-------
group_means : list of ndarray or None
List of views into mean, one per group. None if mean is None.
group_bs : list of ndarray or None
List of views into b, one per group. None if b is None.
"""
group_means = None
if self.mean is not None:
if np.sum(y_dims) != len(self.mean):
msg = "The sum of y_dims must equal the length of mean."
raise ValueError(msg)
group_means = np.split(self.mean, np.cumsum(y_dims)[:-1], axis=0)
group_bs = None
if self.b is not None:
if np.sum(y_dims) != len(self.b):
msg = "The sum of y_dims must equal the length of b."
raise ValueError(msg)
group_bs = np.split(self.b, np.cumsum(y_dims)[:-1], axis=0)
return group_means, group_bs
[docs]
def compute_mean(self, in_place: bool = True) -> np.ndarray:
"""Compute posterior mean a/b.
Parameters
----------
in_place : bool, default True
If True, store in self.mean and return reference.
If False, return new array.
Returns
-------
ndarray
Posterior mean, shape (y_dim,).
"""
floor = stability_floor(self.b.dtype)
if in_place:
if self.mean is None:
self.mean = np.zeros_like(self.b)
self.mean[:] = self.a / np.maximum(self.b, floor)
return self.mean
return self.a / np.maximum(self.b, floor)
[docs]
class ObsParamsPosterior:
"""Posterior distributions over all observation model parameters.
Bundle of posteriors q(C), q(alpha), q(d), q(phi) with methods for
sampling, computing point estimates, and analysis.
Parameters
----------
x_dim : int or None, default None
Number of latent dimensions.
y_dims : ndarray of int, shape (n_groups,) or None, default None
Dimensionalities of each observed group.
C : LoadingPosterior or None, default None
Posterior over loading matrices.
alpha : ARDPosterior or None, default None
Posterior over ARD parameters.
d : ObsMeanPosterior or None, default None
Posterior over observation means.
phi : ObsPrecPosterior or None, default None
Posterior over observation precisions.
Attributes
----------
x_dim : int or None
Number of latent dimensions.
y_dims : ndarray of int, shape (n_groups,) or None
Dimensionalities of each observed group.
C : LoadingPosterior
Posterior over loading matrices.
alpha : ARDPosterior
Posterior over ARD parameters.
d : ObsMeanPosterior
Posterior over observation means.
phi : ObsPrecPosterior
Posterior over observation precisions.
Examples
--------
After fitting a GFA model, access the observation posterior:
>>> model = GFAModel()
>>> model.fit(Y)
>>> obs_post = model.obs_posterior
Get posterior mean as a realization:
>>> params = obs_post.posterior_mean
>>> params.C.shape
(20, 5)
Sample from the posterior:
>>> rng = np.random.default_rng(42)
>>> sample = obs_post.sample(rng)
"""
def __init__(
self,
x_dim: int | None = None,
y_dims: np.ndarray | None = None,
C: LoadingPosterior | None = None,
alpha: ARDPosterior | None = None,
d: ObsMeanPosterior | None = None,
phi: ObsPrecPosterior | None = None,
):
if x_dim is not None and not isinstance(x_dim, int):
msg = "x_dim must be an integer."
raise TypeError(msg)
self.x_dim = x_dim
if y_dims is not None and not isinstance(y_dims, np.ndarray):
msg = "y_dims must be a numpy.ndarray of integers."
raise TypeError(msg)
self.y_dims = y_dims
if C is None:
self.C = LoadingPosterior()
elif not isinstance(C, LoadingPosterior):
msg = "C must be a LoadingPosterior object."
raise TypeError(msg)
else:
self.C = C
if alpha is None:
self.alpha = ARDPosterior()
elif not isinstance(alpha, ARDPosterior):
msg = "alpha must be an ARDPosterior object."
raise TypeError(msg)
else:
self.alpha = alpha
if d is None:
self.d = ObsMeanPosterior()
elif not isinstance(d, ObsMeanPosterior):
msg = "d must be an ObsMeanPosterior object."
raise TypeError(msg)
else:
self.d = d
if phi is None:
self.phi = ObsPrecPosterior()
elif not isinstance(phi, ObsPrecPosterior):
msg = "phi must be an ObsPrecPosterior object."
raise TypeError(msg)
else:
self.phi = phi
def __repr__(self) -> str:
return (
f"{type(self).__name__}(x_dim={self.x_dim}, "
f"y_dims={self.y_dims}, "
f"C={self.C}, "
f"alpha={self.alpha}, "
f"d={self.d}, "
f"phi={self.phi})"
)
@property
def posterior_mean(self) -> ObsParamsRealization:
"""Return posterior means as a realization.
Returns
-------
ObsParamsRealization
Realization with posterior mean values.
"""
return ObsParamsRealization(
C=self.C.mean.copy(),
d=self.d.mean.copy(),
phi=self.phi.mean.copy(),
alpha=self.alpha.mean.copy(),
y_dims=self.y_dims.copy(),
x_dim=self.x_dim,
)
[docs]
def sample(self, rng: np.random.Generator) -> ObsParamsRealization:
"""Draw a sample from the posterior distributions.
Samples from q(C), q(alpha), q(d), q(phi) independently.
Parameters
----------
rng : numpy.random.Generator
Random number generator.
Returns
-------
ObsParamsRealization
Sampled parameter values.
"""
y_dim = int(self.y_dims.sum())
n_groups = len(self.y_dims)
# Sample C: each row independently from N(mean_i, cov_i)
C_sample = np.zeros_like(self.C.mean)
for i in range(y_dim):
C_sample[i, :] = rng.multivariate_normal(self.C.mean[i, :], self.C.cov[i])
# Sample alpha: Gamma(a, b) for each element
# a has shape (n_groups,), b has shape (n_groups, x_dim)
alpha_sample = np.zeros_like(self.alpha.mean)
for g in range(n_groups):
alpha_sample[g, :] = rng.gamma(
shape=self.alpha.a[g], scale=1 / self.alpha.b[g, :]
)
# Sample d: N(mean, cov) element-wise (diagonal covariance)
d_sample = rng.normal(self.d.mean, np.sqrt(self.d.cov))
# Sample phi: Gamma(a, b) for each element
phi_sample = rng.gamma(shape=self.phi.a, scale=1 / self.phi.b)
return ObsParamsRealization(
C=C_sample,
d=d_sample,
phi=phi_sample,
alpha=alpha_sample,
y_dims=self.y_dims.copy(),
x_dim=self.x_dim,
)
[docs]
def is_initialized(self) -> bool:
"""Check if all posterior parameters have been initialized.
Returns
-------
bool
True if all required arrays are non-None.
"""
return (
self.x_dim is not None
and self.y_dims is not None
and self.C.mean is not None
and self.C.moment is not None
and self.alpha.a is not None
and self.alpha.b is not None
and self.alpha.mean is not None
and self.d.mean is not None
and self.d.cov is not None
and self.phi.a is not None
and self.phi.b is not None
and self.phi.mean is not None
)
[docs]
def get_subset_dims(
self,
x_indices: np.ndarray,
in_place: bool = True,
) -> Self:
"""Keep only specified latent dimensions.
Parameters
----------
x_indices : ndarray of int
Indices of latent dimensions to keep.
in_place : bool, default True
If True, modify self. If False, return new instance.
Returns
-------
Self
Modified or new instance.
"""
if in_place:
self.x_dim = len(x_indices)
self.C.get_subset_dims(x_indices, in_place=True)
self.alpha.get_subset_dims(x_indices, in_place=True)
return self
return self.__class__(
x_dim=len(x_indices),
y_dims=self.y_dims.copy(),
C=self.C.get_subset_dims(x_indices, in_place=False),
alpha=self.alpha.get_subset_dims(x_indices, in_place=False),
d=self.d.copy(),
phi=self.phi.copy(),
)
[docs]
def copy(self) -> Self:
"""Return a deep copy.
Returns
-------
Self
Deep copy of this instance.
"""
return self.__class__(
x_dim=self.x_dim,
y_dims=self.y_dims.copy(),
C=self.C.copy(),
alpha=self.alpha.copy(),
d=self.d.copy(),
phi=self.phi.copy(),
)
[docs]
def compute_snr(self, y_dims: np.ndarray | None = None) -> np.ndarray:
"""Compute signal-to-noise ratio for each group.
Parameters
----------
y_dims : ndarray or None, default None
Dimensionalities of observed groups. If None, uses self.y_dims.
Returns
-------
ndarray
SNR for each group, shape (n_groups,).
"""
if y_dims is None:
y_dims = self.y_dims
_, _, C_moments = self.C.get_groups(y_dims)
phi_means, _ = self.phi.get_groups(y_dims)
return np.array(
[
np.trace(np.sum(C_moments[group_idx], axis=0))
/ np.sum(1 / phi_means[group_idx])
for group_idx in range(len(y_dims))
]
)
[docs]
@staticmethod
def get_dim_types(n_groups: int) -> np.ndarray:
"""Generate all dimension types for n_groups.
Parameters
----------
n_groups : int
Number of observed groups.
Returns
-------
ndarray
Boolean array, shape (n_groups, 2^n_groups). Column j indicates
which groups are involved in dimension type j.
"""
n_dim_types = 2**n_groups
dim_types = np.empty((n_groups, n_dim_types))
for dim_idx in range(n_dim_types):
dim_str = format(dim_idx, f"0{n_groups}b")
dim_types[:, dim_idx] = np.array([int(b) for b in dim_str], dtype=bool)
return dim_types
[docs]
def compute_dimensionalities(
self,
cutoff_shared_var: float = 0.02,
cutoff_snr: float = 0.001,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Compute dimensionalities and variance explained by dimension type.
Parameters
----------
cutoff_shared_var : float, default 0.02
Minimum fraction of shared variance for significance.
cutoff_snr : float, default 0.001
Minimum SNR for any latents to be significant.
Returns
-------
num_dim : ndarray of int, shape (n_dim_types,)
Number of each dimension type.
sig_dims : ndarray of bool, shape (n_groups, x_dim)
Significant dimensions.
var_exp : ndarray of float, shape (n_groups, n_dim_types)
Variance explained by type.
dim_types : ndarray of bool, shape (n_groups, n_dim_types)
Dimension type indicators.
"""
n_groups = len(self.y_dims)
dim_types = self.get_dim_types(n_groups)
n_dim_types = dim_types.shape[1]
snr = self.compute_snr()
alpha_inv = 1 / self.alpha.mean
alpha_inv_rel = alpha_inv / np.sum(alpha_inv, axis=1, keepdims=True)
# snr: (n_groups,) -> (n_groups, 1) for broadcast with
# alpha_inv_rel: (n_groups, x_dim)
sig_dims = (alpha_inv_rel > cutoff_shared_var) & (snr > cutoff_snr)[
:, np.newaxis
] # (n_groups, x_dim)
num_dim = np.zeros(n_dim_types)
var_exp = np.zeros((n_groups, n_dim_types))
for dim_idx in range(n_dim_types):
# dim_types[:, dim_idx]: (n_groups,) -> (n_groups, 1) for broadcast with
# sig_dims
dims = np.all(sig_dims == dim_types[:, dim_idx, np.newaxis], axis=0)
num_dim[dim_idx] = np.sum(dims)
var_exp[:, dim_idx] = np.sum(alpha_inv_rel[:, dims], axis=1)
return num_dim, sig_dims, var_exp, dim_types
[docs]
@staticmethod
def compute_dims_pairs(
num_dim: np.ndarray,
dim_types: np.ndarray,
var_exp: np.ndarray,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Analyze shared dimensionalities between pairs of groups.
Parameters
----------
num_dim : ndarray
Number of each dimension type, shape (n_dim_types,).
dim_types : ndarray
Dimension type indicators, shape (n_groups, n_dim_types).
var_exp : ndarray
Variance explained by type, shape (n_groups, n_dim_types).
Returns
-------
pair_dims : ndarray of int, shape (n_pairs, 3)
Dimensionalities per pair.
pair_var_exp : ndarray of float, shape (n_pairs, 2)
Variance explained per pair.
pairs : ndarray of int, shape (n_pairs, 2)
Pair indices.
"""
n_groups = dim_types.shape[0]
group_idxs = [np.nonzero(dim_types[g, :])[0] for g in range(n_groups)]
pairs = list(combinations(range(n_groups), 2))
num_pairs = len(pairs)
pair_dims = np.zeros((num_pairs, 3), dtype=int)
pair_var_exp = np.zeros((num_pairs, 2))
for pair_idx, pair in enumerate(pairs):
pair_dims[pair_idx, 0] = num_dim[group_idxs[pair[0]]].sum()
pair_dims[pair_idx, 2] = num_dim[group_idxs[pair[1]]].sum()
shared_idxs = np.intersect1d(group_idxs[pair[0]], group_idxs[pair[1]])
pair_dims[pair_idx, 1] = num_dim[shared_idxs].sum()
pair_var_exp[pair_idx, 0] = var_exp[pair[0], shared_idxs].sum()
pair_var_exp[pair_idx, 1] = var_exp[pair[1], shared_idxs].sum()
return pair_dims, pair_var_exp, np.array(pairs)
