"""GFAModel class for Group Factor Analysis."""
from __future__ import annotations
import os
import sys
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
from latents.data import ObsStatic
from latents.gfa.config import GFAFitConfig
from latents.gfa.inference import fit, infer_latents, infer_loadings, init_posteriors
from latents.gfa.tracking import (
GFAFitFlags,
GFAFitTracker,
load_gfa_state,
save_gfa_state,
)
from latents.observation import (
ObsParamsHyperPrior,
ObsParamsPosterior,
ObsParamsPrior,
)
from latents.state import LatentsPosteriorStatic, LatentsPriorStatic
[docs]
class GFAModel:
"""High-level interface for Group Factor Analysis.
Parameters
----------
config : GFAFitConfig or None, default None
Fitting configuration. If None, uses default GFAFitConfig().
obs_hyperprior : ObsParamsHyperPrior or None, default None
Prior hyperparameters for observation model. If None, uses default.
Attributes
----------
config : GFAFitConfig
Fitting configuration (immutable after construction).
obs_posterior : ObsParamsPosterior or None
Posterior over observation model parameters. None until fit.
latents_posterior : LatentsPosteriorStatic or None
Posterior over latent variables. None until fit.
tracker : GFAFitTracker or None
Fitting progress tracker. None until fit.
flags : GFAFitFlags or None
Fitting status flags. None until fit.
Examples
--------
**Basic usage: fit and infer**
>>> from latents.gfa import GFAModel, GFAFitConfig
>>> from latents.callbacks import ProgressCallback
>>> config = GFAFitConfig(x_dim_init=10)
>>> model = GFAModel(config=config)
>>> model.fit(Y, callbacks=[ProgressCallback()])
>>> X_new = model.infer_latents(Y_new)
**Persistence: save and load**
>>> model.save("fitted_model.safetensors")
>>> loaded = GFAModel.load("fitted_model.safetensors")
"""
def __init__(
self,
config: GFAFitConfig | None = None,
obs_hyperprior: ObsParamsHyperPrior | None = None,
):
# Configuration (immutable)
self.config = config or GFAFitConfig()
# Prior specification (private, fixed at construction)
self._obs_prior = ObsParamsPrior(
hyperprior=obs_hyperprior or ObsParamsHyperPrior()
)
self._latents_prior = LatentsPriorStatic()
# Posterior estimates (populated by fit)
self.obs_posterior: ObsParamsPosterior | None = None
self.latents_posterior: LatentsPosteriorStatic | None = None
# Fitting state (populated by fit)
self.tracker: GFAFitTracker | None = None
self.flags: GFAFitFlags | None = None
def __repr__(self) -> str:
fitted = self.obs_posterior is not None
return f"GFAModel(fitted={fitted}, config={self.config})"
@property
def obs_hyperprior(self) -> ObsParamsHyperPrior:
"""Observation model hyperprior parameters."""
return self._obs_prior.hyperprior
@property
def obs_prior(self) -> ObsParamsPrior:
"""Observation model prior."""
return self._obs_prior
@property
def latents_prior(self) -> LatentsPriorStatic:
"""Latent variable prior."""
return self._latents_prior
[docs]
def fit(self, Y: ObsStatic, callbacks: list | None = None) -> Self:
"""Fit model to data via variational inference.
Resets tracker and flags. Warm-starts from posteriors if present,
otherwise initializes from scratch.
Parameters
----------
Y : ObsStatic
Observed data.
callbacks : list of Callback or None, default None
List of callback objects for progress, logging, checkpointing, etc.
See :mod:`~latents.callbacks` module.
Returns
-------
Self
The fitted model (for method chaining).
Examples
--------
>>> from latents.gfa import GFAModel, GFAFitConfig
>>> from latents.callbacks import ProgressCallback
>>> config = GFAFitConfig(x_dim_init=10, max_iter=100)
>>> model = GFAModel(config=config)
>>> model.fit(Y, callbacks=[ProgressCallback()])
"""
# Initialize posteriors if not present (cold start)
if self.obs_posterior is None:
self._init_posteriors(Y)
# Fit (resets tracker/flags for fresh convergence tracking)
self.obs_posterior, self.latents_posterior, self.tracker, self.flags = fit(
Y,
self.obs_posterior,
self.latents_posterior,
config=self.config,
obs_hyperprior=self.obs_hyperprior,
callbacks=callbacks,
)
return self
[docs]
def resume_fit(
self,
Y: ObsStatic,
max_iter: int | None = None,
callbacks: list | None = None,
) -> Self:
"""Resume an interrupted fit.
Appends to tracker, preserves convergence baseline.
Parameters
----------
Y : ObsStatic
Observed data.
max_iter : int or None, default None
Maximum iterations for this resume run. If None, uses config.max_iter.
callbacks : list of Callback or None, default None
List of callback objects for progress, logging, checkpointing, etc.
Returns
-------
Self
The model (for method chaining).
Raises
------
ValueError
If no fit to resume (posteriors not initialized).
"""
if self.obs_posterior is None:
msg = "No fit to resume. Use fit() instead."
raise ValueError(msg)
if self.tracker is None or self.flags is None:
msg = "No tracking state to resume. Use fit() instead."
raise ValueError(msg)
if self.flags.converged:
# Already converged, nothing to do
return self
self.obs_posterior, self.latents_posterior, self.tracker, self.flags = fit(
Y,
self.obs_posterior,
self.latents_posterior,
config=self.config,
obs_hyperprior=self.obs_hyperprior,
tracker=self.tracker,
flags=self.flags,
max_iter=max_iter,
callbacks=callbacks,
)
return self
[docs]
def clear_fit(self) -> Self:
"""Clear fit results for fresh initialization on next fit().
Returns
-------
Self
The model (for method chaining).
"""
self.obs_posterior = None
self.latents_posterior = None
self.tracker = None
self.flags = None
return self
[docs]
def infer_latents(self, Y: ObsStatic) -> LatentsPosteriorStatic:
"""Infer latent posterior for new data given fitted parameters.
Does not modify the model's stored latents_posterior.
Parameters
----------
Y : ObsStatic
Observed data.
Returns
-------
LatentsPosteriorStatic
Posterior over latent variables for the given data.
Raises
------
ValueError
If model has not been fitted.
Examples
--------
>>> model.fit(Y_train)
>>> X_test = model.infer_latents(Y_test)
>>> X_test.mean.shape
(5, 50)
"""
if self.obs_posterior is None:
msg = "Model must be fitted before inferring latents."
raise ValueError(msg)
return infer_latents(Y, self.obs_posterior)
[docs]
def recompute_latents(self, Y: ObsStatic) -> Self:
"""Recompute latents from data. Updates self.latents_posterior.
Use this to restore latents after loading a model saved with
save_x=False.
Parameters
----------
Y : ObsStatic
Observed data (typically the training data).
Returns
-------
Self
The model (for method chaining).
Raises
------
ValueError
If model has not been fitted.
"""
if self.obs_posterior is None:
msg = "Model must be fitted before recomputing latents."
raise ValueError(msg)
if self.latents_posterior is None:
self.latents_posterior = LatentsPosteriorStatic()
infer_latents(Y, self.obs_posterior, self.latents_posterior)
return self
[docs]
def recompute_loadings(self, Y: ObsStatic) -> Self:
"""Recompute loading posterior from data. Updates self.obs_posterior.C.
Use this to restore C.cov after loading a model saved with
save_c_cov=False.
Note: At convergence, the recomputed values are essentially identical
to the original fitted values. Pre-convergence, there may be non-negligible
differences because the reconstruction uses the final latents posterior rather
than the latents posterior from the previous iteration.
Parameters
----------
Y : ObsStatic
Observed data (typically the training data).
Returns
-------
Self
The model (for method chaining).
Raises
------
ValueError
If model has not been fitted, or if latents are not available.
"""
if self.obs_posterior is None:
msg = "Model must be fitted before recomputing loadings."
raise ValueError(msg)
if (
self.latents_posterior is None
or not self.latents_posterior.is_initialized()
):
msg = "Latents must be available. Call recompute_latents(Y) first."
raise ValueError(msg)
infer_loadings(Y, self.obs_posterior, self.latents_posterior)
return self
def _init_posteriors(self, Y: ObsStatic) -> None:
"""Initialize posteriors from data."""
self.obs_posterior, self.latents_posterior = init_posteriors(
Y, config=self.config, obs_hyperprior=self.obs_hyperprior
)
[docs]
def save(self, path: str | os.PathLike[str]) -> None:
"""Save model to a safetensors file.
Uses safetensors format for secure serialization (no arbitrary code
execution on load). Arrays are stored as tensors; scalars and config
are stored as JSON in metadata.
Parameters
----------
path : str or PathLike
Output file path (conventionally ends in .safetensors).
"""
save_gfa_state(
path,
config=self.config,
obs_hyperprior=self.obs_hyperprior,
obs_posterior=self.obs_posterior,
latents_posterior=self.latents_posterior,
tracker=self.tracker,
flags=self.flags,
)
[docs]
@classmethod
def load(cls, path: str | os.PathLike[str]) -> GFAModel:
"""Load model from a safetensors file.
Parameters
----------
path : str or PathLike
Path to .safetensors file.
Returns
-------
GFAModel
Loaded model, ready for inference or continued fitting.
"""
config, obs_hyperprior, obs_posterior, latents_posterior, tracker, flags = (
load_gfa_state(path)
)
# Create model with loaded config and hyperprior
model = cls(config=config, obs_hyperprior=obs_hyperprior)
# Restore posteriors and tracking state
model.obs_posterior = obs_posterior
model.latents_posterior = latents_posterior
model.tracker = tracker
model.flags = flags
return model
