cca_zoo.probabilistic

Probabilistic CCA via MCMC, black-box variational inference, or closed-form coordinate-ascent variational Bayes. GFA has no extra dependencies; ProbabilisticCCA and VariationalBayesCCA require pip install cca-zoo[probabilistic].

---

GFA

GFA(
    n_components: int = 1,
    *,
    center: bool = True,
    max_iter: int = 10000,
    tol: float = 0.0001,
    drop_k: bool = True,
    n_posterior_samples: int = 1000,
    random_state: int | None = None,
)

Bases: BaseProbabilistic

Group Factor Analysis: Bayesian CCA with per-view ARD.

A port of GFA() from the R package CCAGFA. A shared latent \(z\) generates every view, with an ARD precision per view and dimension:

\[ \begin{aligned} \alpha_{i,k} &\sim \mathrm{Gamma}(a_0, b_0), & W_i[:, k] &\sim \mathcal{N}(0, \alpha_{i,k}^{-1} I), \\ z &\sim \mathcal{N}(0, I), & x_i \mid z &\sim \mathcal{N}(W_i z, \tau_i^{-1} I), \end{aligned} \]

so a dimension is shared when \(\alpha_{i,k}\) is small in several views and private when small in one. Inference is closed-form coordinate-ascent variational Bayes. n_components is an upper bound: dimensions with vanishing loadings are pruned, leaving n_components_. Convergence is judged on the change in \(z\) rather than the lower bound, so raise max_iter if n_components_ is larger than expected.

Parameters:

Name	Type	Description	Default
n_components	int	Upper bound on the number of latent dimensions. Default is 1.	1
center	bool	Whether to centre each view. Default is True.	True
max_iter	int	Maximum coordinate-ascent iterations. Default is 10000.	10000
tol	float	Tolerance on the relative change in \(z\), held for 1000 iterations. Default is 1e-4.	0.0001
drop_k	bool	Whether to prune unused dimensions, as CCAGFA's dropK. Default is True.	True
n_posterior_samples	int	Draws from the fitted posterior. Default is 1000.	1000
random_state	int | None	Seed for the initialisation. Default is None.	None

Attributes:

Name	Type	Description
weights_	list[ndarray]	Posterior mean loadings of each view, shape (n_features_i, n_components_).
ard_precision_	list[ndarray]	Posterior mean ARD precisions, shape (n_views, n_components_); large means shrunk away.
n_components_	list[ndarray]	Number of dimensions kept.
posterior_samples_	dict[str, Any]	Posterior draws keyed W_{i}, noise_sd_{i} and alpha.
n_iter_	dict[str, Any]	Number of iterations run.

References

Klami, A., Virtanen, S., Leppäaho, E., & Kaski, S. (2015). Group factor analysis. IEEE Transactions on Neural Networks and Learning Systems, 26(9), 2136-2147. Klami, A., Virtanen, S., & Kaski, S. (2013). Bayesian Canonical Correlation Analysis. Journal of Machine Learning Research, 14, 965-1003.

Examples:

>>> import numpy as np
>>> from cca_zoo.probabilistic import GFA
>>> rng = np.random.default_rng(0)
>>> X1 = rng.standard_normal((50, 4))
>>> X2 = rng.standard_normal((50, 3))
>>> model = GFA(n_components=2, max_iter=50).fit([X1, X2])

Source code in cca_zoo/probabilistic/_gfa.py

def __init__(
    self,
    n_components: int = 1,
    *,
    center: bool = True,
    max_iter: int = 10000,
    tol: float = 1e-4,
    drop_k: bool = True,
    n_posterior_samples: int = 1000,
    random_state: int | None = None,
) -> None:
    super().__init__(n_components=n_components, center=center)
    self.max_iter = max_iter
    self.tol = tol
    self.drop_k = drop_k
    self.n_posterior_samples = n_posterior_samples
    self.random_state = random_state

fit

fit(views: list[ArrayLike], y: None = None) -> GFA

Fit the model.

Parameters:

Name	Type	Description	Default
views	list[ArrayLike]	Arrays of shape (n_samples, n_features_i), one per view.	required
y	None	Ignored.	None

Returns:

Type	Description
GFA	self.

Source code in cca_zoo/probabilistic/_gfa.py

def fit(self, views: list[ArrayLike], y: None = None) -> GFA:
    """Fit the model.

    Args:
        views: Arrays of shape (n_samples, n_features_i), one per view.
        y: Ignored.

    Returns:
        self.
    """
    views_ = self._setup_fit(views)
    rng = np.random.default_rng(self.random_state)
    q = _VariationalPosterior.initialise(views_, self.n_components, rng)
    # Convergence needs `_PATIENCE` consecutive small changes in Z and no
    # pruning: the change can dip below tol for hundreds of iterations
    # while one dimension's ARD decay waits on another's.
    previous: np.ndarray | None = None
    stable, converged, self.n_iter_ = 0, False, self.max_iter
    for iteration in range(self.max_iter):
        q.update_loadings(views_)
        q.update_factors(views_)
        q.update_ard()
        q.update_noise(views_)
        if self.drop_k and q.prune():
            stable = 0
        elif previous is not None and previous.shape == q.z.shape:
            change = np.linalg.norm(q.z - previous) / np.linalg.norm(previous)
            stable = stable + 1 if change < self.tol else 0
        previous = q.z.copy()
        if stable >= _PATIENCE:
            converged, self.n_iter_ = True, iteration + 1
            break
    warn_if_not_converged(self, converged)

    self.weights_: list[np.ndarray] = q.w
    self.ard_precision_: np.ndarray = np.array(q.alpha)
    self.posterior_samples_ = q.sample(rng, self.n_posterior_samples)
    self._fit_maps_and_importances(views_)
    return self

---

ProbabilisticCCA

ProbabilisticCCA(
    n_components: int = 1,
    *,
    center: bool = True,
    n_warmup: int = 500,
    n_posterior_samples: int = 1000,
    random_state: int | None = None,
)

Bases: BaseProbabilistic

Probabilistic CCA with posterior sampling by NUTS.

\[ z \sim \mathcal{N}(0, I), \qquad x_i \mid z \sim \mathcal{N}(W_i z + \mu_i, \Psi_i), \]

with each \(\Psi_i\) a full covariance, so that \(z\) models only what the views share: at the maximum likelihood each view's posterior mean spans its canonical variates (Bach & Jordan, 2005). \(\Psi_i\) has an LKJ prior on its correlation and log-normal scales, and \(z\) is integrated out. The likelihood is invariant to a shared rotation of the loadings, so the draws are aligned by generalized Procrustes before averaging. Requires the probabilistic extra.

Parameters:

Name	Type	Description	Default
n_components	int	Number of latent dimensions. Default is 1.	1
center	bool	Whether to centre each view. Default is True.	True
n_warmup	int	NUTS warm-up steps. Default is 500.	500
n_posterior_samples	int	NUTS draws. Default is 1000.	1000
random_state	int | None	Seed for the JAX PRNG. Default is None.	None

Attributes:

Name	Type	Description
weights_	list[ndarray]	Posterior mean loadings of each view, shape (n_features_i, n_components).
posterior_samples_	dict[str, Any]	Aligned posterior draws keyed by site name.

References

Bach, F. R., & Jordan, M. I. (2005). A probabilistic interpretation of canonical correlation analysis. Technical Report 688, UC Berkeley.

Examples:

>>> import numpy as np
>>> from cca_zoo.probabilistic import ProbabilisticCCA
>>> rng = np.random.default_rng(0)
>>> X1 = rng.standard_normal((50, 4))
>>> X2 = rng.standard_normal((50, 3))
>>> model = ProbabilisticCCA(
...     n_components=2, n_warmup=10, n_posterior_samples=10
... ).fit([X1, X2])

Source code in cca_zoo/probabilistic/_pcca.py

def __init__(
    self,
    n_components: int = 1,
    *,
    center: bool = True,
    n_warmup: int = 500,
    n_posterior_samples: int = 1000,
    random_state: int | None = None,
) -> None:
    super().__init__(n_components=n_components, center=center)
    self.n_warmup = n_warmup
    self.n_posterior_samples = n_posterior_samples
    self.random_state = random_state

fit

fit(
    views: list[ArrayLike], y: None = None
) -> ProbabilisticCCA

Fit the model.

Parameters:

Name	Type	Description	Default
views	list[ArrayLike]	Arrays of shape (n_samples, n_features_i), one per view.	required
y	None	Ignored.	None

Returns:

Type	Description
ProbabilisticCCA	self.

Source code in cca_zoo/probabilistic/_pcca.py

def fit(self, views: list[ArrayLike], y: None = None) -> ProbabilisticCCA:
    """Fit the model.

    Args:
        views: Arrays of shape (n_samples, n_features_i), one per view.
        y: Ignored.

    Returns:
        self.
    """
    import jax
    from numpyro.infer import MCMC, NUTS, init_to_value

    validated = self._setup_fit(views)

    # Sampling starts at the maximum likelihood.
    start = maximum_likelihood_start(validated, self.n_components)
    nuts_kernel = NUTS(self._model, init_strategy=init_to_value(values=start))
    mcmc = MCMC(
        nuts_kernel,
        num_warmup=self.n_warmup,
        num_samples=self.n_posterior_samples,
        progress_bar=False,
    )
    rng_key = jax.random.PRNGKey(_integer_seed(self.random_state))
    mcmc.run(rng_key, validated)
    self.posterior_samples_ = {
        k: np.array(v) for k, v in mcmc.get_samples().items()
    }

    # Resolve the model's rotational symmetry (see class docstring)
    # before averaging: align every view's stacked W_i draws to a common
    # reference.
    w_stack = np.concatenate(
        [self.posterior_samples_[f"W_{i}"] for i in range(self.n_views_)], axis=1
    )  # (n_posterior_samples, P, k)
    aligned_w, _ = align_posterior_rotation(w_stack)
    splits = np.cumsum(self.n_features_per_view_)[:-1]
    for i, w_i_aligned in enumerate(np.split(aligned_w, splits, axis=1)):
        self.posterior_samples_[f"W_{i}"] = w_i_aligned

    # Set weights_ to posterior mean W matrices (p_i x k) for each view
    self.weights_: list[np.ndarray] = [
        self.posterior_samples_[f"W_{i}"].mean(axis=0) for i in range(self.n_views_)
    ]
    self._fit_maps_and_importances(validated)
    return self

---

VariationalBayesCCA

VariationalBayesCCA(
    n_components: int = 1,
    *,
    center: bool = True,
    n_iter: int = 2000,
    learning_rate: float = 0.01,
    n_posterior_samples: int = 1000,
    random_state: int | None = None,
)

Bases: BaseProbabilistic

Variational Bayesian CCA with automatic relevance determination.

The probabilistic CCA model with an ARD prior shared across views:

\[ \begin{aligned} \alpha_k &\sim \mathrm{Gamma}(a_0, b_0), & W_i[:, k] &\sim \mathcal{N}(0, \alpha_k^{-1} I), \\ z &\sim \mathcal{N}(0, I), & x_i \mid z &\sim \mathcal{N}(W_i z + \mu_i, \Psi_i), \end{aligned} \]

with each \(\Psi_i\) a full covariance, as in :class:ProbabilisticCCA.

Dimensions no view supports are shrunk away, so set n_components generously and read ard_precision_. Inference is mean-field SVI in numpyro, cheaper than :class:ProbabilisticCCA's NUTS. Requires the probabilistic extra.

Parameters:

Name	Type	Description	Default
n_components	int	Upper bound on the number of latent dimensions. Default is 1.	1
center	bool	Whether to centre each view. Default is True.	True
n_iter	int	SVI steps; there is no stopping rule, so all are run. Default is 2000.	2000
learning_rate	float	Adam learning rate for SVI. Default is 1e-2.	0.01
n_posterior_samples	int	Draws from the fitted posterior. Default is 1000.	1000
random_state	int | None	Seed for the JAX PRNG. Default is None.	None

Attributes:

Name	Type	Description
weights_	list[ndarray]	Posterior mean loadings of each view, shape (n_features_i, n_components).
ard_precision_	list[ndarray]	Posterior mean ARD precision of each dimension; large means shrunk away.
posterior_samples_	dict[str, Any]	Posterior draws keyed by site name.
losses_	dict[str, Any]	SVI loss at each step.

References

Wang, C. (2007). Variational Bayesian approach to canonical correlation analysis. IEEE Transactions on Neural Networks, 18(3), 905-910.

Examples:

>>> import numpy as np
>>> from cca_zoo.probabilistic import VariationalBayesCCA
>>> rng = np.random.default_rng(0)
>>> X1 = rng.standard_normal((50, 4))
>>> X2 = rng.standard_normal((50, 3))
>>> model = VariationalBayesCCA(n_components=2, n_iter=50).fit([X1, X2])

Source code in cca_zoo/probabilistic/_vbcca.py

def __init__(
    self,
    n_components: int = 1,
    *,
    center: bool = True,
    n_iter: int = 2000,
    learning_rate: float = 1e-2,
    n_posterior_samples: int = 1000,
    random_state: int | None = None,
) -> None:
    super().__init__(n_components=n_components, center=center)
    self.n_iter = n_iter
    self.learning_rate = learning_rate
    self.n_posterior_samples = n_posterior_samples
    self.random_state = random_state

fit

fit(
    views: list[ArrayLike], y: None = None
) -> VariationalBayesCCA

Fit the model.

Parameters:

Name	Type	Description	Default
views	list[ArrayLike]	Arrays of shape (n_samples, n_features_i), one per view.	required
y	None	Ignored.	None

Returns:

Type	Description
VariationalBayesCCA	self.

Source code in cca_zoo/probabilistic/_vbcca.py

def fit(self, views: list[ArrayLike], y: None = None) -> VariationalBayesCCA:
    """Fit the model.

    Args:
        views: Arrays of shape (n_samples, n_features_i), one per view.
        y: Ignored.

    Returns:
        self.
    """
    import jax
    import numpyro.optim as optim
    from numpyro.infer import SVI, Predictive, Trace_ELBO, init_to_value
    from numpyro.infer.autoguide import AutoNormal

    validated = self._setup_fit(views)

    # From the saddle at zero loadings, where the noise explains each view
    # alone, SVI is slow to find the shared signal: start at the maximum
    # likelihood instead.
    start = maximum_likelihood_start(validated, self.n_components)
    guide = AutoNormal(self._model, init_loc_fn=init_to_value(values=start))
    svi = SVI(self._model, guide, optim.Adam(self.learning_rate), Trace_ELBO())

    rng_key, predictive_key = jax.random.split(
        jax.random.PRNGKey(_integer_seed(self.random_state))
    )
    svi_result = svi.run(rng_key, self.n_iter, validated, progress_bar=False)
    self.losses_: np.ndarray = np.array(svi_result.losses)

    predictive = Predictive(
        guide, params=svi_result.params, num_samples=self.n_posterior_samples
    )
    self.posterior_samples_: dict[str, np.ndarray] = {
        site: np.array(draws)
        for site, draws in predictive(predictive_key, validated).items()
    }

    # Set weights_ to variational posterior mean W matrices (p_i x k)
    self.weights_: list[np.ndarray] = [
        self.posterior_samples_[f"W_{i}"].mean(axis=0) for i in range(self.n_views_)
    ]
    # Posterior mean ARD precision per latent dimension: larger means
    # "more shrunk / less relevant".
    self.ard_precision_: np.ndarray = self.posterior_samples_["alpha"].mean(axis=0)
    self._fit_maps_and_importances(validated)
    return self