diff --git a/gpcr/covariance_cov_shrinkage.py b/gpcr/covariance_cov_shrinkage.py new file mode 100644 index 0000000..9208ff0 --- /dev/null +++ b/gpcr/covariance_cov_shrinkage.py @@ -0,0 +1,367 @@ +""" +Covariance Matrix Estimation with Shrinkage Techniques + +This module implements three different shrinkage techniques for covariance +matrix estimation: Linear Inverse Shrinkage (LIS), Quadratic Inverse +Shrinkage (QIS), and Geometric Inverse Shrinkage (GIS). These methods are +designed to improve the estimation of covariance matrices in scenarios +where the number of features may be comparable to or exceed the number +of observations. + +Linear Inverse Shrinkage (LIS): +-------------------------------- +LIS aims to improve the conditioning of the sample covariance matrix by +shrinking its eigenvalues. The shrinkage intensity is determined by the +ratio of the number of features to the number of observations. +Mathematically, the shrunk covariance matrix is computed as: + + S_hat = U * diag(d_i) * U^T + +where U is the matrix of eigenvectors, d_i are the shrunk eigenvalues +calculated from the original eigenvalues λ_i of the sample covariance +matrix, and the shrinkage targets are based on the overall variance. + +Quadratic Inverse Shrinkage (QIS): +---------------------------------- +QIS extends the LIS by considering a quadratic form of the eigenvalue +adjustment. It is particularly useful when dealing with outliers or +heavy-tailed distributions. The shrunk eigenvalues in QIS are calculated +as: + + d_i = 1 / (c^2 * l_i + (1 - c^2) / p * S(1/l_j)) + +where c is the concentration ratio (p/n), l_i are the eigenvalues, and p +is the number of features. + +Geometric Inverse Shrinkage (GIS): +---------------------------------- +GIS uses a geometric mean of the eigenvalues to determine the shrinkage +intensity. This method is less sensitive to the specific distribution of +eigenvalues and provides a balance between the largest and smallest +variances. The formula for GIS is given by: + + d_i = (l_i * h) / (l_i + h - l_i * h) + +where h is a parameter typically depending on the geometric mean of the +eigenvalues. + +These techniques are implemented as classes that inherit from a base +covariance class, allowing them to be easily integrated with other +statistical analysis pipelines. + +Import Dependencies: + numpy as np + numpy.typing as NDArray + math + from bystro.covariance._base_covariance import BaseCovariance + + +This code is a minimally modified version of the code in the CovShrinkage +package (the code differs solely to remove pandas). However, the advantage +of incorporating the package into our framework is that it automatically +inherits all the nice bells and whistles you can do with a covariance matrix +in terms of prediction/imputation with multivariate Gaussians. + +LICENSE FROM CovShrinkage + +https://github.com/pald22/covShrinkage/blob/main/LICENSE + +MIT License + +Copyright (c) 2022 Anonymous Panda + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +""" + +import math +from typing import Optional + +import numpy as np +from numpy.typing import NDArray + +from ._base_covariance import ( + BaseCovariance, + _symmeterize_and_warning, +) + + +class GeometricInverseShrinkage(BaseCovariance): + def __init__(self) -> None: + """ + Initialize the Geometric Inverse Shrinkage covariance estimator. + """ + super().__init__() + + def fit(self, X: NDArray[np.float64]) -> "GeometricInverseShrinkage": + """ + Fit the Geometric Inverse Shrinkage model to the given data. + + Parameters + ---------- + X : NDArray[np.float64] + Input data with shape (n_samples, n_features). + + Returns + ------- + self : GeometricInverseShrinkage + The instance itself. + """ + covariance = gis(X) + self.covariance = _symmeterize_and_warning(covariance) + return self + + +class LinearInverseShrinkage(BaseCovariance): + def __init__(self) -> None: + """ + Initialize the Linear Inverse Shrinkage covariance estimator. + """ + super().__init__() + + def fit(self, X: NDArray[np.float64]) -> "LinearInverseShrinkage": + """ + Fit the Linear Inverse Shrinkage model to the given data. + + Parameters + ---------- + X : NDArray[np.float64] + Input data with shape (n_samples, n_features). + + Returns + ------- + self : LinearInverseShrinkage + The instance itself. + """ + covariance = lis(X) + self.covariance = _symmeterize_and_warning(covariance) + return self + + +class QuadraticInverseShrinkage(BaseCovariance): + def __init__(self) -> None: + """ + Initialize the Quadratic Inverse Shrinkage covariance estimator. + """ + super().__init__() + + def fit(self, X: NDArray[np.float64]) -> "QuadraticInverseShrinkage": + """ + Fit the Quadratic Inverse Shrinkage model to the given data. + + Parameters + ---------- + X : NDArray[np.float64] + Input data with shape (n_samples, n_features). + + Returns + ------- + self : QuadraticInverseShrinkage + The instance itself. + """ + covariance = qis(X) + self.covariance = _symmeterize_and_warning(covariance) + return self + + +def gis(Y: NDArray[np.float64], k: Optional[int] = None) -> NDArray[np.float64]: + """ + Compute the Geometric Inverse Shrinkage covariance matrix. + + Parameters + ---------- + Y : NDArray[np.float64] + Input data matrix with shape (n_samples, n_features). + k : int, optional + Adjustment to the degrees of freedom. Default is None. + + Returns + ------- + NDArray[np.float64] + The shrunk covariance matrix. + """ + N, p = Y.shape + if N <= p: + raise ValueError( + "p must be <= n for the Symmetrized Kullback-Leibler divergence" + ) + + if k is None or math.isnan(k): + Y = Y - Y.mean(axis=0) + k = 1 + + n = N - k + c = p / n + + sample = np.dot(Y.T, Y) / n + sample = (sample + sample.T) / 2 + + lambda1, u = np.linalg.eigh(sample) + lambda1 = lambda1.clip(min=0) + indices = np.argsort(lambda1) + lambda1 = lambda1[indices] + u = u[:, indices] + + h = (min(c**2, 1 / c**2) ** 0.35) / p**0.35 + invlambda = 1 / lambda1[max(1, p - n + 1) - 1 : p] + + Lj = np.tile(invlambda, (len(invlambda), 1)) + Lj = Lj.T + Lj_i = Lj - Lj.T + + num = Lj * Lj_i + den = Lj_i**2 + Lj**2 * h**2 + theta = np.mean(num / den, axis=0) + Htheta = np.mean(Lj * Lj * h / den, axis=0) + Atheta2 = theta**2 + Htheta**2 + + deltahat_1 = (1 - c) * invlambda + 2 * c * invlambda * theta + delta = 1 / ( + (1 - c) ** 2 * invlambda + + 2 * c * (1 - c) * invlambda * theta + + c**2 * invlambda * Atheta2 + ) + + deltaLIS_1 = np.maximum(deltahat_1, np.min(invlambda)) + + temp2 = np.diag((delta / deltaLIS_1) ** 0.5) + sigmahat = np.dot(np.dot(u, temp2), u.T.conjugate()) + + return sigmahat + + +def lis(Y: NDArray[np.float64], k: Optional[int] = None) -> NDArray[np.float64]: + """ + Compute the Linear Inverse Shrinkage covariance matrix. + + Parameters + ---------- + Y : NDArray[np.float64] + Input data matrix with shape (n_samples, n_features). + k : int, optional + Adjustment to the degrees of freedom. Default is None. + + Returns + ------- + NDArray[np.float64] + The shrunk covariance matrix. + """ + N, p = Y.shape + if N <= p: + raise ValueError("p must be <= n for Stein's loss") + + if k is None or math.isnan(k): + Y = Y - np.mean(Y, axis=0) # demean + k = 1 + + n = N - k # adjust effective sample size + c = p / n # concentration ratio + + sample = np.dot(Y.T, Y) / n + sample = (sample + sample.T) / 2 # make symmetrical + + lambda1, u = np.linalg.eigh(sample) # use symmetric decomposition + lambda1 = np.clip(lambda1, 0, None) # reset negative values to 0 + + h = (min(c**2, 1 / c**2) ** 0.35) / p**0.35 + + valid_range = max(1, p - n + 1) - 1 + invlambda = 1 / lambda1[valid_range:p] + + # Matrix operations to calculate theta + Lj = np.tile(invlambda, (len(invlambda), 1)) + Lj = Lj.T + Lj_i = Lj - Lj.T + numerator = Lj * Lj_i + denominator = Lj_i**2 + (Lj**2) * h**2 + theta = np.mean(numerator / denominator, axis=0) + + deltahat_1 = (1 - c) * invlambda + 2 * c * invlambda * theta + + # Ensure no eigenvalue shrinkage below minimum + deltaLIS_1 = np.maximum(deltahat_1, np.min(invlambda)) + + # Reconstruct covariance matrix + temp2 = np.diag(1 / deltaLIS_1) + sigmahat = np.dot(np.dot(u, temp2), u.T.conjugate()) + + return sigmahat + + +def qis(Y: NDArray[np.float64], k: Optional[int] = None) -> NDArray[np.float64]: + """ + Compute the Quadratic Inverse Shrinkage covariance matrix. + + Parameters + ---------- + Y : NDArray[np.float64] + Input data matrix with shape (n_samples, n_features). + k : int, optional + Adjustment to the degrees of freedom. Default is None. + + Returns + ------- + NDArray[np.float64] + The shrunk covariance matrix. + """ + N, p = Y.shape + + if k is None or math.isnan(k): + Y = Y - np.mean(Y, axis=0) + k = 1 + + n = N - k + c = p / n + + sample = np.matmul(Y.T, Y) / n + sample = (sample + sample.T) / 2 + + lambda1, u = np.linalg.eigh(sample) + lambda1 = np.clip(lambda1, 0, None) + + h = (min(c**2, 1 / c**2) ** 0.35) / p**0.35 + invlambda = 1 / lambda1[-min(p, n) :] + + Lj = np.tile(invlambda, (len(invlambda), 1)) + Lj_i = Lj - Lj.T + Lj_i = Lj_i.T + Lj = Lj.T + + num = Lj * Lj_i + den = Lj_i**2 + Lj**2 * h**2 + theta = np.mean(num / den, axis=0) + Htheta = np.mean(Lj * Lj * h / den, axis=0) + Atheta2 = theta**2 + Htheta**2 + + if p <= n: + delta = 1 / ( + (1 - c) ** 2 * invlambda + + 2 * c * (1 - c) * invlambda * theta + + c**2 * invlambda * Atheta2 + ) + else: + delta0 = 1 / ((c - 1) * np.mean(invlambda)) + delta = np.repeat(delta0, p - n) + delta = np.concatenate((delta, 1 / (invlambda * Atheta2))) + + deltaQIS = delta * (np.sum(lambda1) / np.sum(delta)) + + temp2 = np.diag(deltaQIS) + sigmahat = np.dot(np.dot(u, temp2), u.T.conjugate()) + + return sigmahat diff --git a/gpcr/generative_fa.py b/gpcr/generative_fa.py new file mode 100644 index 0000000..5adacf1 --- /dev/null +++ b/gpcr/generative_fa.py @@ -0,0 +1,559 @@ +from typing import Any, Callable + +import numpy as np +import torch +from numpy.typing import NDArray +from sklearn.decomposition import PCA # type: ignore +from torch import Tensor, nn +from torch.distributions.gamma import Gamma +from torch.distributions.multivariate_normal import MultivariateNormal +from tqdm import trange + +from ._base import BasePCASGDModel +from ._misc_np import softplus_inverse_np +from ._sherman_woodbury_pt import mvn_log_prob_sw + + +class PPCA(BasePCASGDModel): + def __init__( + self, + n_components: int = 2, + prior_options: dict | None = None, + training_options: dict | None = None, + ): + super().__init__( + n_components=n_components, + prior_options=prior_options, + training_options=training_options, + ) + self._initialize_save_losses() + + self.W_: NDArray[np.float_] | None = None + self.sigmas_: NDArray[np.float_] | None = None + self.sigma2_: np.float_ | None = None + self.p: int | None = None + + def __repr__(self): + return f"PPCApt(n_components={self.n_components})" + + def fit( + self, + X: NDArray[np.float_], + progress_bar: bool = True, + seed: int = 2021, + sherman_woodbury: bool = False, + ) -> "PPCA": + """ + Fits a model given covariates X as well as option labels y in the + supervised methods + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + progress_bar : bool,default=True + Whether to print the progress bar to monitor time + + seed : int,default=2021 + The seed of the random number generator + + sherman_woodbury : bool,default=False + Whether to use the Sherman Woodbury identity to calculate + the likelihood. Advantageous in high-p situations + + Returns + ------- + self : PPCA + The model + """ + self._test_inputs(X) + training_options = self.training_options + N, p = X.shape + self.p = p + device = torch.device( + "cuda" + if torch.cuda.is_available() and training_options["use_gpu"] + else "cpu" + ) + + rng = np.random.default_rng(int(seed)) + + W_, sigmal_ = self._initialize_variables(device, X) + + X_tensor = self._transform_training_data(device, X)[0] + + trainable_variables = [W_, sigmal_] + + optimizer = torch.optim.SGD( + trainable_variables, + lr=training_options["learning_rate"], + momentum=training_options["momentum"], + ) + eye = torch.tensor(np.eye(p).astype(np.float32), device=device) + eye_L = torch.tensor( + np.eye(self.n_components).astype(np.float32), device=device + ) + zeros_p = torch.tensor( + np.zeros(p).astype(np.float32), dtype=torch.float32, device=device + ) + softplus = nn.Softplus() + + _prior = self._create_prior(device) + + for i in trange( + training_options["n_iterations"], disable=not progress_bar + ): + idx = rng.choice( + X_tensor.shape[0], + size=training_options["batch_size"], + replace=False, + ) + X_batch = X_tensor[idx].float() + + sigma = softplus(sigmal_) + + if sherman_woodbury: + Lambda = sigma * eye + like_tot = mvn_log_prob_sw(X_batch, zeros_p, Lambda, W_, eye_L) + else: + WWT = torch.matmul(torch.transpose(W_, 0, 1), W_) + Sigma = WWT + sigma * eye + m = MultivariateNormal(zeros_p, Sigma) + like_tot = torch.mean(m.log_prob(X_batch)) + + like_prior = _prior(trainable_variables) + posterior = like_tot + like_prior / N + loss = -1 * posterior + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + self._save_losses(i, device, like_tot, like_prior, posterior) + + self._store_instance_variables(device, trainable_variables) + + return self + + def get_covariance(self): + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigma2_ is None or self.p is None: + raise ValueError("Fit model first") + + return np.dot(self.W_.T, self.W_) + self.sigma2_ * np.eye(self.p) + + def get_noise(self): + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigma2_ is None or self.p is None: + raise ValueError("Fit model first") + + return self.sigma2_ * np.eye(self.p) + + def _create_prior(self, device): + """ + This creates the function representing prior on pararmeters + + Parameters + ---------- + log_prior : function + The function representing the log density of the prior + """ + prior_options = self.prior_options + + def log_prior(trainable_variables: list[Tensor]): + W_ = trainable_variables[0] + sigmal_ = trainable_variables[1] + sigma_ = nn.Softplus()(sigmal_) + + part1 = ( + -1 * prior_options["weight_W"] * torch.mean(torch.square(W_)) + ) + part2 = ( + Gamma(prior_options["alpha"], prior_options["beta"]) + .log_prob(sigma_) + .to(device) + ) + out = torch.mean(part1 + part2) + return out + + return log_prior + + def _initialize_variables( + self, device: Any, X: NDArray[np.float_] + ) -> tuple[Tensor, Tensor]: + """ + Initializes the variables of the model. Right now fits a PCA model + in sklearn, uses the loadings and sets sigma^2 to be unexplained + variance. + + Parameters + ---------- + device ; pytorch.device + The device used for trainging (gpu or cpu) + + X : NDArray,(n_samples,p) + The data + + Returns + ------- + W_ : torch.tensor,shape=(n_components,p) + The loadings of our latent factor model + + sigmal_ : torch.tensor + The unrectified variance of the model + """ + model = PCA(self.n_components) + S_hat = model.fit_transform(X) + W_init = model.components_.astype(np.float32) + W_ = torch.tensor(W_init, requires_grad=True, device=device) + X_recon = np.dot(S_hat, W_init) + diff = np.mean((X - X_recon) ** 2) + sinv = softplus_inverse_np(diff * np.ones(1).astype(np.float32)) + sigmal_ = torch.tensor(sinv, requires_grad=True, device=device) + return W_, sigmal_ + + def _store_instance_variables( # type: ignore[override] + self, device: Any, trainable_variables: list[Tensor] + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list[Tensor] + List of tensorflow variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : np.float_ + The isotropic variance + """ + if device.type == "cuda": + self.W_ = trainable_variables[0].detach().cpu().numpy() + self.sigma2_ = ( + nn.Softplus()(trainable_variables[1]).detach().cpu().numpy() + ) + else: + self.W_ = trainable_variables[0].detach().numpy() + self.sigma2_ = ( + nn.Softplus()(trainable_variables[1]).detach().numpy() + ) + + def _test_inputs(self, X: NDArray[np.float_]) -> None: + """ + Just tests to make sure data is numpy array + """ + if not isinstance(X, np.ndarray): + raise ValueError("Data is numpy array") + if self.training_options["batch_size"] > X.shape[0]: + raise ValueError("Batch size exceeds number of samples") + + def _fill_prior_options( + self, prior_options: dict[str, Any] + ) -> dict[str, Any]: + """ + Fills in options for prior parameters + + Paramters + --------- + new_dict : dictionary + The prior parameters used to specify the prior + """ + default_dict = {"weight_W": 0.01, "alpha": 3.0, "beta": 3.0} + + return {**default_dict, **prior_options} + + +class FactorAnalysis(BasePCASGDModel): + def __init__( + self, + n_components=2, + prior_options: dict | None = None, + training_options: dict | None = None, + ): + """ + This implements factor analysis which allows for each covariate to + have it's own isotropic noise. No analytic solution that I know of + but fortunately with SGD it doesn't matter. + + Parameters + ---------- + n_components : int,default=2 + The latent dimensionality + + training_options : dict,default={} + The options for gradient descent + + prior_options : dict,default={} + The options for priors on model parameters + """ + super().__init__( + n_components=n_components, + prior_options=prior_options, + training_options=training_options, + ) + self._initialize_save_losses() + self.p: int | None = None + self.W_: NDArray[np.float_] | None = None + self.sigmas_: NDArray[np.float_] | None = None + + def __repr__(self) -> str: + return f"FactorAnalysispt(n_components={self.n_components})" + + def fit( + self, + X: NDArray[np.float_], + progress_bar: bool = True, + seed: int = 2021, + sherman_woodbury: bool = False, + ) -> "FactorAnalysis": + """ + Fits a model given covariates X + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + Returns + ------- + self : FactorAnalysis + The model + """ + self._test_inputs(X) + training_options = self.training_options + device = torch.device( + "cuda" + if torch.cuda.is_available() and training_options["use_gpu"] + else "cpu" + ) + + N, p = X.shape + self.p = p + + rng = np.random.default_rng(int(seed)) + + W_, sigmal_ = self._initialize_variables(device, X) + + X_ = self._transform_training_data(device, X)[0] + + trainable_variables = [W_, sigmal_] + + optimizer = torch.optim.SGD( + trainable_variables, + lr=training_options["learning_rate"], + momentum=training_options["momentum"], + ) + softplus = nn.Softplus() + + _prior = self._create_prior(device) + + eye = torch.tensor(np.eye(p).astype(np.float32), device=device) + zeros_p = torch.zeros(p, device=device) + + for i in trange( + training_options["n_iterations"], disable=not progress_bar + ): + idx = rng.choice( + X_.shape[0], size=training_options["batch_size"], replace=False + ) + X_batch = X_[idx] + + sigmas = softplus(sigmal_) + Lambda = torch.diag(sigmas) + + if sherman_woodbury: + like_tot = mvn_log_prob_sw(X_batch, zeros_p, Lambda, W_, eye) + else: + WWT = torch.matmul(torch.transpose(W_, 0, 1), W_) + Sigma = WWT + Lambda + m = MultivariateNormal(zeros_p, Sigma) + like_tot = torch.mean(m.log_prob(X_batch)) + like_prior = _prior(trainable_variables) + posterior = like_tot + like_prior / N + loss = -1 * posterior + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + self._save_losses(i, device, like_tot, like_prior, posterior) + + self._store_instance_variables(device, trainable_variables) + + return self + + def get_covariance(self) -> NDArray[np.float_]: + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigmas_ is None: + raise ValueError("Fit model first") + + return np.dot(self.W_.T, self.W_) + np.diag(self.sigmas_) + + def get_noise(self) -> NDArray[np.float_]: + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigmas_ is None: + raise ValueError("Fit model first") + + return np.diag(self.sigmas_) + + def _create_prior(self, device) -> Callable[[list[Tensor]], Tensor]: + """ + This creates the function representing prior on pararmeters + + Parameters + ---------- + log_prior : function + The function representing the negative log density of the prior + """ + + def log_prior(trainable_variables: list[Tensor]) -> Tensor: + sigma_ = nn.Softmax()(trainable_variables[1]) + return torch.mean( + Gamma(self.prior_options["alpha"], self.prior_options["beta"]) + .log_prob(sigma_) + .to(device) + ) + + return log_prior + + def _fill_prior_options( + self, prior_options: dict[str, Any] + ) -> dict[str, Any]: + """ + Fills in options for prior parameters + + Paramters + --------- + new_dict : dictionary + The prior parameters used to specify the prior + """ + default_dict = {"alpha": 3.0, "beta": 3.0} + return {**default_dict, **prior_options} + + def _initialize_variables(self, device: Any, X: NDArray[np.float_]): + """ + Initializes the variables of the model by fitting PCA model in + sklearn and using those loadings + + Parameters + ---------- + X : NDArray,(n_samples,p) + The data + + Returns + ------- + W_ : torch.tensor,(n_components,p) + The loadings of our latent factor model + + sigmal_ : torch.tensor,(p,) + The noise of each covariate, unrectified + """ + if self.p is None: + raise ValueError("Fit model first") + + model = PCA(self.n_components) + S_hat = model.fit_transform(X) + W_init = model.components_.astype(np.float32) + W_ = torch.tensor(W_init, requires_grad=True, device=device) + X_recon = np.dot(S_hat, W_init) + diff = np.mean((X - X_recon) ** 2) + sinv = softplus_inverse_np(diff * np.ones(1)) + sigmal_ = torch.tensor( + sinv[0] * np.ones(self.p).astype(np.float32), + requires_grad=True, + device=device, + ) + + return W_, sigmal_ + + def _store_instance_variables( # type: ignore[override] + self, device: Any, trainable_variables: list[Tensor] + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of tensorflow variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigmas_ : NDArray,(n_components,p) + The diagonal variances + """ + if device.type == "cuda": + self.W_ = trainable_variables[0].detach().cpu().numpy() + self.sigmas_ = ( + nn.Softplus()(trainable_variables[1]).detach().cpu().numpy() + ) + else: + self.W_ = trainable_variables[0].detach().numpy() + self.sigmas_ = ( + nn.Softplus()(trainable_variables[1]).detach().numpy() + ) + + def _test_inputs(self, X: NDArray[np.float_]): + """ + Just tests to make sure data is numpy array + """ + if not isinstance(X, np.ndarray): + raise ValueError("Data is numpy array") + if self.training_options["batch_size"] > X.shape[0]: + raise ValueError("Batch size exceeds number of samples") diff --git a/gpcr/gpcr.py b/gpcr/gpcr.py new file mode 100644 index 0000000..e124ec4 --- /dev/null +++ b/gpcr/gpcr.py @@ -0,0 +1,552 @@ +""" +This implements Factor analysis but with variational inference to allow us +to supervise a subset of latent variables to be predictive of an auxiliary +outcome. Currently only implements logistic regression but in the future +will be modified to allow for more general predictive outcomes. + +Note that in the code W is a L x p array, where L is the latent dimensionality +and p is the covariate dimension, for implementation convenience. However, +mathematically it is often pxL for notational convenience. Given that the +most insidious errors are mathematical in nature rather than coding, (as faulty +math is difficult to detect in unit tests), our notation matches mathematics +rather than code, specifically when naming WWT and WTW to match the Bishop 2006 +notation rather than code. + + +Objects +------- +GPCRvi(PPCA) + This is PPCA with SGD but supervises a single factor to predict y. + +Methods +------- +None +""" + +import numpy as np +import torch +from numpy.typing import NDArray +from sklearn.decomposition import PCA +from sklearn.linear_model import LogisticRegression +from torch import Tensor, nn +from torch.distributions.gamma import Gamma +from torch.distributions.multivariate_normal import MultivariateNormal +from tqdm import trange + +from ._misc_np import softplus_inverse_np +from .generative_fa import PPCA +from .ppca_augmented import PPCAsupervisedRandomized + + +def _get_projection_matrix(W_: Tensor, sigma_: Tensor): + """ + This is currently just implemented for PPCA due to nicer formula. Will + modify for broader application later. + + Computes the parameters for p(S|X) + + Description in future to be released paper + + Parameters + ---------- + W_ : Tensor(n_components,p) + The loadings + + sigma_ : Tensor + Isotropic noise + + Returns + ------- + Proj_X : Tensor(n_components,p) + Beta such that np.dot(Proj_X, X) = E[S|X] + + Cov : Tensor(n_components,n_components) + Var(S|X) + """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + n_components = int(W_.shape[0]) + eye = torch.tensor(np.eye(n_components).astype(np.float32), device=device) + M_init = torch.matmul(W_, torch.transpose(W_, 0, 1)) + M_end = sigma_ * eye + M = M_init + M_end + Proj_X = torch.linalg.solve(M, W_) + Cov = torch.linalg.inv(M) * sigma_ + return Proj_X, Cov + + +class GPCRvi(PPCA): + """ + This implements supervised PPCA according to the paper draft in + prepration (Talbot et al, 2023), robust variational inference with + variational objectivesThat is the generative mechanism matches + probabilistic PCA with isotropic variance. However, a variational + lower bound ona predictive objective is used to ensure that a subset + of latent variables are predictive of an auxiliary task. + + Parameters + ---------- + n_components : int + The latent dimensionality + + n_supervised : int + The number of predictive latent variables + + prior_options : dict + The hyperparameters for the prior on the parameters + + mu : float>0,default=1.0 + + gamma : float,default=10.0 + + d_weight_min : float or None, default=None + Lower bound on each element of the predictive weight vector. + If None, no lower bound is applied. + + d_weight_max : float or None, default=None + Upper bound on each element of the predictive weight vector. + If None, no upper bound is applied. + + d_bias_min : float or None, default=None + Lower bound on the predictive bias term. + If None, no lower bound is applied. + + d_bias_max : float or None, default=None + Upper bound on the predictive bias term. + If None, no upper bound is applied. + + d_freeze_iters : int, default=0 + Number of iterations to keep predictive coefficients D = (d_weights, + d_bias) frozen at their warm-start initialization. During this phase, + only W and sigma^2 receive gradient updates. This implements Phase 1 + of a block-coordinate-ascent strategy: the generative parameters learn + to arrange the latent space so the *fixed* D already predicts well, + which forces the posterior's supervised dimensions to genuinely carry + predictive information. Set to 0 to disable (joint optimization from + the start, matching the original gPCR Algorithm 1). + + d_warmup_iters : int, default=0 + After the freeze phase ends, the effective learning-rate multiplier + for D is linearly ramped from 0 to 1 over this many iterations. + Concretely, at iteration i in the warmup window the D-gradients are + scaled by alpha(i) = (i - d_freeze_iters) / d_warmup_iters. + This prevents a sudden gradient shock from destabilizing the + converged (W, sigma^2). Set to 0 to unfreeze D instantly. + + d_anchor_weight : float, default=0.0 + Coefficient of a proximal / L2-anchoring penalty + lambda * (||d_weights - d_weights_init||^2 + ||d_bias - d_bias_init||^2) + that is added to the loss once D is unfrozen. The effective weight is + annealed from d_anchor_weight down to 0 over the warmup window (or + held constant if d_warmup_iters == 0 and d_anchor_weight > 0). + This regularizes early D-updates toward the warm-start, acting as a + trust region. Set to 0.0 to disable. + + + """ + + def __init__( + self, + n_components: int = 2, + n_supervised: int = 1, + prior_options: dict | None = None, + mu: float = 1.0, + gamma: float = 10.0, + d_weight_min: float | None = None, + d_weight_max: float | None = None, + d_bias_min: float | None = None, + d_bias_max: float | None = None, + d_freeze_iters: int = 0, + d_warmup_iters: int = 0, + d_anchor_weight: float = 0.0, + training_options: dict | None = None, + ): + self.mu = float(mu) + self.gamma = float(gamma) + self.d_weight_min = d_weight_min + self.d_weight_max = d_weight_max + self.d_bias_min = d_bias_min + self.d_bias_max = d_bias_max + self.d_freeze_iters = int(d_freeze_iters) + self.d_warmup_iters = int(d_warmup_iters) + self.d_anchor_weight = float(d_anchor_weight) + self.n_supervised = int(n_supervised) + super().__init__( + n_components=n_components, + prior_options=prior_options, + training_options=training_options, + ) + self._initialize_save_losses() + self.losses_supervision = np.empty( + self.training_options["n_iterations"] + ) + + # override needed for mypy to ignore the non-optional `y` argument + def fit( # type: ignore[override] + self, + X: NDArray[np.float_], + y: NDArray[np.float_], + task: str = "classification", + progress_bar: bool = True, + seed: int = 2021, + ) -> "GPCRvi": + """ + Fits a model given covariates X as well as option labels y in the + supervised methods + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + y : NDArray,(n_samples,n_prediction) + Covariates we wish to predict. For now lazy and assuming + logistic regression. + + task : string,default='classification' + Is this prediction, multinomial regression, or classification + + progress_bar : bool,default=True + Whether to print the progress bar to monitor time + + seed : int,default=2021 + The random number generator seed used to ensure reproducibility + + Returns + ------- + self : GPCRvi + The model + """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # self._test_inputs(X, y) + rng = np.random.default_rng(int(seed)) + training_options = self.training_options + N, p = X.shape + self.p = p + + W_, sigmal_ = self._initialize_variables(X, y) + + # X_, y_ = self._transform_training_data(X, 1.0 * y) + X_ = torch.tensor(X.astype(np.float32)) + y_ = torch.tensor(y.astype(np.float32)) + X_ = X_.to(device) + y_ = y_.to(device) + + # ---- Initialize learnable predictive coefficients D ---- + # D = (d_weights, d_bias) is the affine predictive model from the + # supervised latent dimensions to the outcome, jointly optimized + # with W and sigma^2 as described in gPCR Algorithm 1. + # + # d_weights: initialized from a warm-start logistic/linear regression + # d_bias: initialized from the same warm-start model's intercept + if task == "classification": + sigm = nn.Sigmoid() + supervision_loss = nn.BCELoss(reduction="mean") + + mod = LogisticRegression() + mod.fit(X, 1.0 * y) + + # Project the full-space logistic regression coefficients into + # the supervised latent subspace to get a warm-start for d_weights. + # W_sup is (n_supervised, p); coef is (1, p). + W_sup_np = W_.detach().cpu().numpy()[: self.n_supervised] + coef_projected = mod.coef_ @ W_sup_np.T # (1, n_supervised) + d_weights = torch.tensor( + coef_projected.flatten().astype(np.float32), + requires_grad=True, + device=device, + ) + d_bias = torch.tensor( + mod.intercept_.astype(np.float32), + requires_grad=True, + device=device, + ) + elif task == "regression": + supervision_loss = nn.MSELoss() + d_weights = torch.ones( + self.n_supervised, + requires_grad=True, + device=device, + dtype=torch.float32, + ) + d_bias = torch.zeros( + 1, + requires_grad=True, + device=device, + dtype=torch.float32, + ) + else: + err_msg = f"unrecognized_task {task}, must be regression or classification" + raise ValueError(err_msg) + + # ---- Save initial D for proximal anchoring ---- + d_weights_init = d_weights.detach().clone() + d_bias_init = d_bias.detach().clone() + + trainable_variables = [W_, sigmal_, d_weights, d_bias] + + optimizer = torch.optim.SGD( + trainable_variables, + lr=training_options["learning_rate"], + momentum=training_options["momentum"], + ) + + eye = torch.tensor(np.eye(p).astype(np.float32), device=device) + softplus = nn.Softplus() + + _prior = self._create_prior() + + # ---- Precompute unfreezing schedule boundaries ---- + freeze_end = self.d_freeze_iters + warmup_end = freeze_end + self.d_warmup_iters + + for i in trange( + int(training_options["n_iterations"]), disable=not progress_bar + ): + idx = rng.choice( + X_.shape[0], size=training_options["batch_size"], replace=False + ) + X_batch = X_[idx] + y_batch = y_[idx] + + sigma = softplus(sigmal_) + WWT = torch.matmul(torch.transpose(W_, 0, 1), W_) + Sigma = WWT + sigma * eye + + like_prior = _prior(trainable_variables) + + # Generative likelihood + m = MultivariateNormal(torch.zeros(p, device=device), Sigma) + like_gen = torch.mean(m.log_prob(X_batch)) + + # Predictive lower bound + P_x, Cov = _get_projection_matrix(W_, sigma) + mean_z = torch.matmul(X_batch, torch.transpose(P_x, 0, 1)) + + # Reparameterization trick: eps ~ N(0, I_L) + eps = torch.randn_like(mean_z) + + C1_2 = torch.linalg.cholesky(Cov) + z_samples = mean_z + torch.matmul(eps, C1_2) + + z_supervised = z_samples[:, : self.n_supervised] + + if task == "regression": + y_hat = torch.matmul(z_supervised, d_weights) + d_bias + loss_y = supervision_loss( + torch.squeeze(y_hat), torch.squeeze(y_batch) + ) + else: + y_hat = torch.matmul(z_supervised, d_weights) + d_bias + loss_y = supervision_loss(sigm(y_hat), y_batch) + + WTW = torch.matmul(W_, torch.transpose(W_, 0, 1)) + off_diag = WTW - torch.diag(torch.diag(WTW)) + loss_i = torch.linalg.matrix_norm(off_diag) + + posterior = like_gen + 1 / N * like_prior + loss = -1 * posterior + self.mu * loss_y + self.gamma * loss_i + + # ---- Proximal anchoring penalty on D ---- + # Active only after the freeze phase, annealed over the warmup. + if self.d_anchor_weight > 0.0 and i >= freeze_end: + if self.d_warmup_iters > 0 and i < warmup_end: + # Linear decay from d_anchor_weight -> 0 over warmup + anchor_frac = 1.0 - (i - freeze_end) / self.d_warmup_iters + else: + # No warmup or past warmup: no anchoring + anchor_frac = 0.0 + if anchor_frac > 0.0: + anchor_loss = ( + anchor_frac + * self.d_anchor_weight + * ( + torch.sum(torch.square(d_weights - d_weights_init)) + + torch.sum(torch.square(d_bias - d_bias_init)) + ) + ) + loss = loss + anchor_loss + + optimizer.zero_grad() + loss.backward() + + # ---- Staged unfreezing: modify D-gradients before step ---- + if i < freeze_end: + # Phase 1: D is frozen. Zero out D-gradients entirely so + # the optimizer (including its momentum buffer) sees no + # D-signal. W and sigma^2 train freely. + if d_weights.grad is not None: + d_weights.grad.zero_() + if d_bias.grad is not None: + d_bias.grad.zero_() + elif i < warmup_end and self.d_warmup_iters > 0: + # Phase 2 (warmup): linearly ramp D's effective learning + # rate from 0 to 1. Scaling the gradient is equivalent to + # a time-varying LR multiplier for these parameters only. + alpha = (i - freeze_end) / self.d_warmup_iters + if d_weights.grad is not None: + d_weights.grad.mul_(alpha) + if d_bias.grad is not None: + d_bias.grad.mul_(alpha) + # else: Phase 3 (full optimization) — no modification needed. + + optimizer.step() + + # ---- Project predictive coefficients onto feasible set ---- + with torch.no_grad(): + d_weights.clamp_(min=self.d_weight_min, max=self.d_weight_max) + d_bias.clamp_(min=self.d_bias_min, max=self.d_bias_max) + + self._save_losses(i, like_gen, like_prior, posterior) + self.losses_supervision[i] = loss_y.detach().cpu().numpy() + + self._store_instance_variables(trainable_variables) + + return self + + def _store_instance_variables(self, trainable_variables: list[Tensor]): + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of saved variables of type Tensor + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : float + The isotropic variance + + d_weights_ : NDArray,(n_supervised,) + The learned predictive coefficients + + d_bias_ : NDArray,(1,) + The learned predictive bias + + """ + self.W_ = trainable_variables[0].detach().cpu().numpy() + self.sigma2_ = ( + nn.Softplus()(trainable_variables[1]).detach().cpu().numpy() + ) + self.d_weights_ = trainable_variables[2].detach().cpu().numpy() + self.d_bias_ = trainable_variables[3].detach().cpu().numpy() + + def _initialize_variables(self, X: NDArray, Y: NDArray): + """ + Initializes the variables of the model. Right now fits a PCA model + in sklearn, uses the loadings and sets sigma^2 to be unexplained + variance. + + Parameters + ---------- + X : NDArray,(n_samples,p) + The data + + Returns + ------- + W_ : torch.tensor,shape=(n_components,p) + The loadings of our latent factor model + + sigmal_ : torch.tensor + The unrectified variance of the model + + """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = PPCAsupervisedRandomized(n_components=1, mu=100000.0) + model.fit(X, Y.reshape((-1, 1))) + W0 = model.W_ + S_hat = model.transform(X) + model = PCA(self.n_components) + model.fit(X) + W_init = model.components_ + if np.dot(np.squeeze(S_hat), np.squeeze(Y)) > 0: + W_init[0] = W0 + else: + W_init[0] = -1 * W0 + + S_hat = model.fit_transform(X) + W_ = torch.tensor(W_init, requires_grad=True, device=device) + X_recon = np.dot(S_hat, W_init) + diff = np.mean((X - X_recon) ** 2) + sinv = softplus_inverse_np(diff * np.ones(1)) + sigmal_ = torch.tensor( + sinv, requires_grad=True, dtype=torch.float32, device=device + ) + return W_, sigmal_ + + def _test_inputs(self, X, y): + """ + Just tests to make sure data is numpy array and dimensions match + """ + if not isinstance(X, np.ndarray): + raise ValueError("Data must be numpy array") + if self.training_options["batch_size"] > X.shape[0]: + raise ValueError("Batch size exceeds number of samples") + if X.shape[0] != len(y): + err_msg = "Length of data matrix X must equal length of labels y" + raise ValueError(err_msg) + + def _create_prior(self): + """ + This creates the function representing prior on pararmeters + + Parameters + ---------- + log_prior : function + The function representing the log density of the prior + """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + prior_options = self.prior_options + + def log_prior(trainable_variables: list[Tensor]): + W_ = trainable_variables[0] + sigmal_ = trainable_variables[1] + sigma_ = nn.Softplus()(sigmal_) + + part1 = ( + -1 * prior_options["weight_W"] * torch.mean(torch.square(W_)) + ) + part2 = ( + Gamma(prior_options["alpha"], prior_options["beta"]) + .log_prob(sigma_) + .to(device) + ) + out = torch.mean(part1 + part2) + return out + + return log_prior + + def _save_losses( + self, + i, + log_likelihood: Tensor, + log_prior: NDArray[np.float_] | Tensor, + log_posterior, + ) -> None: + """ + Saves the values of the losses at each iteration + + Parameters + ----------- + i : int + Current training iteration + + losses_likelihood : Tensor + The log likelihood + + losses_prior : NDArray[np.float_] | Tensor + The log prior + + losses_posterior : Tensor + The log posterior + """ + self.losses_likelihood[i] = log_likelihood.detach().cpu().numpy() + if isinstance(log_prior, Tensor): + self.losses_prior[i] = log_prior.detach().cpu().numpy() + else: + self.losses_prior[i] = log_prior + self.losses_posterior[i] = log_posterior.detach().cpu().numpy() diff --git a/gpcr/ppca_augmented.py b/gpcr/ppca_augmented.py new file mode 100644 index 0000000..7b6dfea --- /dev/null +++ b/gpcr/ppca_augmented.py @@ -0,0 +1,1077 @@ +from typing import Union + +import numpy as np +import numpy.linalg as la +from numpy.typing import NDArray +from sklearn.utils.extmath import randomized_range_finder # type: ignore + +from ._base import BaseGaussianFactorModel +from ._covariance_np import ( + EmpiricalCovariance, + LinearShrinkageCovariance, + NonLinearShrinkageCovariance, +) +from .covariance_cov_shrinkage import ( + GeometricInverseShrinkage, + LinearInverseShrinkage, + QuadraticInverseShrinkage, +) + + +class BasePPCAAugmented(BaseGaussianFactorModel): + """ + This performs transform according to standard PCA where we don't scale + based on noise + """ + + def transform(self, X): + S = np.dot(X, self.W_.T) + return S + + +class PPCAadversarial(BasePPCAAugmented): + """ + Probabilistic PCA that mitigates the influence of confounding + variables in the latent representation. + + This class extends the traditional PCA model by incorporating an + adversarial mechanism that penalizes the representation of confounding + variables. This approach not only reduces dimensionality but also + ensures that the resulting components are more representative of the + true underlying signals rather than artifacts introduced by confounders. + + Parameters + ---------- + mu : float + The adversarial strength which controls the penalty for + representing the confounding variables in the latent representation. + eps : float + A small constant added to improve the conditioning of + matrix inverses during computation. + + Attributes + ---------- + components_ : array, shape (n_components, n_features) + Principal axes in feature space, representing the directions of + maximum variance in the data. + + Notes + ----- + This implementation is based on modifications to the + 'augmented-pca' package by Carson, Talbot and Carlson + """ + + def __init__( + self, + n_components: int = 2, + mu=1.0, + eps=1e-8, + regularization="Linear", + ): + super().__init__(n_components=n_components) + self.regularization = regularization + self.mu: float = mu + self.eps: float = eps + self.W_: NDArray[np.float_] | None = None + self.D_: NDArray[np.float_] | None = None + self.p: int | None = None + self.sigma2_: np.float_ | None = None + + def fit( + self, X: NDArray[np.float_], Y: NDArray[np.float_] + ) -> "PPCAadversarial": + """ + Fits a model given covariates X + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + Y : NDArray,(n_samples,n_confounders) + The concommitant data we want to remove + + Returns + ------- + self : PPCAadversarial + The model + """ + self.mean_x = np.mean(X, axis=0) + self.mean_y = np.mean(Y, axis=0) + X_dm = X - self.mean_x + Y_dm = Y - self.mean_y + N, self.p = X.shape + p = self.p + q = Y.shape[1] + + model_cov = _select_covariance_estimator(self.regularization) + XX = np.zeros((N, p + q)) + XX[:, :p] = X_dm + if q == 1: + XX[:, -1] = np.squeeze(Y_dm) + else: + XX[:, p:] = Y_dm + model_cov.fit(XX) + cov = model_cov.covariance + if cov is None: + raise ValueError("Covariance matrix failed to fit") + B_11 = cov[:p, :p] + B_12 = cov[:p, p:] + B_22 = B_12.T @ la.solve(B_11 + self.eps * np.eye(p), B_12) + + B = np.zeros((p + q, p + q)) + B[:p, :p] = B_11 + B[:p, p:] = B_12 + B[p:, :p] = -self.mu * B_12.T + B[p:, p:] = -self.mu * B_22 + B = B.T + + eigvals, eigvecs = la.eig(B) + eigvals = np.real(eigvals) + eigvecs = np.real(eigvecs) + idx = eigvals.argsort()[::-1] + self.eigvals_ = eigvals[idx] + V = eigvecs[:, idx] + W = V[:p, : self.n_components] + D = V[p:, : self.n_components] + A = W.T @ W + B = W.T @ X_dm.T + S = la.solve(A, B).T + X_recon = S @ W.T + var = np.mean((X - X_recon) ** 2) + self._store_instance_variables((W, var, D)) + + return self + + def get_covariance(self) -> NDArray[np.float_]: + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray,(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return np.dot(self.W_.T, self.W_) + self.sigma2_ * np.eye(self.p) + + def get_noise(self): + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return self.sigma2_ * np.eye(self.p) + + def _store_instance_variables( + self, + trainable_variables: tuple[ + NDArray[np.float_], np.float_, NDArray[np.float_] + ], + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : float + The isotropic variance + """ + self.W_ = trainable_variables[0].T + self.sigma2_ = trainable_variables[1] + self.D_ = trainable_variables[2].T + + def _initialize_save_losses(self): + pass + + def _save_losses(self): + pass + + def _test_inputs(self): + pass + + def _transform_training_data(self): + pass + + def _save_variables(self): + pass + + +class PPCAsupervised(BasePPCAAugmented): + """ + Probabilistic PCA that mitigates the influence of confounding + variables in the latent representation. + + This class extends the traditional PCA model by incorporating an + adversarial mechanism that penalizes the representation of confounding + variables. This approach not only reduces dimensionality but also + ensures that the resulting components are more representative of the + true underlying signals rather than artifacts introduced by confounders. + + Parameters + ---------- + mu : float + The supervision strength which controls the penalty for + failing to predict the supervised variables in the latent + representation. + + eps : float + A small constant added to improve the conditioning of + matrix inverses during computation. + + Attributes + ---------- + components_ : array, shape (n_components, n_features) + Principal axes in feature space, representing the directions of + maximum variance in the data. + + Notes + ----- + This implementation is based on modifications to the + 'augmented-pca' package by Carson, Talbot and Carlson + """ + + def __init__( + self, + n_components: int = 2, + mu=1.0, + eps=1e-8, + regularization="Linear", + ): + super().__init__(n_components=n_components) + self.regularization = regularization + self.mu: float = mu + self.eps: float = eps + self.W_: NDArray[np.float_] | None = None + self.Phi_: NDArray[np.float_] | None = None + self.p: int | None = None + self.sigma2_: np.float_ | None = None + + def fit( + self, X: NDArray[np.float_], Y: NDArray[np.float_] + ) -> "PPCAsupervised": + """ + Fits a model given covariates X + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + Y : NDArray,(n_samples,n_predicted) + The data we want to predict + + Returns + ------- + self : PPCAadversarial + The model + """ + self.mean_x = np.mean(X, axis=0) + self.mean_y = np.mean(Y, axis=0) + X_dm = X - self.mean_x + Y_dm = Y - self.mean_y + N, self.p = X.shape + p = self.p + q = Y.shape[1] + + model_cov = _select_covariance_estimator(self.regularization) + XX = np.zeros((N, p + q)) + XX[:, :p] = X_dm + if q == 1: + XX[:, -1] = np.squeeze(Y_dm) + else: + XX[:, p:] = Y_dm + model_cov.fit(XX) + cov = model_cov.covariance + if cov is None: + raise ValueError("Covariance matrix failed to fit") + B_11 = cov[:p, :p] + B_12 = cov[:p, p:] + B_22 = B_12.T @ la.solve(B_11 + self.eps * np.eye(p), B_12) + + B = np.zeros((p + q, p + q)) + B[:p, :p] = B_11 + B[:p, p:] = B_12 + B[p:, :p] = self.mu * B_12.T + B[p:, p:] = self.mu * B_22 + B = B.T + + eigvals, eigvecs = la.eig(B) + eigvals = np.real(eigvals) + eigvecs = np.real(eigvecs) + idx = eigvals.argsort()[::-1] + self.eigvals_ = eigvals[idx] + V = eigvecs[:, idx] + W = V[:p, : self.n_components] + D = V[p:, : self.n_components] + A = W.T @ W + B = W.T @ X_dm.T + S = la.solve(A, B).T + X_recon = S @ W.T + var = np.mean((X - X_recon) ** 2) + self._store_instance_variables((W, var, D)) + + return self + + def get_covariance(self) -> NDArray[np.float_]: + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray,(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return np.dot(self.W_.T, self.W_) + self.sigma2_ * np.eye(self.p) + + def get_noise(self): + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return self.sigma2_ * np.eye(self.p) + + def _store_instance_variables( + self, + trainable_variables: tuple[ + NDArray[np.float_], np.float_, NDArray[np.float_] + ], + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : float + The isotropic variance + """ + self.W_ = trainable_variables[0].T + self.sigma2_ = trainable_variables[1] + self.Phi_ = trainable_variables[2].T + + def _initialize_save_losses(self): + pass + + def _save_losses(self): + pass + + def _test_inputs(self): + pass + + def _transform_training_data(self): + pass + + def _save_variables(self): + pass + + +class PPCASupAdversarial(BasePPCAAugmented): + """ + Probabilistic PCA that mitigates the influence of confounding + variables in the latent representation. + + This class extends the traditional PCA model by incorporating an + adversarial mechanism that penalizes the representation of confounding + variables. This approach not only reduces dimensionality but also + ensures that the resulting components are more representative of the + true underlying signals rather than artifacts introduced by confounders. + + Parameters + ---------- + mu : float + The supervision strength which controls the penalty for + predicting the predicted variables in the latent representation. + + eps : float + A small constant added to improve the conditioning of + matrix inverses during computation. + + Attributes + ---------- + components_ : array, shape (n_components, n_features) + Principal axes in feature space, representing the directions of + maximum variance in the data. + + Notes + ----- + This implementation is based on modifications to the + 'augmented-pca' package by Carson, Talbot and Carlson + """ + + def __init__( + self, + n_components: int = 2, + mu_s=1.0, + mu_a=1.0, + eps=1e-8, + regularization="Linear", + ): + super().__init__(n_components=n_components) + self.regularization = regularization + self.mu_s: float = mu_s + self.mu_a: float = mu_a + self.eps: float = eps + self.W_: NDArray[np.float_] | None = None + self.Phi_: NDArray[np.float_] | None = None + self.D_: NDArray[np.float_] | None = None + self.p: int | None = None + self.sigma2_: np.float_ | None = None + + def fit( + self, + X: NDArray[np.float_], + Y: NDArray[np.float_], + Z: NDArray[np.float_], + ) -> "PPCASupAdversarial": + """ + Fits a model given covariates X + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + Y : NDArray,(n_samples,n_predictors) + The data we want to predict + + Z : NDArray,(n_samples,n_confounders) + The concommitant data we want to remove + + Returns + ------- + self : PPCA_sup_adversarial + The model + """ + self.mean_x = np.mean(X, axis=0) + self.mean_y = np.mean(Y, axis=0) + self.mean_z = np.mean(Z, axis=0) + X_dm = X - self.mean_x + Y_dm = Y - self.mean_y + Z_dm = Z - self.mean_z + N, self.p = X.shape + p = self.p + q1 = Y.shape[1] + q2 = Z.shape[1] + + model_cov = _select_covariance_estimator(self.regularization) + XX = np.zeros((N, p + q1 + q2)) + XX[:, :p] = X_dm + + if q1 == 1: + XX[:, p] = np.squeeze(Y_dm) + else: + XX[:, p : (p + q1)] = Y_dm + + if q2 == 1: + XX[:, -1] = np.squeeze(Z_dm) + else: + XX[:, (p + q1) :] = Z_dm + + model_cov.fit(XX) + cov = model_cov.covariance + if cov is None: + raise ValueError("Covariance matrix failed to fit") + B_11 = cov[:p, :p] + B_12 = cov[:p, p : (p + q1)] + B_13 = cov[:p, (p + q1) :] + + B_22 = B_12.T @ la.solve(B_11 + self.eps * np.eye(p), B_12) + B_33 = B_13.T @ la.solve(B_11 + self.eps * np.eye(p), B_13) + + B_23 = B_12.T @ la.solve(B_11 + self.eps * np.eye(p), B_13) + + B = np.zeros((p + q1 + q2, p + q1 + q2)) + B[:p, :p] = B_11 + B[:p, p : (p + q1)] = B_12 + B[:p, (p + q1) :] = B_13 + + B[p : (p + q1), :p] = self.mu_s * B_12.T + B[p : (p + q1), p : (p + q1)] = self.mu_s * B_22 + B[p : (p + q1), (p + q1) :] = self.mu_s * B_23 + + B[(p + q1 + q2) :, :p] = -self.mu_a * B_13.T + B[(p + q1 + q2) :, p : (p + q1)] = -self.mu_a * B_23.T + B[(p + q1 + q2) :, (p + q1) :] = -self.mu_a * B_33 + + eigvals, eigvecs = la.eig(B) + eigvals = np.real(eigvals) + eigvecs = np.real(eigvecs) + idx = eigvals.argsort()[::-1] + self.eigvals_ = eigvals[idx] + V = eigvecs[:, idx] + W = V[:p, : self.n_components] + D1 = V[p : (p + q1), : self.n_components] + D2 = V[(p + q1) :, : self.n_components] + A = W.T @ W + B = W.T @ X_dm.T + S = la.solve(A, B).T + X_recon = S @ W.T + var = np.mean((X - X_recon) ** 2) + self._store_instance_variables((W, var, D1, D2)) + + return self + + def get_covariance(self) -> NDArray[np.float_]: + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray,(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return np.dot(self.W_.T, self.W_) + self.sigma2_ * np.eye(self.p) + + def get_noise(self): + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return self.sigma2_ * np.eye(self.p) + + def _store_instance_variables( + self, + trainable_variables: tuple[ + NDArray[np.float_], + np.float_, + NDArray[np.float_], + NDArray[np.float_], + ], + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : float + The isotropic variance + """ + self.W_ = trainable_variables[0].T + self.sigma2_ = trainable_variables[1] + self.D_ = trainable_variables[2].T + + def _initialize_save_losses(self): + pass + + def _save_losses(self): + pass + + def _test_inputs(self): + pass + + def _transform_training_data(self): + pass + + def _save_variables(self): + pass + + +def _select_covariance_estimator(regularization): + """ + This is a tiny method for selecting the estimator for the covariance + matrix. + """ + model_cov: Union[ + EmpiricalCovariance, + LinearShrinkageCovariance, + NonLinearShrinkageCovariance, + LinearInverseShrinkage, + GeometricInverseShrinkage, + QuadraticInverseShrinkage, + ] + if regularization == "Empirical": + model_cov = EmpiricalCovariance() + elif regularization == "Linear": + model_cov = LinearShrinkageCovariance() + elif regularization == "LinearInverse": + model_cov = LinearInverseShrinkage() + elif regularization == "QuadraticInverse": + model_cov = QuadraticInverseShrinkage() + elif regularization == "GeometricInverse": + model_cov = GeometricInverseShrinkage() + elif regularization == "NonLinear": + model_cov = NonLinearShrinkageCovariance() + elif regularization == "Bayesian": + raise ValueError("Bayesian currently not supported") + else: + raise ValueError("Unrecognized regularization %s" % regularization) + return model_cov + + +class PPCAadversarialRandomized(BasePPCAAugmented): + def __init__( + self, + n_components: int = 2, + mu: float = 1.0, + eps: float = 1e-8, + n_oversamples: int = 20, + random_state: int = 2021, + regularization: str = "Linear", + ): + super().__init__(n_components=n_components) + self.regularization = regularization + self.mu: float = mu + self.eps: float = eps + self.W_: NDArray[np.float_] | None = None + self.D_: NDArray[np.float_] | None = None + self.p: int | None = None + self.n_oversamples: int = n_oversamples + self.random_state: int = random_state + self.sigma2_: np.float_ | None = None + + def fit( + self, X: NDArray[np.float_], Y: NDArray[np.float_] + ) -> "PPCAadversarialRandomized": + """ + Fits a model given covariates X + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + Y : NDArray,(n_samples,n_confounders) + The concommitant data we want to remove + + Returns + ------- + self : PPCAadversarial + The model + """ + self.mean_x = np.mean(X, axis=0) + self.mean_y = np.mean(Y, axis=0) + X_dm = X - self.mean_x + Y_dm = Y - self.mean_y + N, self.p = X.shape + q = Y.shape[1] + + n_random = self.n_components + self.n_oversamples + Q = randomized_range_finder( + X_dm.T, + n_iter=7, + size=n_random, + power_iteration_normalizer="auto", + random_state=self.random_state, + ) + + X_tilde = Q.T @ X_dm.T + X_tilde = X_tilde.T + + model_cov = _select_covariance_estimator(self.regularization) + XX = np.zeros((N, n_random + q)) + XX[:, :n_random] = X_tilde + if q == 1: + XX[:, -1] = np.squeeze(Y_dm) + else: + XX[:, n_random:] = Y_dm + model_cov.fit(XX) + cov = model_cov.covariance + + B_11 = cov[:n_random, :n_random] + B_12 = cov[n_random:, :n_random].T + B_22 = B_12.T @ la.solve(B_11 + self.eps * np.eye(n_random), B_12) + + B = np.zeros((n_random + q, n_random + q)) + B[:n_random, :n_random] = B_11 + B[:n_random, n_random:] = B_12 + B[n_random:, :n_random] = -self.mu * B_12.T + B[n_random:, n_random:] = -self.mu * B_22 + B = B.T + + eigvals, eigvecs = la.eig(B) + eigvals = np.real(eigvals) + eigvecs = np.real(eigvecs) + idx = eigvals.argsort()[::-1] + + self.eigvals_ = eigvals[idx] + V_latent = eigvecs[:, idx] + + self.V_latent = V_latent + + W_latent = V_latent[:n_random, : self.n_components] + + W = Q @ W_latent + + A = W.T @ W + B = W.T @ X_dm.T + S = la.solve(A, B).T + self.explained_variance_ = np.mean(S**2, axis=0) + X_recon = S @ W.T + var = np.mean((X - X_recon) ** 2) + self._store_instance_variables((W, var)) + + return self + + def get_covariance(self) -> NDArray[np.float_]: + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray,(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + exp_var_diff = self.explained_variance_ - self.sigma2_ + W_mod = self.W_.T * np.sqrt(exp_var_diff) + return np.dot(W_mod, W_mod.T) + self.sigma2_ * np.eye(self.p) + + def get_noise(self): + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return self.sigma2_ * np.eye(self.p) + + def _store_instance_variables( + self, + trainable_variables: tuple[ + NDArray[np.float_], + np.float_, + ], + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : float + The isotropic variance + """ + self.W_ = trainable_variables[0].T + self.sigma2_ = trainable_variables[1] + + def _initialize_save_losses(self): + pass + + def _save_losses(self): + pass + + def _test_inputs(self): + pass + + def _transform_training_data(self): + pass + + def _save_variables(self): + pass + + +class PPCAsupervisedRandomized(BasePPCAAugmented): + def __init__( + self, + n_components: int = 2, + mu: float = 1.0, + eps: float = 1e-8, + n_oversamples: int = 20, + random_state: int = 2021, + regularization: str = "Linear", + ): + super().__init__(n_components=n_components) + self.regularization = regularization + self.mu: float = mu + self.eps: float = eps + self.W_: NDArray[np.float_] | None = None + self.D_: NDArray[np.float_] | None = None + self.p: int | None = None + self.n_oversamples: int = n_oversamples + self.random_state: int = random_state + self.sigma2_: np.float_ | None = None + + def fit( + self, X: NDArray[np.float_], Y: NDArray[np.float_] + ) -> "PPCAsupervisedRandomized": + """ + Fits a model given covariates X + + Parameters + ---------- + X : NDArray,(n_samples,n_covariates) + The data + + Y : NDArray,(n_samples,n_confounders) + The concommitant data we want to remove + + Returns + ------- + self : PPCAadversarial + The model + """ + self.mean_x = np.mean(X, axis=0) + self.mean_y = np.mean(Y, axis=0) + X_dm = X - self.mean_x + Y_dm = Y - self.mean_y + N, self.p = X.shape + q = Y.shape[1] + + n_random = self.n_components + self.n_oversamples + Q = randomized_range_finder( + X_dm.T, + n_iter=7, + size=n_random, + power_iteration_normalizer="auto", + random_state=self.random_state, + ) + + X_tilde = Q.T @ X_dm.T + X_tilde = X_tilde.T + + model_cov = _select_covariance_estimator(self.regularization) + XX = np.zeros((N, n_random + q)) + XX[:, :n_random] = X_tilde + if q == 1: + XX[:, -1] = np.squeeze(Y_dm) + else: + XX[:, n_random:] = Y_dm + model_cov.fit(XX) + cov = model_cov.covariance + + B_11 = cov[:n_random, :n_random] + B_12 = cov[n_random:, :n_random].T + B_22 = B_12.T @ la.solve(B_11 + self.eps * np.eye(n_random), B_12) + + B = np.zeros((n_random + q, n_random + q)) + B[:n_random, :n_random] = B_11 + B[:n_random, n_random:] = B_12 + B[n_random:, :n_random] = self.mu * B_12.T + B[n_random:, n_random:] = self.mu * B_22 + B = B.T + + eigvals, eigvecs = la.eig(B) + eigvals = np.real(eigvals) + eigvecs = np.real(eigvecs) + idx = eigvals.argsort()[::-1] + + self.eigvals_ = eigvals[idx] + V_latent = eigvecs[:, idx] + + self.V_latent = V_latent + + W_latent = V_latent[:n_random, : self.n_components] + + W = Q @ W_latent + + A = W.T @ W + B = W.T @ X_dm.T + S = la.solve(A, B).T + self.explained_variance_ = np.mean(S**2, axis=0) + X_recon = S @ W.T + var = np.mean((X - X_recon) ** 2) + self._store_instance_variables((W, var)) + + return self + + def get_covariance(self) -> NDArray[np.float_]: + """ + Gets the covariance matrix + + Sigma = W^TW + sigma2*I + + Parameters + ---------- + None + + Returns + ------- + covariance : NDArray,(p,p) + The covariance matrix + """ + if self.W_ is None or self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + exp_var_diff = self.explained_variance_ - self.sigma2_ + W_mod = self.W_.T * np.sqrt(exp_var_diff) + return np.dot(W_mod, W_mod.T) + self.sigma2_ * np.eye(self.p) + + def get_noise(self): + """ + Returns the observational noise as a diagonal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : NDArray,(p,p) + The observational noise + """ + if self.sigma2_ is None or self.p is None: + raise ValueError("Model has not been fit yet") + + return self.sigma2_ * np.eye(self.p) + + def _store_instance_variables( + self, + trainable_variables: tuple[ + NDArray[np.float_], + np.float_, + ], + ) -> None: + """ + Saves the learned variables + + Parameters + ---------- + trainable_variables : list + List of variables saved + + Sets + ---- + W_ : NDArray,(n_components,p) + The loadings + + sigma2_ : float + The isotropic variance + """ + self.W_ = trainable_variables[0].T + self.sigma2_ = trainable_variables[1] + + def _initialize_save_losses(self): + pass + + def _save_losses(self): + pass + + def _test_inputs(self): + pass + + def _transform_training_data(self): + pass + + def _save_variables(self): + pass + + +""" +https://github.com/wecarsoniv/augmented-pca/tree/main + +Copyright (c) 2021 The Python Packaging Authority + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +"""