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(
    latent_dimensions: int = 1,
    center: bool = True,
    max_iter: int = 10000,
    tol: float = 0.0001,
    drop_k: bool = True,
    num_posterior_samples: int = 1000,
    random_state: int = 0,
)

Bases: PosteriorMeanTransformMixin, BaseModel

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

Ported faithfully from the reference implementation, GFA() in the R package CCAGFA <https://github.com/cran/CCAGFA>_ (Klami, Virtanen & Kaski) — the update equations below are transliterated directly from that source rather than re-derived. Fits a single shared latent variable \(z\), but unlike :class:~cca_zoo.probabilistic.ProbabilisticCCA and :class:~cca_zoo.probabilistic.VariationalBayesCCA (which tie every view to the same ARD precision per latent dimension), GFA gives each view \(i\) its own ARD precision \(\alpha_{i,k}\) per latent dimension \(k\):

\[ \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_K) \\ \tau_i &\sim \mathrm{Gamma}(a_{0\tau}, b_{0\tau}) \\ x_i \mid z &\sim \mathcal{N}(W_i z,\ \tau_i^{-1} I) \end{aligned} \]

"Shared" vs. "private" latent dimensions are therefore emergent, not a fixed split of \(z\) into blocks: a dimension \(k\) ends up shared if \(\alpha_{i,k}\) stays small (loadings retained) in several views at once, and private to view \(i\) if \(\alpha_{i,k}\) shrinks toward zero loadings in every other view. view_relevance_ (posterior mean of \(\alpha_{i,k}\), shape (n_views, n_components_)) exposes this directly.

Note also the noise model: \(\tau_i\) is a single scalar precision per view (homoscedastic — every feature in a view shares the same noise variance), not a per-feature diagonal like the other two classes — this matches the R package exactly, not an approximation.

Inference is closed-form coordinate-ascent mean-field variational Bayes (conjugate throughout, so no black-box SVI is needed here unlike :class:~cca_zoo.probabilistic.VariationalBayesCCA). latent_dimensions is an upper bound: dimensions whose posterior mean squared value falls below 1e-7 in every view are pruned during fitting (drop_k=True, the R package's default), so the fitted number of components, n_components_, can end up smaller than latent_dimensions — every output array's last axis has size n_components_, not latent_dimensions.

Note

This port omits the R package's optional orthogonal-rotation step (opts$rotate, on by default in R) that speeds convergence and helps escape poor local optima; it doesn't change the fitted model class, only the optimization path, and is deferred to a follow-up rather than risk porting it incorrectly without a reference R run to check against.

Convergence is monitored via relative change in \(z\), sustained for 1000 consecutive iterations, rather than the R package's full variational lower bound (which is guaranteed monotonically non-decreasing under exact coordinate ascent — provably immune to the issue below). This is a best-effort speed heuristic, not a correctness guarantee: checking against a run with early stopping disabled entirely caught this proxy dipping below tolerance for 700+ consecutive iterations in the middle of a slow ARD pruning process (one dimension's decay temporarily dominating a still-shrinking one), before rising again once that pruning actually needed hundreds more iterations to finish — a patience window can make this less likely but, unlike the true ELBO, can't rule it out for an arbitrarily slow case. max_iter (default 10000) is the actual safety net: raise it if n_components_ looks larger than expected, rather than trusting early stopping alone on a hard pruning problem.

References

Klami, A., Virtanen, S., & Kaski, S. (2013). "Bayesian Canonical Correlation Analysis." Journal of Machine Learning Research, 14, 965-1003. Virtanen, S., Klami, A., & Kaski, S. (2011). "Bayesian CCA via Group Sparsity." ICML.

Parameters:

Name	Type	Description	Default
latent_dimensions	int	Upper bound on the number of latent components. Default is 1.	1
center	bool	Whether to center each view before fitting. Default is True.	True
max_iter	int	Maximum number of coordinate-ascent iterations, and the actual safety net for correctness (see the class-level note on early stopping being best-effort). Default is 10000 (the R package defaults to 1e5, using it purely as a cap around tol-based early stopping). Raise this if n_components_ comes out larger than expected on a hard problem.	10000
tol	float	Relative Frobenius-norm change in the latent variable \(z\) between iterations. Fitting stops once this stays below tol for 1000 consecutive iterations with no pruning event — see the class-level note for why even that isn't a full correctness guarantee. This is also a different quantity from the R package's iter.crit (a relative change in the full variational lower bound), so the two aren't numerically comparable; 1e-4 is calibrated against this specific proxy instead of copying the R default value. Default is 1e-4.	0.0001
drop_k	bool	Whether to prune latent dimensions with near-zero posterior mean squared value across the whole run. Default is True (matches the R package's dropK).	True
num_posterior_samples	int	Number of samples drawn from the fitted variational posterior to populate posterior_samples_. Default is 1000.	1000
random_state	int	Integer seed for reproducible initialization. Default is 0.	0

Example

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

Source code in cca_zoo/probabilistic/_gfa.py

def __init__(
    self,
    latent_dimensions: int = 1,
    center: bool = True,
    max_iter: int = 10000,
    tol: float = 1e-4,
    drop_k: bool = True,
    num_posterior_samples: int = 1000,
    random_state: int = 0,
) -> None:
    super().__init__(latent_dimensions=latent_dimensions, center=center)
    self.max_iter = max_iter
    self.tol = tol
    self.drop_k = drop_k
    self.num_posterior_samples = num_posterior_samples
    self.random_state = random_state

fit

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

Run coordinate-ascent variational Bayes to fit the GFA model.

Parameters:

Name	Type	Description	Default
views	list[ArrayLike]	List of arrays, each of shape (n_samples, n_features_i). All arrays must have the same number of rows.	required
y	None	Ignored. Present for scikit-learn API compatibility.	None

Returns:

Name	Type	Description
self	GFA	Fitted estimator.

Raises:

Type	Description
ValueError	If fewer than 2 views are provided.
ValueError	If views have inconsistent numbers of samples.

Source code in cca_zoo/probabilistic/_gfa.py

def fit(self, views: list[ArrayLike], y: None = None) -> GFA:
    """Run coordinate-ascent variational Bayes to fit the GFA model.

    Args:
        views: List of arrays, each of shape (n_samples, n_features_i).
            All arrays must have the same number of rows.
        y: Ignored.  Present for scikit-learn API compatibility.

    Returns:
        self: Fitted estimator.

    Raises:
        ValueError: If fewer than 2 views are provided.
        ValueError: If views have inconsistent numbers of samples.
    """
    validated = self._setup_fit(views)
    rng = np.random.default_rng(self.random_state)

    views_arr = [np.asarray(v, dtype=float) for v in validated]
    m_views = len(views_arr)
    n = views_arr[0].shape[0]
    d = [v.shape[1] for v in views_arr]
    k = self.latent_dimensions

    # --- initialization (CCAGFA::GFA(), matching getDefaultOpts()) ---
    z = rng.standard_normal((n, k))
    cov_z = np.eye(k)
    w = [np.zeros((d[m], k)) for m in range(m_views)]
    cov_w = [np.eye(k) for m in range(m_views)]
    tau = np.full(m_views, _INIT_TAU)
    datavar = np.array(
        [np.var(views_arr[m], axis=0, ddof=1).sum() for m in range(m_views)]
    )
    alpha = [
        np.full(k, k * d[m] / max(datavar[m] - 1.0 / tau[m], 1e-8))
        for m in range(m_views)
    ]

    y_const = np.array([np.sum(views_arr[m] ** 2) for m in range(m_views)])
    a_ard = _ARD_ALPHA_0 + np.array(d) / 2.0  # (M,), constant across iters
    a_tau = _TAU_ALPHA_0 + n * np.array(d) / 2.0  # (M,), constant

    ww = [w[m].T @ w[m] + d[m] * cov_w[m] for m in range(m_views)]
    zz = z.T @ z + n * cov_z
    b_ard = [np.full(k, _ARD_BETA_0) for _ in range(m_views)]
    b_tau = np.full(m_views, _TAU_BETA_0)

    # Requires `_PATIENCE` consecutive iterations with small relative
    # change AND no pruning event, not just one: rel_change can dip
    # below tol for a few hundred iterations in the middle of a slow
    # ARD-driven pruning process (one dimension's decay temporarily
    # dominated by a faster-settling one finishing first) before
    # rising again once that's the only signal left. A single-iteration
    # check was caught declaring convergence during exactly such a lull,
    # 2500+ iterations before the pruning it was still waiting on.
    prev_z: np.ndarray | None = None
    n_iter = self.max_iter
    stable_count = 0
    for iteration in range(self.max_iter):
        # --- W update (per view), using the current zz ---
        for m in range(m_views):
            tmp = 1.0 / np.sqrt(alpha[m])
            inner = np.outer(tmp, tmp) * zz + np.eye(k) / tau[m]
            cho_w = np.linalg.cholesky(inner)
            inv_inner = np.linalg.solve(cho_w.T, np.linalg.solve(cho_w, np.eye(k)))
            cov_w[m] = (1.0 / tau[m]) * np.outer(tmp, tmp) * inv_inner
            w[m] = views_arr[m].T @ z @ cov_w[m] * tau[m]
            ww[m] = w[m].T @ w[m] + d[m] * cov_w[m]

        # --- Z update, using the just-updated W ---
        precision_z = np.eye(k)
        for m in range(m_views):
            precision_z = precision_z + tau[m] * ww[m]
        cho_z = np.linalg.cholesky(precision_z)
        cov_z = np.linalg.solve(cho_z.T, np.linalg.solve(cho_z, np.eye(k)))
        rhs = np.zeros((n, k))
        for m in range(m_views):
            rhs = rhs + views_arr[m] @ w[m] * tau[m]
        z = rhs @ cov_z
        zz = z.T @ z + n * cov_z

        # --- alpha update (per view), using the just-updated ww ---
        for m in range(m_views):
            b_ard[m] = _ARD_BETA_0 + np.diag(ww[m]) / 2.0
            alpha[m] = a_ard[m] / b_ard[m]

        # --- tau update (per view) ---
        for m in range(m_views):
            b_tau[m] = (
                _TAU_BETA_0
                + (
                    y_const[m]
                    + np.sum(ww[m] * zz)
                    - 2.0 * np.sum(z * (views_arr[m] @ w[m]))
                )
                / 2.0
            )
            tau[m] = a_tau[m] / b_tau[m]

        # --- dynamic component pruning (dropK) ---
        pruned = False
        if self.drop_k:
            keep = np.where(np.mean(z**2, axis=0) > _DROP_TOL)[0]
            if 0 < len(keep) != k:
                pruned = True
                k = len(keep)
                z = z[:, keep]
                cov_z = cov_z[np.ix_(keep, keep)]
                zz = zz[np.ix_(keep, keep)]
                for m in range(m_views):
                    w[m] = w[m][:, keep]
                    cov_w[m] = cov_w[m][np.ix_(keep, keep)]
                    ww[m] = ww[m][np.ix_(keep, keep)]
                    alpha[m] = alpha[m][keep]
                    b_ard[m] = b_ard[m][keep]

        # --- convergence check (sustained small relative change in z) ---
        if pruned:
            stable_count = 0
        elif prev_z is not None and prev_z.shape == z.shape:
            rel_change = np.linalg.norm(z - prev_z) / max(
                np.linalg.norm(prev_z), 1e-300
            )
            stable_count = stable_count + 1 if rel_change < self.tol else 0
        prev_z = z.copy()
        if stable_count >= _PATIENCE:
            n_iter = iteration + 1
            break

    self.n_iter_ = n_iter
    self.n_components_ = k
    self._draw_posterior_samples(
        rng, z, cov_z, w, cov_w, a_ard, b_ard, a_tau, b_tau, d
    )
    self.weights_: list[np.ndarray] = list(w)
    self.view_relevance_: np.ndarray = np.array(alpha)
    return self

---

ProbabilisticCCA

ProbabilisticCCA(
    latent_dimensions: int = 1,
    center: bool = True,
    num_warmup: int = 500,
    num_samples: int = 1000,
    random_state: int = 0,
)

Bases: PosteriorMeanTransformMixin, BaseModel

Probabilistic Canonical Correlation Analysis via NUTS MCMC.

Fits a Bayesian latent variable model with the following generative process for \(V\) views:

\[ \begin{aligned} z &\sim \mathcal{N}(0, I) \\ x_i \mid z &\sim \mathcal{N}(W_i z + \mu_i,\ \Psi_i), \quad i = 1, \dots, V \end{aligned} \]

MCMC sampling is performed with the No-U-Turn Sampler (NUTS) from numpyro. After fitting, :meth:transform returns the posterior mean of z conditioned on the observed views (computed analytically using the posterior mean formula for linear Gaussian models).

This model has an exact rotational symmetry (\(z \to zR\), \(W_i \to W_i R\) for any orthogonal \(R\) shared across views leaves the likelihood unchanged), and different NUTS draws can settle on different rotations along that ridge of equal density. Averaging un-aligned draws for a point estimate is then biased toward zero (draws along different rotations partially cancel), so fit aligns every draw's loadings (and correspondingly, that draw's \(z\)) to a common reference via generalized Procrustes analysis (see :func:~cca_zoo.probabilistic._utils.align_posterior_rotation) before computing weights_ or storing posterior_samples_.

The weights_ attribute is set to the (rotation-aligned) posterior mean of each W_i matrix so that :class:~cca_zoo._base.BaseModel's scoring utilities work without modification.

References

Bach, F. R. & Jordan, M. I. "A probabilistic interpretation of canonical correlation analysis." (2005). Wang, C. "Variational Bayesian approach to canonical correlation analysis." IEEE Transactions on Neural Networks 18.3 (2007).

Parameters:

Name	Type	Description	Default
latent_dimensions	int	Dimensionality of the latent space. Default is 1.	1
center	bool	Whether to center each view before fitting. Default is True.	True
num_warmup	int	Number of NUTS warm-up (burn-in) steps. Default is 500.	500
num_samples	int	Number of NUTS posterior samples to draw. Default is 1000.	1000
random_state	int	Integer seed for JAX PRNG. Default is 0.	0

Example

import numpy as np rng = np.random.default_rng(0) X1 = rng.standard_normal((50, 4)) X2 = rng.standard_normal((50, 3)) model = ProbabilisticCCA( ... latent_dimensions=2, num_warmup=10, num_samples=10 ... ).fit([X1, X2])

Source code in cca_zoo/probabilistic/_pcca.py

def __init__(
    self,
    latent_dimensions: int = 1,
    center: bool = True,
    num_warmup: int = 500,
    num_samples: int = 1000,
    random_state: int = 0,
) -> None:
    super().__init__(latent_dimensions=latent_dimensions, center=center)
    self.num_warmup = num_warmup
    self.num_samples = num_samples
    self.random_state = random_state

fit

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

Run NUTS MCMC to infer posterior over model parameters and latents.

Parameters:

Name	Type	Description	Default
views	list[ArrayLike]	List of arrays, each of shape (n_samples, n_features_i). All arrays must have the same number of rows.	required
y	None	Ignored. Present for scikit-learn API compatibility.	None

Returns:

Name	Type	Description
self	ProbabilisticCCA	Fitted estimator.

Raises:

Type	Description
ValueError	If fewer than 2 views are provided.
ValueError	If views have inconsistent numbers of samples.

Source code in cca_zoo/probabilistic/_pcca.py

def fit(self, views: list[ArrayLike], y: None = None) -> ProbabilisticCCA:
    """Run NUTS MCMC to infer posterior over model parameters and latents.

    Args:
        views: List of arrays, each of shape (n_samples, n_features_i).
            All arrays must have the same number of rows.
        y: Ignored.  Present for scikit-learn API compatibility.

    Returns:
        self: Fitted estimator.

    Raises:
        ValueError: If fewer than 2 views are provided.
        ValueError: If views have inconsistent numbers of samples.
    """
    import jax
    from numpyro.infer import MCMC, NUTS

    validated = self._setup_fit(views)

    nuts_kernel = NUTS(self._model)
    mcmc = MCMC(
        nuts_kernel,
        num_warmup=self.num_warmup,
        num_samples=self.num_samples,
    )
    rng_key = jax.random.PRNGKey(self.random_state)
    mcmc.run(rng_key, validated)
    self.mcmc_ = mcmc
    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 any cross-draw averaging: stack every view's W_i draws,
    # align them to a common reference, then rotate that draw's z by
    # the same rotation to keep it internally consistent.
    w_stack = np.concatenate(
        [self.posterior_samples_[f"W_{i}"] for i in range(self.n_views_)], axis=1
    )  # (num_samples, P, k)
    aligned_w, rotations = align_posterior_rotation(w_stack)
    splits = np.cumsum(self.n_features_in_)[:-1]
    for i, w_i_aligned in enumerate(np.split(aligned_w, splits, axis=1)):
        self.posterior_samples_[f"W_{i}"] = w_i_aligned
    self.posterior_samples_["z"] = np.einsum(
        "snk,skj->snj", self.posterior_samples_["z"], rotations
    )

    # 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_)
    ]
    return self

---

VariationalBayesCCA

VariationalBayesCCA(
    latent_dimensions: int = 1,
    center: bool = True,
    num_steps: int = 2000,
    learning_rate: float = 0.01,
    num_posterior_samples: int = 1000,
    random_state: int = 0,
)

Bases: PosteriorMeanTransformMixin, BaseModel

Variational Bayesian CCA with automatic relevance determination.

Fits the same probabilistic CCA generative model as :class:~cca_zoo.probabilistic.ProbabilisticCCA, extended with a hierarchical automatic relevance determination (ARD) prior over the columns of the loading matrices, shared across views:

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

Because \(\alpha_k\) is shared across every view's \(k\)-th loading column, a latent dimension is only retained if some view finds it useful; dimensions unsupported by the data are shrunk towards zero in every view simultaneously. The posterior mean of \(\alpha_k\) (exposed as ard_relevance_) is therefore a direct, per-dimension usefulness score: large values indicate a dimension that has been shrunk away and can be dropped, giving automatic latent-dimensionality selection instead of a GridSearchCV sweep over latent_dimensions.

Inference uses mean-field stochastic variational inference (SVI) via numpyro, rather than the closed-form conjugate coordinate-ascent updates derived in Wang (2007) for this model: SVI reuses the exact same numpyro generative-model machinery as :class:~cca_zoo.probabilistic.ProbabilisticCCA, and (unlike a hand-derived conjugate solver) extends unmodified to non-conjugate variants of the model. It is a substantially cheaper alternative to that class's full NUTS MCMC, at the cost of the mean-field independence assumption between latent variables.

The weights_ attribute is set to the variational posterior mean of each \(W_i\) matrix so that :class:~cca_zoo._base.BaseModel's scoring utilities work without modification.

References

Bach, F. R. & Jordan, M. I. "A probabilistic interpretation of canonical correlation analysis." (2005). Wang, C. "Variational Bayesian approach to canonical correlation analysis." IEEE Transactions on Neural Networks 18.3 (2007).

Parameters:

Name	Type	Description	Default
latent_dimensions	int	Dimensionality of the latent space. Default is 1. Because of the ARD prior, this should be set generously (an upper bound on the number of shared factors you expect); use ard_relevance_ after fitting to see how many were retained.	1
center	bool	Whether to center each view before fitting. Default is True.	True
num_steps	int	Number of SVI gradient steps. Default is 2000.	2000
learning_rate	float	Adam learning rate for SVI. Default is 1e-2.	0.01
num_posterior_samples	int	Number of samples drawn from the fitted variational posterior to populate posterior_samples_. Default is 1000.	1000
random_state	int	Integer seed for JAX PRNG. Default is 0.	0

Example

import numpy as np rng = np.random.default_rng(0) X1 = rng.standard_normal((50, 4)) X2 = rng.standard_normal((50, 3)) model = VariationalBayesCCA( ... latent_dimensions=2, num_steps=50 ... ).fit([X1, X2])

Source code in cca_zoo/probabilistic/_vbcca.py

def __init__(
    self,
    latent_dimensions: int = 1,
    center: bool = True,
    num_steps: int = 2000,
    learning_rate: float = 1e-2,
    num_posterior_samples: int = 1000,
    random_state: int = 0,
) -> None:
    super().__init__(latent_dimensions=latent_dimensions, center=center)
    self.num_steps = num_steps
    self.learning_rate = learning_rate
    self.num_posterior_samples = num_posterior_samples
    self.random_state = random_state

fit

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

Run mean-field SVI to infer an approximate posterior.

Parameters:

Name	Type	Description	Default
views	list[ArrayLike]	List of arrays, each of shape (n_samples, n_features_i). All arrays must have the same number of rows.	required
y	None	Ignored. Present for scikit-learn API compatibility.	None

Returns:

Name	Type	Description
self	VariationalBayesCCA	Fitted estimator.

Raises:

Type	Description
ValueError	If fewer than 2 views are provided.
ValueError	If views have inconsistent numbers of samples.

Source code in cca_zoo/probabilistic/_vbcca.py

def fit(self, views: list[ArrayLike], y: None = None) -> VariationalBayesCCA:
    """Run mean-field SVI to infer an approximate posterior.

    Args:
        views: List of arrays, each of shape (n_samples, n_features_i).
            All arrays must have the same number of rows.
        y: Ignored.  Present for scikit-learn API compatibility.

    Returns:
        self: Fitted estimator.

    Raises:
        ValueError: If fewer than 2 views are provided.
        ValueError: If views have inconsistent numbers of samples.
    """
    import jax
    import numpyro.optim as optim
    from numpyro.infer import SVI, Predictive, Trace_ELBO
    from numpyro.infer.autoguide import AutoNormal

    validated = self._setup_fit(views)

    guide = AutoNormal(self._model)
    svi = SVI(self._model, guide, optim.Adam(self.learning_rate), Trace_ELBO())

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

    predictive = Predictive(
        guide, params=svi_result.params, num_samples=self.num_posterior_samples
    )
    self.posterior_samples_: dict[str, Any] = predictive(predictive_key, validated)

    # Set weights_ to variational posterior mean W matrices (p_i x k)
    self.weights_: list[np.ndarray] = [
        np.array(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_relevance_: np.ndarray = np.array(
        self.posterior_samples_["alpha"].mean(axis=0)
    )
    return self