From bc841530cbec1b1a92fadc6e7d4af71a432e2c7f Mon Sep 17 00:00:00 2001 From: Lobo Date: Tue, 16 Jun 2026 16:05:05 -0400 Subject: [PATCH 1/3] Added templates --- gpcr/_base.py | 778 +++++++++++++++++++++++++++++ gpcr/_base_covariance.py | 1015 ++++++++++++++++++++++++++++++++++++++ gpcr/_covariance_np.py | 286 +++++++++++ 3 files changed, 2079 insertions(+) create mode 100644 gpcr/_base.py create mode 100644 gpcr/_base_covariance.py create mode 100644 gpcr/_covariance_np.py diff --git a/gpcr/_base.py b/gpcr/_base.py new file mode 100644 index 0000000..9e61c31 --- /dev/null +++ b/gpcr/_base.py @@ -0,0 +1,778 @@ +""" +This provides the base class of any Gaussian factor model, such as +(probabilistic) PCA, supervised PCA, and factor analysis, as well as +supervised and adversarial derivatives. + +Implementing an extension model requires that the following methods be +implemented + fit - Learns the model given data (and optionally supervised/adversarial + labels + get_covariance - Given a fitted model returns the covariance matrix + +For the remaining shared methods computing likelihoods etc, there are two +options, compute the covariance matrix then use the default methods from +the covariance module, or use the Sherman-Woodbury matrix identity to +invert the matrix more efficiently. Currently only the first is implemented +but left the options for future implementation. + +Objects +------- +BaseGaussianFactorModel(_BaseSGDModel) + Base class of all factor analysis models implementing any shared + Gaussian methods. + +BaseSGDModel(BaseGaussianFactorModel) + This is the base class of models that use Tensorflow stochastic + gradient descent for inference. This reduces boilerplate code + associated with some of the standard methods etc. + +Methods +------- +None +""" + +from abc import ABC, abstractmethod +from typing import Any + +import numpy as np +import torch +from numpy import linalg as la +from numpy.typing import NDArray +from torch import Tensor + +from ._base_covariance import ( + _conditional_score, + _conditional_score_samples, + _conditional_score_samples_sherman_woodbury, + _conditional_score_sherman_woodbury, + _entropy, + _entropy_subset, + _get_stable_rank, + _marginal_score, + _marginal_score_samples, + _marginal_score_samples_sherman_woodbury, + _marginal_score_sherman_woodbury, + _mutual_information, + _score, + _score_samples, + _score_samples_sherman_woodbury, + _score_sherman_woodbury, + inv_sherman_woodbury_fa, +) + + +def _get_projection_matrix(W_: Tensor, sigma_: Tensor, device: Any): + """ + 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) + """ + 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 + + +def kl_divergence_vae( + mu1: torch.Tensor, + sigma1: torch.Tensor, +) -> torch.Tensor: + """ + Computes the Kullback-Leibler divergence between two multivariate + normal distributions. + + This function assumes the first distribution, N(mu1, sigma1), + represents the conditional distribution of latent variables given + observations (approximate posterior), and the second distribution + is the standard multivariate normal distribution N(0, I) representing + the marginal latent distribution (prior). + + The KL divergence is computed using the formula: + KL(N(mu1, sigma1) || N(0, I)) = 0.5 * (tr(sigma1) + + mu1^T * mu1 - log(det(sigma1)) - d) + where `d` is the dimensionality of the latent space. + + Parameters: + - mu1 (torch.Tensor): Mean vector of the first Gaussian + distribution (approximate posterior), shape (batch_size, d). + - sigma1 (torch.Tensor): Covariance matrix of the first Gaussian + distribution (approximate posterior), expected to be a diagonal + matrix represented as a tensor of shape (batch_size, d). + + Returns: + - torch.Tensor: A tensor containing the KL divergence for each + instance in the batch, shape (batch_size,). + """ + d = sigma1.shape[1] # Dimensionality of the latent space + term1 = -torch.sum( + torch.log(sigma1), dim=1 + ) # Log determinant of diagonal covariance + term2 = -d # Negative of the latent space dimensionality + term3 = torch.sum( + sigma1, dim=1 + ) # Trace of diagonal covariance (sum of diagonal) + term4 = torch.sum(torch.square(mu1), dim=1) # Sum of squares + kl_div = 0.5 * (term1 + term2 + term3 + term4) + return kl_div + + +class BaseGaussianFactorModel(ABC): + def __init__(self, n_components=2): + """ + This is the base class of the model. Will never be called directly. + + Parameters + ---------- + n_components : int,default=2 + The latent dimensionality + + Sets + ---- + creationDate : datetime + The date/time that the object was created + """ + self.n_components = int(n_components) + self.W_ = None + + @abstractmethod + def fit(self, *args, **kwargs): + """ + Fits a model given covariates X as well as option labels y in the + supervised methods + + Parameters + ---------- + X : np.array-like,(n_samples,n_covariates) + The data + + other arguments + + Returns + ------- + self : object + The model + """ + + @abstractmethod + def get_covariance(self) -> NDArray[np.float64]: + """ + Gets the covariance matrix defined by the model parameters + + Parameters + ---------- + None + + Returns + ------- + covariance : np.array-like(p,p) + The covariance matrix + """ + + @abstractmethod + def get_noise(self) -> NDArray[np.float64]: + """ + Returns the observational noise as a diagnoal matrix + + Parameters + ---------- + None + + Returns + ------- + Lambda : np.array-like(p,p) + The observational noise + """ + + def get_precision( + self, sherman_woodbury: bool = False + ) -> NDArray[np.float64]: + """ + Gets the precision matrix defined as the inverse of the covariance + + Parameters + ---------- + None + + Returns + ------- + precision : np.array-like(p,p) + The inverse of the covariance matrix + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + covariance = self.get_covariance() + return la.inv(covariance) + + return inv_sherman_woodbury_fa(self.get_noise(), self.W_) + + def get_stable_rank(self) -> np.float64: + """ + Returns the stable rank defined as + ||A||_F^2/||A||^2 + + Parameters + ---------- + None + + Returns + ------- + srank : np.float64 + The stable rank. See Vershynin High dimensional probability for + discussion, but this is a statistically stable approximation to rank + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + return _get_stable_rank(self.get_covariance()) + + def transform(self, X: NDArray[np.float64], sherman_woodbury: bool = False): + """ + This returns the latent variable estimates given X + + Parameters + ---------- + X : np array-like,(N_samples,p) + The data to transform. + + Returns + ------- + S : NDArray,(N_samples,n_components) + The factor estimates + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + prec = self.get_precision() + coefs = np.dot(self.W_, prec) + else: + W_ = self.W_ + M = np.dot(W_, W_.T) + self.sigma2_ * np.eye(self.n_components) + Mi = la.inv(M) + coefs = np.dot(Mi, W_) + # raise NotImplementedError("Subclass PCA required for this Sherman Woodbury") + + return np.dot(X, coefs.T) + + def transform_subset( + self, + X: NDArray[np.float64], + observed_feature_idxs: NDArray[np.float64], + sherman_woodbury: bool = False, + ) -> NDArray[np.float64]: + """ + This returns the latent variable estimates given partial observations + contained in X + + Parameters + ---------- + X : NDArray,(N_samples,sum(observed_feature_idxs)) + The data to transform. + + observed_feature_idxs: np.array-like,(sum(p),) + The observation locations + + Returns + ------- + S : NDArray,(N_samples,n_components) + The factor estimates + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + Wo = self.W_[:, observed_feature_idxs] + if sherman_woodbury is False: + covariance = self.get_covariance() + cov_sub = covariance[ + np.ix_(observed_feature_idxs, observed_feature_idxs) + ] + coefs = np.dot(Wo, la.inv(cov_sub)) + else: + M = np.dot(Wo, Wo.T) + self.sigma2_ * np.eye(self.n_components) + Mi = la.inv(M) + coefs = np.dot(Mi, Wo) + return np.dot(X, coefs.T) + + def conditional_score( + self, + X: NDArray[np.float64], + observed_feature_idxs: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, + sherman_woodbury: bool = False, + ) -> np.float64: + """ + Returns the predictive log-likelihood of a subset of data. + + mean(log p(X[idx==1]|X[idx==0],covariance)) + + Parameters + ---------- + X : NDArray,(N,sum(observed_feature_idxs)) + The data + + observed_feature_idxs: NDArray,(sum(p),) + The observation locations + + weights : NDArray,(N,),default=None + The optional weights on the samples + + Returns + ------- + avg_score : float + Average log likelihood + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + covariance = self.get_covariance() + return _conditional_score( + covariance, X, observed_feature_idxs, weights=weights + ) + + return _conditional_score_sherman_woodbury( + self.get_noise(), + self.W_, + X, + observed_feature_idxs, + weights=weights, + ) + + def conditional_score_samples( + self, + X: NDArray[np.float64], + observed_feature_idxs: NDArray[np.float64], + sherman_woodbury: bool = False, + ) -> NDArray[np.float64]: + """ + Return the conditional log likelihood of each sample, that is + + log p(X[idx==1]|X[idx==0],covariance) + + Parameters + ---------- + X : np.array-like,(N,p) + The data + + observed_feature_idxs: np.array-like,(p,) + The observation locations + + sherman_woodbury : bool,default=False + Whether to use the sherman_woodbury matrix identity + + Returns + ------- + scores : float + Log likelihood for each sample + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + covariance = self.get_covariance() + return _conditional_score_samples( + covariance, X, observed_feature_idxs + ) + + return _conditional_score_samples_sherman_woodbury( + self.get_noise(), self.W_, X, observed_feature_idxs + ) + + def marginal_score( + self, + X: NDArray[np.float64], + observed_feature_idxs: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, + sherman_woodbury: bool = False, + ) -> np.float64: + """ + Returns the marginal log-likelihood of a subset of data + + Parameters + ---------- + X : np.array-like,(N,sum(idxs)) + The data + + observed_feature_idxs: np.array-like,(sum(p),) + The observation locations + + weights : np.array-like,(N,),default=None + The optional weights on the samples + + Returns + ------- + avg_score : float + Average log likelihood + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + covariance = self.get_covariance() + return _marginal_score( + covariance, X, observed_feature_idxs, weights=weights + ) + + return _marginal_score_sherman_woodbury( + self.get_noise(), + self.W_, + X, + observed_feature_idxs, + weights=weights, + ) + + def marginal_score_samples( + self, + X: NDArray, + observed_feature_idxs: NDArray, + sherman_woodbury: bool = False, + ): + """ + Returns the marginal log-likelihood of a subset of data + + Parameters + ---------- + X : np.array-like,(N,sum(observed_feature_idxs)) + The data + + observed_feature_idxs: np.array-like,(sum(p),) + The observation locations + + Returns + ------- + scores : float + Average log likelihood + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + covariance = self.get_covariance() + return _marginal_score_samples(covariance, X, observed_feature_idxs) + + return _marginal_score_samples_sherman_woodbury( + self.get_noise(), self.W_, X, observed_feature_idxs + ) + + def score( + self, + X: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, + sherman_woodbury: bool = False, + ) -> np.float64: + """ + Returns the average log liklihood of data. + + Parameters + ---------- + X : np.array-like,(N,sum(p)) + The data + + weights : np.array-like,(N,),default=None + The optional weights on the samples + + Returns + ------- + avg_score : float + Average log likelihood + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + return _score(self.get_covariance(), X, weights=weights) + + return _score_sherman_woodbury( + self.get_noise(), self.W_, X, weights=weights + ) + + def score_samples( + self, X: NDArray[np.float64], sherman_woodbury: bool = False + ) -> NDArray[np.float64]: + """ + Return the log likelihood of each sample + + Parameters + ---------- + X : NDArray[np.float64],(N,sum(p)) + The data + + Returns + ------- + scores : np.float64 + Log likelihood for each sample + """ + if self.W_ is None: + raise ValueError("Model has not been fit yet") + + if sherman_woodbury is False: + return _score_samples(self.get_covariance(), X) + + return _score_samples_sherman_woodbury(self.get_noise(), self.W_, X) + + def get_entropy(self) -> np.float64: + """ + Computes the entropy of a Gaussian distribution parameterized by + covariance. + + Parameters + ---------- + None + + Returns + ------- + entropy : np.float64 + The differential entropy of the distribution + """ + return _entropy(self.get_covariance()) + + def get_entropy_subset( + self, observed_feature_idxs: NDArray[np.float64] + ) -> np.float64: + """ + Computes the entropy of a subset of the Gaussian distribution + parameterized by covariance. + + Parameters + ---------- + observed_feature_idxs: NDArray[np.float64],(sum(p),) + The observation locations + + Returns + ------- + entropy : np.float64 + The differential entropy of the distribution + """ + return _entropy_subset(self.get_covariance(), observed_feature_idxs) + + def mutual_information( + self, + observed_feature_idxs1: NDArray[np.float64], + observed_feature_idxs2: NDArray[np.float64], + ) -> np.float64: + """ + This computes the mutual information bewteen the two sets of + covariates based on the model. + + Parameters + ---------- + observed_feature_idxs1 : np.array-like,(p,) + First group of variables + + observed_feature_idxs2 : np.array-like,(p,) + Second group of variables + + Returns + ------- + mutual_information : np.float64 + The mutual information between the two variables + """ + covariance = self.get_covariance() + return _mutual_information( + covariance, observed_feature_idxs1, observed_feature_idxs2 + ) + + +class BasePCASGDModel(BaseGaussianFactorModel): + def __init__( + self, + n_components: int = 2, + training_options: dict[str, Any] | None = None, + prior_options: dict[str, Any] | None = None, + ) -> None: + """ + This is the base class of models that use stochastic + gradient descent for inference. This reduces boilerplate code + associated with some of the standard methods etc. + + 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 + + Sets + ---- + creationDate : datetime + The date/time that the object was created + """ + super().__init__(n_components=n_components) + + if training_options is None: + training_options = {} + if prior_options is None: + prior_options = {} + + self.training_options = self._fill_training_options(training_options) + self.prior_options = self._fill_prior_options(prior_options) + + def _fill_training_options( + self, training_options: dict[str, Any] + ) -> dict[str, Any]: + """ + This sets the default parameters for stochastic gradient descent, + our inference strategy for the model. + + Parameters + ---------- + training_options : dict + The original options set by the user passed as a dictionary + + Options + ------- + n_iterations : int, default=3000 + Number of iterations to train using stochastic gradient descent + + learning_rate : float, default=1e-4 + Learning rate of gradient descent + + method : string {'Nadam'}, default='Nadam' + The learning algorithm + + batch_size : int, default=None + The number of observations to use at each iteration. If none + corresponds to batch learning + + gpu_memory : int, default=1024 + The amount of memory you wish to use during training + """ + default_options = { + "n_iterations": 3000, + "learning_rate": 1e-2, + "use_gpu": True, + "method": "Nadam", + "batch_size": 100, + "momentum": 0.9, + } + + unexpected_but_present_keys = set(training_options.keys()) - set( + default_options.keys() + ) + if unexpected_but_present_keys: + raise ValueError( + "the following training options were unrecognized but provided..." + ) + + tops = {**default_options, **training_options} + return tops + + def _initialize_save_losses(self) -> None: + """ + This method initializes the arrays to track relevant variables + during training at each iteration + + Sets + ---- + losses_likelihood : np.array(n_iterations) + The log likelihood + + losses_prior : np.array(n_iterations) + The log prior + + losses_posterior : np.array(n_iterations) + The log posterior + """ + n_iterations = self.training_options["n_iterations"] + self.losses_likelihood = np.empty(n_iterations) + self.losses_prior = np.empty(n_iterations) + self.losses_posterior = np.empty(n_iterations) + + def _save_losses( + self, + i: int, + device: Any, + log_likelihood: Tensor, + log_prior: Tensor | NDArray[np.float64], + log_posterior: Tensor, + ) -> None: + """ + Saves the values of the losses at each iteration + + Parameters + ----------- + device ; pytorch.device + The device used for trainging (gpu or cpu) + + i : int + Current training iteration + + losses_likelihood : Tensor + The log likelihood + + losses_prior : Tensor | NDArray[np.float64] + The log prior + + losses_posterior : Tensor + The log posterior + """ + if device.type == "cuda": + self.losses_likelihood[i] = log_likelihood.detach().cpu().numpy() + if isinstance(log_prior, np.ndarray): + self.losses_prior[i] = log_prior + else: + self.losses_prior[i] = log_prior.detach().cpu().numpy() + self.losses_posterior[i] = log_posterior.detach().cpu().numpy() + else: + self.losses_likelihood[i] = log_likelihood.detach().numpy() + if isinstance(log_prior, np.ndarray): + self.losses_prior[i] = log_prior + else: + self.losses_prior[i] = log_prior.detach().numpy() + self.losses_posterior[i] = log_posterior.detach().numpy() + + def _transform_training_data( + self, device: Any, *args: NDArray + ) -> list[Tensor]: + """ + Convert a list of numpy arrays to tensors + """ + out = [] + for arg in args: + out.append(torch.tensor(arg.astype(np.float32)).to(device)) + return out + + @abstractmethod + def _create_prior(self, device: Any): + """ + Creates a prior on the parameters taking your trainable variable + dictionary as input + + Parameters + ---------- + None + + Returns + ------- + log_prior : function + The function representing the negative log density of the prior + """ diff --git a/gpcr/_base_covariance.py b/gpcr/_base_covariance.py new file mode 100644 index 0000000..d2780e2 --- /dev/null +++ b/gpcr/_base_covariance.py @@ -0,0 +1,1015 @@ +""" +This module provides a range of functions and methods for advanced statistical +analysis and manipulation of covariance matrices. It includes tools for matrix +symmetrization, precision and stable rank calculations, along with predictive +modeling and entropy calculations using covariance matrices. The core +functionalities are encapsulated within the BaseCovariance object, enhancing +ease of use and integration in statistical and machine learning workflows. + +Classes: +- BaseCovariance: Encapsulates covariance matrix operations. + Methods: + - __init__: Initialize with data validation. + - get_precision: Retrieve precision matrices. + - get_stable_rank: Calculate stable ranks. + - predict: Predict missing data values. + - conditional_score: Compute conditional scores for subsets. + - conditional_score_samples: Score individual samples conditionally. + - marginal_score: Calculate marginal scores for subsets. + - marginal_score_samples: Score individual samples marginally. + - score: General scoring function. + - score_samples: Score individual samples. + - entropy: Calculate entropy of the distribution. + - entropy_subset: Calculate entropy for subsets. + - mutual_information: Compute mutual information between sets. + +Standalone Functions: +- _symmeterize_and_warning: Symmetrize covariance matrices with warnings. +- ldet_sherman_woodbury_fa: Calculate log determinant using factor analysis. +- ldet_sherman_woodbury_full: Calculate log determinant of a complex matrix. +- inv_sherman_woodbury_fa: Invert a matrix with diagonal and low-rank structure. +- inv_sherman_woodbury_full: Invert a complex matrix structured as A + UBV. +""" + +from datetime import datetime as dt +from typing import Tuple + +import numpy as np +import pytz +from numpy import linalg as la +from numpy.typing import NDArray + + +def _symmeterize_and_warning(cov: np.ndarray) -> np.ndarray: + """ + Symmetrize the given square covariance matrix and warn if the symmetrization + introduces significant changes. + + Parameters + ---------- + cov : ndarray of shape (n_features, n_features) + The covariance matrix to be symmetrized. It must be a square matrix. + + Returns + ------- + cov_new : ndarray of shape (n_features, n_features) + The symmetric matrix obtained by averaging the matrix with its transpose. + + Raises + ------ + Warning + If the relative mean squared error between the original and symmetrized + covariance matrix exceeds 2%, indicating potential issues with the + symmetrization process. + """ + cov_new = (cov + cov.T) / 2 + diff_mse = np.mean((cov_new - cov) ** 2) + cov_mse = np.mean(cov**2) + if diff_mse / cov_mse > 0.02: + print( + "Warning: original matrix estimate deviated substantially from symmetric" + ) + return cov_new + + +class BaseCovariance: + def __init__(self) -> None: + self.creationDate = dt.now(pytz.timezone("US/Pacific")) + self.covariance: NDArray | None = None + + def get_precision(self) -> NDArray[np.float64]: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _get_precision(self.covariance) + + def get_stable_rank(self) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _get_stable_rank(self.covariance) + + def predict(self, Xobs: NDArray, idxs: NDArray): + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _predict(self.covariance, Xobs, idxs) + + def conditional_score( + self, X: NDArray, idxs: NDArray, weights=None + ) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _conditional_score(self.covariance, X, idxs, weights=weights) + + def conditional_score_samples( + self, X: NDArray, idxs: NDArray + ) -> NDArray[np.float64]: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _conditional_score_samples(self.covariance, X, idxs) + + def marginal_score( + self, X: NDArray, idxs: NDArray, weights=None + ) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _marginal_score(self.covariance, X, idxs, weights=weights) + + def marginal_score_samples( + self, X: NDArray, idxs: NDArray + ) -> NDArray[np.float64]: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _marginal_score_samples(self.covariance, X, idxs) + + def score(self, X: NDArray, weights=None) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _score(self.covariance, X, weights=weights) + + def score_samples(self, X: NDArray) -> NDArray[np.float64]: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _score_samples(self.covariance, X) + + def entropy(self) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _entropy(self.covariance) + + def entropy_subset(self, idxs: NDArray[np.float64]) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _entropy_subset(self.covariance, idxs) + + def mutual_information( + self, idxs1: NDArray[np.float64], idxs2: NDArray[np.float64] + ) -> np.float64: + if self.covariance is None: + raise ValueError("Covariance matrix has not been fit") + + return _mutual_information(self.covariance, idxs1, idxs2) + + def _test_inputs(self, X: NDArray[np.float64]) -> None: + """ + Just tests to make sure data is numpy array + """ + if not isinstance(X, np.ndarray): + raise ValueError("Data is not a numpy array") + if np.sum(np.isnan(X)) > 0: + raise ValueError("Data has nans") + + +def _get_precision(covariance: NDArray[np.float64]) -> NDArray[np.float64]: + """ + Gets the precision matrix defined as the inverse of the covariance + + Parameters + ---------- + covariance : NDArray,(p,p) + The covariance matrix + + Returns + ------- + precision : NDArray,(sum(p),sum(p)) + The inverse of the covariance matrix + """ + return la.inv(covariance) + + +def _get_stable_rank(covariance: NDArray[np.float64]) -> np.float64: + """ + Returns the stable rank defined as + ||A||_F^2/||A||^2 + + Parameters + ---------- + covariance : NDArray,(p,p) + The covariance matrix + + Returns + ------- + srank : np.float64 + The stable rank. See Vershynin High dimensional probability for + discussion, but this is a statistically stable approximation to rank + """ + singular_values = la.svd(covariance, compute_uv=False) + return np.sum(singular_values**2) / singular_values[0] ** 2 + + +def _predict( + covariance: NDArray[np.float64], + Xobs: NDArray[np.float64], + idxs: NDArray[np.float64], +) -> NDArray[np.float64]: + """ + Predicts missing data using observed data. This uses the conditional + Gaussian formula (see wikipedia multivariate gaussian). Solve allows + us to avoid an explicit matrix inversion. + + Parameters + ---------- + covariance : NDArray,(p,p) + The covariance matrix + + Xobs : NDArray,(N_samples,\\sum idxs) + The observed data + + idxs: NDArray,(sum(p),) + The observation locations + + Returns + ------- + preds : NDArray,(N_samples,p-\\sum idxs) + The predicted values + """ + covariance_22 = covariance[np.ix_(idxs == 1, idxs == 1)] + covariance_21 = covariance[np.ix_(idxs == 1, idxs == 0)] + beta_bar = la.solve(covariance_22, covariance_21) + + return np.dot(Xobs[:, idxs == 1], beta_bar) + + +def _conditional_score( + covariance: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, +) -> np.float64: + """ + Returns the predictive log-likelihood of a subset of data. + + mean(log p(X[idx==1]|X[idx==0],covariance)) + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + X : NDArray[np.float64],(N,sum(idxs)) + The centered data + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + weights : NDArray[np.float64],(N,),default=None + The optional weights on the samples. Don't have negative values. + Average value forced to 1. + + Returns + ------- + avg_score : np.float64 + Average log likelihood + """ + weights = ( + np.ones(X.shape[0]) if weights is None else weights / np.mean(weights) + ) + + return np.mean(weights * _conditional_score_samples(covariance, X, idxs)) + + +def _conditional_score_sherman_woodbury( + Lambda: NDArray[np.float64], + W: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, +) -> np.float64: + """ + Returns the predictive log-likelihood of a subset of data. + + mean(log p(X[idx==1]|X[idx==0],covariance)) + + Parameters + ---------- + Lambda : NDArray[np.float64],(p,p) + The diagonal noise matrix + + W : NDArray[np.float64],(L,p) + The low rank component + + X : NDArray[np.float64],(N,sum(idxs)) + The centered data + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + weights : NDArray[np.float64] | None,(N,),default=None + The optional weights on the samples. Don't have negative values. + Average value forced to 1. + + Returns + ------- + avg_score : np.float64 + Average log likelihood + """ + weights = ( + np.ones(X.shape[0]) if weights is None else weights / np.mean(weights) + ) + + return np.mean( + weights + * _conditional_score_samples_sherman_woodbury(Lambda, W, X, idxs) + ) + + +def _conditional_score_samples( + covariance: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], +) -> NDArray[np.float64]: + """ + Return the conditional log likelihood of each sample, that is + + log p(X[idx==1]|X[idx==0],covariance) = N(Sigma_10Sigma_00^{-1}x_0, + Sigma_11-Sigma_10Sigma_00^{-1}Sigma_01) + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + X : NDArray[np.float64],(N,p) + The centered data + + idxs: NDArray[np.float64],(p,) + The observation locations + + Returns + ------- + scores : np.float64 + Log likelihood for each sample + """ + covariance_22 = covariance[np.ix_(idxs == 1, idxs == 1)] + covariance_21 = covariance[np.ix_(idxs == 1, idxs == 0)] + covariance_11 = covariance[np.ix_(idxs == 0, idxs == 0)] + + beta_bar = la.solve(covariance_22, covariance_21) + Second_part = np.dot(covariance_21.T, beta_bar) + + covariance_bar = covariance_11 - Second_part + + mu_ = np.dot(X[:, idxs == 1], beta_bar) + + return _score_samples(covariance_bar, X[:, idxs == 0] - mu_) + + +def _conditional_score_samples_sherman_woodbury( + Lambda: NDArray[np.float64], + W: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], +) -> NDArray[np.float64]: + """ + Return the conditional log likelihood of each sample, that is + + log p(X[idx==1]|X[idx==0],covariance) = N(Sigma_10Sigma_00^{-1}x_0, + Sigma_11-Sigma_10Sigma_00^{-1}Sigma_01) + + Parameters + ---------- + Lambda : NDArray[np.float64],(p,p) + The diagonal noise matrix + + W : NDArray[np.float64],(L,p) + The low rank component + + X : NDArray[np.float64],(N,p) + The centered data + + idxs: NDArray[np.float64],(p,) + The observation locations + + Returns + ------- + scores : np.float64 + Log likelihood for each sample + """ + I_L = np.eye(W.shape[0]) + X_obs = X[:, idxs == 1] + X_miss = X[:, idxs == 0] + W_obs = W[:, idxs == 1] + W_miss = W[:, idxs == 0] + Lo = Lambda[np.ix_(idxs == 1, idxs == 1)] + Lm = Lambda[np.ix_(idxs == 0, idxs == 0)] + + C = np.dot(W_obs, la.inv(Lo)) + B = np.dot(C, W_obs.T) + + IpB = I_L + B + back = la.solve(IpB, C) + second = C - np.dot(B, back) + coef = np.dot(W_miss.T, second) + + mu_ = np.dot(X_obs, coef.T) + middle = B - np.dot(B, la.solve(IpB, B)) + term2 = np.dot(W_miss.T, np.dot(middle, W_miss)) + + covariance11 = Lm + np.dot(W_miss.T, W_miss) + covariance_bar = covariance11 - term2 + + return _score_samples(covariance_bar, X_miss - mu_) + + +def _get_conditional_parameters( + covariance: NDArray[np.float64], idxs: NDArray[np.float64] +) -> Tuple[NDArray[np.float64], NDArray[np.float64]]: + """ + Computes the distribution parameters p(X_miss|X_obs) + + Parameters + ---------- + covariance : NDArray,(p,p) + The covariance matrix + + idxs: NDArray,(p,) + The observed covariates + + Returns + ------- + beta_bar : array-like + The predictive covariates + + covariance_bar : array-like + Conditional covariance + """ + covariance_sub = covariance[idxs == 1] + covariance_22 = covariance_sub[:, idxs == 1] + covariance_21 = covariance_sub[:, idxs == 0] + covariance_nonsub = covariance[idxs == 0] + covariance_11 = covariance_nonsub[:, idxs == 0] + beta_bar = la.solve(covariance_22, covariance_21) + Second_part = np.dot(covariance_21.T, beta_bar) + + covariance_bar = covariance_11 - Second_part + + return beta_bar.T, covariance_bar + + +def _get_conditional_parameters_sherman_woodbury( + Lambda: NDArray[np.float64], + W: NDArray[np.float64], + idxs: NDArray[np.float64], +) -> Tuple[NDArray[np.float64], NDArray[np.float64]]: + """ + Computes the distribution parameters p(X_miss|X_obs) + given that Sigma = WWT + Lambda + + Parameters + ---------- + Lambda : NDArray,(p,p) + The diagonal noise matrix + + W : NDArray,(L,p) + The low rank component + + idxs: NDArray,(p,) + The observed covariates + + Returns + ------- + beta_bar : NDArray,(p,) + The predictive covariates + + covariance_bar : NDArray,(p,p) + Conditional covariance + """ + I_L = np.eye(W.shape[0]) + W_obs = W[:, idxs == 1] + W_miss = W[:, idxs == 0] + Lo = Lambda[np.ix_(idxs == 1, idxs == 1)] + Lm = Lambda[np.ix_(idxs == 0, idxs == 0)] + + C = np.dot(W_obs, la.inv(Lo)) + B = np.dot(C, W_obs.T) + + IpB = I_L + B + back = la.solve(IpB, C) + second = C - np.dot(B, back) + coef = np.dot(W_miss.T, second) + + middle = B - np.dot(B, la.solve(IpB, B)) + term2 = np.dot(W_miss.T, np.dot(middle, W_miss)) + + covariance11 = Lm + np.dot(W_miss.T, W_miss) + covariance_bar = covariance11 - term2 + + return coef, covariance_bar + + +def _marginal_score( + covariance: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, +) -> np.float64: + """ + Returns the marginal log-likelihood of a subset of data + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + X : NDArray[np.float64],(N,sum(idxs)) + The centered data + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + weights : NDArray[np.float64],(N,),default=None + The optional weights on the samples + + Returns + ------- + avg_score : np.float64 + Average log likelihood + """ + weights = ( + np.ones(X.shape[0]) if weights is None else weights / np.mean(weights) + ) + + return np.mean(weights * _marginal_score_samples(covariance, X, idxs)) + + +def _marginal_score_sherman_woodbury( + Lambda: NDArray[np.float64], + W: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, +) -> np.float64: + """ + Returns the marginal log-likelihood of a subset of data + given that Sigma = WWT + Lambda + + Parameters + ---------- + Lambda : NDArray[np.float64],(p,p) + The diagonal noise matrix + + W : NDArray[np.float64],(L,p) + The low rank component + + X : NDArray[np.float64],(N,sum(idxs)) + The centered data + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + weights : NDArray[np.float64] | None,(N,),default=None + The optional weights on the samples + + Returns + ------- + avg_score : np.float64 + Average log likelihood + """ + if weights is None: + weights = np.ones(X.shape[0]) + + return np.mean( + weights * _marginal_score_samples_sherman_woodbury(Lambda, W, X, idxs) + ) + + +def _marginal_score_samples( + covariance: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], +) -> NDArray[np.float64]: + """ + Returns the marginal log-likelihood of a subset of data + per window + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + X : NDArray[np.float64],(N,sum(idxs)) + The centered data + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + Returns + ------- + scores : np.float64 + Average log likelihood + """ + cov1 = covariance[idxs == 1] + cov_sub = cov1[:, idxs == 1] + + return _score_samples(cov_sub, X) + + +def _marginal_score_samples_sherman_woodbury( + Lambda: NDArray[np.float64], + W: NDArray[np.float64], + X: NDArray[np.float64], + idxs: NDArray[np.float64], +) -> NDArray[np.float64]: + """ + Returns the marginal log-likelihood of a subset of data + per window given that Sigma = WWT + Lambda + + Parameters + ---------- + Lambda : NDArray[np.float64],(p,p) + The diagonal noise matrix + + W : NDArray[np.float64],(L,p) + The low rank component + + X : NDArray[np.float64],(N,sum(idxs)) + The centered data + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + Returns + ------- + scores : np.float64 + Average log likelihood + """ + Lambda_sub = Lambda[np.ix_(idxs == 1, idxs == 1)] + W_sub = W[:, idxs == 1] + + return _score_samples_sherman_woodbury(Lambda_sub, W_sub, X) + + +def _score( + covariance: NDArray[np.float64], + X: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, +) -> np.float64: + """ + Returns the average log liklihood of data. + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + X : NDArray[np.float64],(N,sum(p)) + The centered data + + weights : NDArray[np.float64] | None,(N,),default=None + The optional weights on the samples + + Returns + ------- + avg_score : np.float64 + Average log likelihood + """ + weights = ( + np.ones(X.shape[0]) if weights is None else weights / np.mean(weights) + ) + + return np.mean(weights * _score_samples(covariance, X)) + + +def _score_sherman_woodbury( + Lambda: NDArray[np.float64], + W: NDArray[np.float64], + X: NDArray[np.float64], + weights: NDArray[np.float64] | None = None, +) -> np.float64: + """ + Calculate the average log likelihood of the data given a specific covariance structure. + + The covariance structure is assumed to be Sigma = WW^T + Lambda, where + WW^T is the low-rank component and Lambda is a diagonal matrix (noise). + + Parameters + ---------- + Lambda : NDArray[np.float64] + The diagonal noise matrix, with shape (p, p), where 'p' represents the number of features. + + W : NDArray[np.float64] + The low rank component matrix, with shape (L, p), where 'L' represents the number of factors + and 'p' is the number of features. + + X : NDArray[np.float64] + The centered data matrix, with shape (N, p), where 'N' is the number of samples and 'p' is + the number of features. It is assumed that the data has already been centered. + + weights : NDArray[np.float64] | None, optional + The weights for the samples, with shape (N,), where 'N' is the number of samples. If None, + all samples are weighted equally. Default is None. + + Returns + ------- + avg_score : np.float64 + The average log likelihood of the data under the given model parameters. + + Notes + ----- + The function _score_samples_sherman_woodbury should be defined elsewhere and should compute + the log likelihood of the samples given the model parameters. + + The likelihood is weighted by the `weights` parameter if provided; otherwise, each sample's + likelihood contributes equally to the average. + """ + if weights is None: + weights = np.ones(X.shape[0]) + + return np.mean(weights * _score_samples_sherman_woodbury(Lambda, W, X)) + + +def _score_samples( + covariance: NDArray[np.float64], X: NDArray[np.float64] +) -> NDArray[np.float64]: + """ + Return the log likelihood of each sample + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + X : NDArray[np.float64],(N,sum(p)) + The centered data + + Returns + ------- + scores : np.float64 + Log likelihood for each sample + """ + p = covariance.shape[0] + term1 = -p / 2 * np.log(2 * np.pi) + + _, logdet = la.slogdet(covariance) + term2 = -0.5 * logdet + + quad_init = la.solve(covariance, np.transpose(X)) + difference = X * np.transpose(quad_init) + term3 = np.sum(difference, axis=1) + + return term1 + term2 - 0.5 * term3 + + +def _score_samples_sherman_woodbury( + Lambda: NDArray[np.float64], W: NDArray[np.float64], X: NDArray[np.float64] +) -> NDArray[np.float64]: + """ + Return the log likelihood of each sample + + Parameters + ---------- + Lambda : NDArray[np.float64],(p,p) + The diagonal noise matrix + + W : NDArray[np.float64],(L,p) + The low rank component + + X : NDArray[np.float64],(N,sum(p)) + The centered data + + Returns + ------- + scores : np.float64 + Log likelihood for each sample + """ + p = Lambda.shape[1] + I_L = np.eye(W.shape[0]) + term1 = -p / 2 * np.log(2 * np.pi) + term2 = -0.5 * ldet_sherman_woodbury_fa(Lambda, W) + + Li = la.inv(Lambda) + C = np.dot(Li, W.T) + CtW = np.dot(C.T, W.T) + middle = I_L + CtW + end = la.solve(middle, C.T) # (I + WLiW^T)^{-1}WLi + prod = np.dot(C, end) + Lip = Li - prod + quad_init = np.dot(Lip, X.T) + difference = X * np.transpose(quad_init) + term3 = np.sum(difference, axis=1) + + return term1 + term2 - 0.5 * term3 + + +def _entropy(covariance: NDArray[np.float64]) -> np.float64: + """ + Computes the entropy of a Gaussian distribution parameterized by + covariance. + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + Returns + ------- + entropy : np.float64 + The differential entropy of the distribution + """ + cov_new = 2 * np.pi * np.e * covariance + _, logdet = la.slogdet(cov_new) + + return 0.5 * logdet + + +def _entropy_subset( + covariance: NDArray[np.float64], idxs: NDArray[np.float64] +) -> np.float64: + """ + Computes the entropy of a subset of the Gaussian distribution + parameterized by covariance. + + Parameters + ---------- + covariance : NDArray[np.float64],(p,p) + The covariance matrix + + idxs: NDArray[np.float64],(sum(p),) + The observation locations + + Returns + ------- + entropy : np.float64 + The differential entropy of the distribution + """ + cov1 = covariance[idxs == 1] + cov_sub = cov1[:, idxs == 1] + entropy = _entropy(cov_sub) + + return entropy + + +def _mutual_information( + covariance: NDArray[np.float64], + idxs1: NDArray[np.float64], + idxs2: NDArray[np.float64], +) -> np.float64: + """ + This computes the mutual information bewteen the two sets of + covariates based on the model. + + Parameters + ---------- + idxs1 : NDArray[np.float64],(p,) + First group of variables + + idxs2 : NDArray[np.float64],(p,) + Second group of variables + + Returns + ------- + mutual_information : np.float64 + The mutual information between the two variables + """ + idxs = np.logical_or(idxs1, idxs2).astype(int) + cov_sub = covariance[np.ix_(idxs == 1, idxs == 1)] + idxs1_sub = idxs1[idxs == 1] + Hy = _entropy(cov_sub[np.ix_(idxs1_sub == 1, idxs1_sub == 1)]) + + _, covariance_conditional = _get_conditional_parameters( + cov_sub, 1 - idxs1_sub + ) + H_y_given_x = _entropy(covariance_conditional) + mutual_information = Hy - H_y_given_x + + return mutual_information + + +def ldet_sherman_woodbury_fa(Lambda: NDArray, W: NDArray) -> float: + """ + This converts the log determinant of a matrix Lambda + W^TW where + Lambda is diagonal. Fa for factor analysis + + Parameters + ---------- + W : np.array(n_components,p) + The PCA loadings matrix + + Lambda : NDArray,(p,p) + An easy to invert (diagonal) noise matrix + + Returns + ------- + log_determinant : float + The log determinant of the covariance matrix + """ + _, ldetL = la.slogdet(Lambda) + LiW = la.solve(Lambda, W.T) + WtLiW = np.dot(W, LiW) + IWtLiW = np.eye(W.shape[0]) + WtLiW + _, ldetP = la.slogdet(IWtLiW) + log_determinant = ldetP + ldetL + return log_determinant + + +def ldet_sherman_woodbury_full( + A: NDArray, U: NDArray, B: NDArray, V: NDArray +) -> float: + """ + The value log|A+ UBV|. Useful when A and B are simple to invert and UBV + is low rank + + Parameters + ---------- + A : NDArray,(p,p) + The first square matrix + + U : NDArray,(p,L) + The wide matrix + + B : NDArray,(L,L) + The tiny square matrix + + V : NDArray,(L,p) + Tall matrix + + Returns + ------- + log_determinant : float + The log determinant of the covariance matrix + """ + _, ldetA = la.slogdet(A) + _, ldetB = la.slogdet(B) + term2 = np.dot(V, la.solve(A, U)) + term1 = la.inv(B) + _, ldetProd = la.slogdet(term1 + term2) + log_determinant = ldetA + ldetB + ldetProd + return log_determinant + + +def inv_sherman_woodbury_fa(Lambda: NDArray, W: NDArray) -> NDArray[np.float64]: + """ + This converts the inverse of a matrix Lambda + W^TW where + Lambda is diagonal. Fa for factor analysis + + Parameters + ---------- + W : np.array(n_components,p) + The PCA loadings matrix + + Lambda : NDArray,(p,p) + An easy to invert (diagonal) noise matrix + + Returns + ------- + Sigma_inv : NDArray,(p,p) + The inverse of the covariance matrix + """ + I_L = np.eye(W.shape[0]) + I_p = np.eye(W.shape[1]) + Lambda_inv = la.inv(Lambda) + WLi = np.dot(W, Lambda_inv) + inner = I_L + np.dot(WLi, W.T) + inner_inv = la.inv(inner) + end = np.dot(inner_inv, WLi) + term2 = np.dot(W.T, end) + Imterm2 = I_p - term2 + Sigma_inv = np.dot(Lambda_inv, Imterm2) + return Sigma_inv + + +def inv_sherman_woodbury_full( + A: NDArray, U: NDArray, B: NDArray, V: NDArray +) -> NDArray[np.float64]: + """ + This converts the inverse of a matrix (A + UBV) + + Parameters + ---------- + A : NDArray,(p,p) + The first square matrix + + U : NDArray,(p,L) + The wide matrix + + B : NDArray,(L,L) + The tiny square matrix + + V : NDArray,(L,p) + Tall matrix + + Returns + ------- + Sigma_inv : NDArray,(p,p) + The inverse of the matrix (A + UBV) + """ + Ainv = la.inv(A) # Needed explicitly anyways + AiU = np.dot(Ainv, U) + Binv = la.inv(B) # Should be easy to invert + VAinv = np.dot(V, Ainv) + VAiU = np.dot(VAinv, U) + middle = Binv + VAiU + end = la.solve(middle, VAinv) + second_term = np.dot(AiU, end) + first_term = Ainv + Sigma_inv = first_term - second_term + return Sigma_inv diff --git a/gpcr/_covariance_np.py b/gpcr/_covariance_np.py new file mode 100644 index 0000000..7df4718 --- /dev/null +++ b/gpcr/_covariance_np.py @@ -0,0 +1,286 @@ +""" +This module provides classes for fitting covariance matrices using various +approaches, including empirical, Bayesian, linear shrinkage, and nonlinear +shrinkage methods. These classes extend BaseCovariance, handling missing data +checks and employing advanced statistical methods for optimal estimation. + +Classes: +- EmpiricalCovariance: Fits standard sample covariance. +- BayesianCovariance: Implements MAP estimation with priors. +- LinearShrinkageCovariance: Combines empirical covariance with a structured + estimator. +- NonLinearShrinkageCovariance: Adjusts eigenvalues using a nonlinear + shrinkage based on sample spectral density. + +Dependencies: +- numpy for matrix operations. +- Preprocessed input data as centered, non-missing values. +""" + +from typing import Any + +import numpy as np +from ._base_covariance import ( + BaseCovariance, + _symmeterize_and_warning, +) +from numpy.typing import NDArray + + +class EmpiricalCovariance(BaseCovariance): + def __init__(self): + """ + This object just fits the covariance matrix as the standard sample + covariance matrix. Does not handle missing values + """ + super().__init__() + + def fit(self, X: NDArray) -> "EmpiricalCovariance": + """ + This fits a covariance matrix using samples X. + + Parameters + ---------- + X : np.array-like,(n_samples,n_covariates) + The centered data + + Returns + ------- + self : EmpiricalCovariance + The model + + Raises + ------ + ValueError: + A value error will be raised if missing data is found in X ( + np.isnan(X) evaluates to true ), or if X is not an NDArray + """ + self._test_inputs(X) + self.N, self.p = X.shape + XTX = np.dot(X.T, X) + covariance = XTX / self.N + self.covariance = _symmeterize_and_warning(covariance) + + return self + + +class BayesianCovariance(BaseCovariance): + def __init__(self, prior_options=None): + """ + This object fits the covariance matrix as the MAP estimator using + user-defined priors. Does not handle missing values + """ + super().__init__() + if prior_options is None: + prior_options = {} + + self.prior_options = self._fill_prior_options(prior_options) + + def fit(self, X: NDArray[np.float_]) -> "BayesianCovariance": + """ + This fits a covariance matrix using samples X with MAP estimation. + + Parameters + ---------- + X : np.array-like,(n_samples,n_covariates) + The data + + Returns + ------- + self : BayesianCovariance + The model + + Raises + ------ + ValueError: + A value error will be raised if missing data is found in X ( + np.isnan(X) evaluates to true ), or if X is not an NDArray + """ + self._test_inputs(X) + self.N, self.p = X.shape + + p_opts = self.prior_options + covariance_empirical = np.dot(X.T, X) + nu = p_opts["iw_params"]["pnu"] + self.p + cov_prior = p_opts["iw_params"]["sigma"] * np.eye(self.p) + posterior_cov = cov_prior + covariance_empirical + posterior_nu = nu + self.N + + covariance = posterior_cov / (posterior_nu + self.p + 1) + self.covariance = _symmeterize_and_warning(covariance) + + return self + + def _fill_prior_options( + self, prior_options: dict[str, Any] + ) -> dict[str, Any]: + """ + This sets the prior options for our inference scheme + + Parameters + ---------- + prior_options : dict + The original prior options passed as a dictionary + + Options + ------- + iw_params : dict,default={'pnu':2,'sigma':1.0} + pnu : int - nu = p + pnu + sigma : float>0 - cov = sigma*I_p + """ + default_options = { + "iw_params": {"pnu": 2, "sigma": 1.0}, + } + return {**default_options, **prior_options} + + +class LinearShrinkageCovariance(BaseCovariance): + def fit(self, X: NDArray[np.float_]) -> "LinearShrinkageCovariance": + """ + This fits a covariance matrix using the linear shrinkage approach + to combine the empirical covariance matrix with a structured + estimator. This uses the 2004 paper from Ledoit and Wolf + + Parameters + ---------- + X : np.array-like, (n_samples, n_covariates) + The centered data + + Returns + ------- + self : LinearShrinkageCovariance + The model, with updated covariance attribute + + Raises + ------ + ValueError: + A value error will be raised if missing data is found in X ( + np.isnan(X) evaluates to true ), or if X is not an NDArray + """ + + self._test_inputs(X) + self.N, self.p = X.shape + + S = np.cov(X.T) + + # Compute shrinkage intensity + evalues = np.linalg.eigvalsh(S) + lambd = np.mean(evalues) + + temp = [ + np.linalg.norm(np.outer(X[i], X[i]) - S, "fro") ** 2 + for i in range(self.N) + ] + + b_bar_sq = np.sum(temp) / self.N**2 + d_sq = np.linalg.norm(S - lambd * np.eye(self.p), "fro") ** 2 + b_sq = np.minimum(b_bar_sq, d_sq) + rho_hat = b_sq / d_sq + + # Shrink covariance matrix + lambda_bar = np.trace(S) / self.p + covariance = rho_hat * lambda_bar * np.eye(self.p) + (1 - rho_hat) * S + self.covariance = _symmeterize_and_warning(covariance) + + return self + + +class NonLinearShrinkageCovariance(BaseCovariance): + """ + This method adjusts each eigenvalue of the sample covariance matrix + individually according to a nonlinear shrinkage formula, based on the + Hilbert transform. This approach contrasts with linear shrinkage methods + that apply a uniform adjustment across all eigenvalues. The nonlinear + technique allows for a more nuanced adjustment that is particularly + beneficial in scenarios with clusters of eigenvalues, optimizing the + shrinkage based on local eigenvalue densities. The method assumes a + sample matrix composed of i.i.d. random variables with mean zero and + finite 16th moment, and is suitable for large-dimensional data sets where + the ratio of variables p to n is significant. + + The theoretical basis of this approach is built around the optimal + estimation of p parameters, which balances between overfitting p^2 + parameters and underfitting with only 1 parameter. The nonlinear shrinkage + is calculated using an oracle estimator that depends on a sample + spectral density + function approximated via kernel estimation using the Epanechnikov kernel. + This estimation facilitates the derivation of shrinkage intensities for each + eigenvalue, tailoring the adjustment to improve estimations under various + data conditions. + + https://www.jstor.org/stable/27028732 + """ + + def fit(self, X: NDArray[np.float_]) -> "NonLinearShrinkageCovariance": + """ + This fits a covariance matrix using the nonlinear shrinkage approach, + which adjusts the empirical eigenvalues based on asymptotic results. + + Parameters + ---------- + X : np.array-like, (n_samples, n_covariates) + The centered data + + Returns + ------- + self : NonLinearShrinkageCovariance + The model, with updated covariance attribute + + Raises + ------ + ValueError: + A value error will be raised if missing data is found in X ( + np.isnan(X) evaluates to true ), or if X is not an NDArray + """ + self._test_inputs(X) + self.N, self.p = X.shape + N, p = X.shape + S = np.cov(X.T) + lambd, u = np.linalg.eigh(S) + + lambd = lambd[np.maximum(0, p - N) :] + L = np.tile(lambd, (np.minimum(p, N), 1)).T + h = N ** (-1 / 3) + + H = h * L.T + x = (L - L.T) / H + ftilde = 3 / 4 / np.sqrt(5) + ftilde *= np.mean(np.maximum(1 - x**2 / 5, 0) / H, axis=1) + + Hftemp = (-3 / 10 / np.pi) * x + (3 / 4 / np.sqrt(5) / np.pi) * ( + 1 - x**2 / 5 + ) * np.log(np.abs((np.sqrt(5) - x) / (np.sqrt(5) + x))) + + Hftemp[np.abs(x) == np.sqrt(5)] = (-3 / 10 / np.pi) * x[ + np.abs(x) == np.sqrt(5) + ] + + Hftilde = np.mean(Hftemp / H, axis=1) + + if p <= N: + dtilde = lambd / ( + (np.pi * (p / N) * lambd * ftilde) ** 2 + + (1 - (p / N) - np.pi * (p / N) * lambd * Hftilde) ** 2 + ) + else: + Hftilde0 = ( + (1 / np.pi) + * ( + 3 / 10 / h**2 + + 3 + / 4 + / np.sqrt(5) + / h + * (1 - 1 / 5 / h**2) + * np.log((1 + np.sqrt(5) * h) / (1 - np.sqrt(5) * h)) + ) + * np.mean(1 / lambd) + ) + dtilde0 = 1 / (np.pi * (p - N) / N * Hftilde0) + dtilde1 = lambd / (np.pi**2 * lambd**2 * (ftilde**2 + Hftilde**2)) + dtilde = np.concatenate([dtilde0 * np.ones((p - N)), dtilde1]) + + covariance = np.dot(np.dot(u, np.diag(dtilde)), u.T) + self.covariance = _symmeterize_and_warning(covariance) + + return self From 51f2699b893e2a06f3c286c25bc7172edfa00151 Mon Sep 17 00:00:00 2001 From: Lobo Date: Tue, 16 Jun 2026 16:10:31 -0400 Subject: [PATCH 2/3] Fixed black --- gpcr/_covariance_np.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gpcr/_covariance_np.py b/gpcr/_covariance_np.py index 7df4718..248a4b2 100644 --- a/gpcr/_covariance_np.py +++ b/gpcr/_covariance_np.py @@ -20,11 +20,12 @@ from typing import Any import numpy as np +from numpy.typing import NDArray + from ._base_covariance import ( BaseCovariance, _symmeterize_and_warning, ) -from numpy.typing import NDArray class EmpiricalCovariance(BaseCovariance): From ad8ce34e5891703ee93ffb5daeb1ad68bc78f668 Mon Sep 17 00:00:00 2001 From: Lobo Date: Tue, 16 Jun 2026 16:14:13 -0400 Subject: [PATCH 3/3] Fixed black version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 918373f..927fab4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,7 +28,7 @@ dev = [ "ruff>=0.6.9", "pytest>=8.3", "pytest-cov>=5.0", - "black==23.12.1", + "black==24.2.0", ] [project.urls]