diff --git a/gpcr/_base.py b/gpcr/_base.py index 9e61c31..6911e73 100644 --- a/gpcr/_base.py +++ b/gpcr/_base.py @@ -63,28 +63,30 @@ 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. + Compute the posterior parameters of the latent distribution p(S|X). - Computes the parameters for p(S|X) - - Description in future to be released paper + Currently implemented only for PPCA due to its closed-form expression. + Computes the projection matrix and posterior covariance via the matrix + M = W W^T + sigma * I. Parameters ---------- - W_ : Tensor(n_components,p) - The loadings + W_ : Tensor of shape (n_components, p) + The factor loadings matrix. sigma_ : Tensor - Isotropic noise + Scalar isotropic noise variance. + + device : Any + The PyTorch device (CPU or CUDA) on which to perform computations. Returns ------- - Proj_X : Tensor(n_components,p) - Beta such that np.dot(Proj_X, X) = E[S|X] + Proj_X : Tensor of shape (n_components, p) + Projection matrix such that ``Proj_X @ X`` gives ``E[S|X]``. - Cov : Tensor(n_components,n_components) - Var(S|X) + Cov : Tensor of shape (n_components, n_components) + Posterior covariance ``Var(S|X)``. """ n_components = int(W_.shape[0]) eye = torch.tensor(np.eye(n_components).astype(np.float32), device=device) @@ -101,30 +103,28 @@ def kl_divergence_vae( 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,). + Compute the KL divergence from N(mu1, diag(sigma1)) to N(0, I). + + Computes ``KL(N(mu1, sigma1) || N(0, I))`` using the closed-form + expression for diagonal Gaussian distributions:: + + KL = 0.5 * (tr(sigma1) + mu1^T mu1 - log(det(sigma1)) - d) + + where ``d`` is the dimensionality of the latent space. + + Parameters + ---------- + mu1 : torch.Tensor of shape (batch_size, d) + Mean vector of the approximate posterior distribution. + + sigma1 : torch.Tensor of shape (batch_size, d) + Diagonal of the covariance matrix of the approximate posterior. + Expected to contain positive values. + + Returns + ------- + kl_div : torch.Tensor of shape (batch_size,) + Per-sample KL divergence values. """ d = sigma1.shape[1] # Dimensionality of the latent space term1 = -torch.sum( @@ -142,17 +142,12 @@ def kl_divergence_vae( class BaseGaussianFactorModel(ABC): def __init__(self, n_components=2): """ - This is the base class of the model. Will never be called directly. + Base class for Gaussian factor models. Not intended to be instantiated directly. Parameters ---------- - n_components : int,default=2 - The latent dimensionality - - Sets - ---- - creationDate : datetime - The date/time that the object was created + n_components : int, default=2 + The dimensionality of the latent space. """ self.n_components = int(n_components) self.W_ = None @@ -160,66 +155,68 @@ def __init__(self, n_components=2): @abstractmethod def fit(self, *args, **kwargs): """ - Fits a model given covariates X as well as option labels y in the - supervised methods + Fit the model to data. + + Subclasses must implement this method. Supervised subclasses + additionally accept a label array ``y``. Parameters ---------- - X : np.array-like,(n_samples,n_covariates) - The data + X : array-like of shape (n_samples, n_features) + Training data matrix. + + *args : additional positional arguments + Passed to the subclass implementation. - other arguments + **kwargs : additional keyword arguments + Passed to the subclass implementation. Returns ------- self : object - The model + Fitted estimator. """ @abstractmethod def get_covariance(self) -> NDArray[np.float64]: """ - Gets the covariance matrix defined by the model parameters - - Parameters - ---------- - None + Compute the covariance matrix implied by the fitted model parameters. Returns ------- - covariance : np.array-like(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The model covariance matrix. """ @abstractmethod def get_noise(self) -> NDArray[np.float64]: """ - Returns the observational noise as a diagnoal matrix - - Parameters - ---------- - None + Return the observational noise as a diagonal matrix. Returns ------- - Lambda : np.array-like(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + Diagonal noise matrix. """ def get_precision( self, sherman_woodbury: bool = False ) -> NDArray[np.float64]: """ - Gets the precision matrix defined as the inverse of the covariance + Compute the precision matrix (inverse of the covariance matrix). Parameters ---------- - None + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury matrix identity for + efficient inversion, exploiting the low-rank-plus-diagonal + structure of the covariance. Requires ``self.get_noise()`` to + return a diagonal matrix. Returns ------- - precision : np.array-like(p,p) - The inverse of the covariance matrix + precision : ndarray of shape (n_features, n_features) + The precision matrix, i.e. the inverse of the covariance matrix. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -232,18 +229,17 @@ def get_precision( def get_stable_rank(self) -> np.float64: """ - Returns the stable rank defined as - ||A||_F^2/||A||^2 + Compute the stable rank of the covariance matrix. - Parameters - ---------- - None + The stable rank is defined as ``||A||_F^2 / ||A||_2^2``, i.e. the + ratio of the squared Frobenius norm to the squared spectral norm. + This provides a statistically robust approximation to the matrix rank. + See Vershynin, *High-Dimensional Probability* for a detailed discussion. Returns ------- srank : np.float64 - The stable rank. See Vershynin High dimensional probability for - discussion, but this is a statistically stable approximation to rank + The stable rank of the covariance matrix. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -252,17 +248,24 @@ def get_stable_rank(self) -> np.float64: def transform(self, X: NDArray[np.float64], sherman_woodbury: bool = False): """ - This returns the latent variable estimates given X + Project data into the latent factor space. + + Computes posterior mean estimates of the latent variables given + the observations ``X``. Parameters ---------- - X : np array-like,(N_samples,p) - The data to transform. + X : ndarray of shape (n_samples, n_features) + Data matrix to transform. + + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity to exploit the + low-rank-plus-diagonal structure for efficient computation. Returns ------- - S : NDArray,(N_samples,n_components) - The factor estimates + S : ndarray of shape (n_samples, n_components) + Estimated latent factor scores. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -286,21 +289,29 @@ def transform_subset( sherman_woodbury: bool = False, ) -> NDArray[np.float64]: """ - This returns the latent variable estimates given partial observations - contained in X + Project partially-observed data into the latent factor space. + + Computes latent variable estimates using only the observed features + indicated by ``observed_feature_idxs``. Parameters ---------- - X : NDArray,(N_samples,sum(observed_feature_idxs)) - The data to transform. + X : ndarray of shape (n_samples, n_observed) + Data matrix containing only the observed features, where + ``n_observed = observed_feature_idxs.sum()``. - observed_feature_idxs: np.array-like,(sum(p),) - The observation locations + observed_feature_idxs : ndarray of shape (n_features,) + Boolean or integer index array indicating which features are + observed (value 1) versus missing (value 0). + + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity for efficient + computation under low-rank-plus-diagonal covariance structure. Returns ------- - S : NDArray,(N_samples,n_components) - The factor estimates + S : ndarray of shape (n_samples, n_components) + Estimated latent factor scores. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -326,25 +337,33 @@ def conditional_score( sherman_woodbury: bool = False, ) -> np.float64: """ - Returns the predictive log-likelihood of a subset of data. + Compute the weighted average conditional log-likelihood. - mean(log p(X[idx==1]|X[idx==0],covariance)) + Evaluates ``mean(log p(X[idx==0] | X[idx==1], covariance))`` + over all samples. Parameters ---------- - X : NDArray,(N,sum(observed_feature_idxs)) - The data + X : ndarray of shape (n_samples, n_features) + The centered data matrix (all features, both observed and + unobserved). + + observed_feature_idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 indicates an observed + feature and 0 indicates a missing feature. - observed_feature_idxs: NDArray,(sum(p),) - The observation locations + weights : ndarray of shape (n_samples,), default=None + Optional sample weights. If ``None``, all samples are weighted + equally. Weights are normalized to have mean 1. - weights : NDArray,(N,),default=None - The optional weights on the samples + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity for efficient + computation under low-rank-plus-diagonal covariance structure. Returns ------- - avg_score : float - Average log likelihood + avg_score : np.float64 + Weighted average conditional log-likelihood. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -370,25 +389,28 @@ def conditional_score_samples( sherman_woodbury: bool = False, ) -> NDArray[np.float64]: """ - Return the conditional log likelihood of each sample, that is + Compute the per-sample conditional log-likelihood. - log p(X[idx==1]|X[idx==0],covariance) + Evaluates ``log p(X[idx==0] | X[idx==1], covariance)`` for each + sample individually. Parameters ---------- - X : np.array-like,(N,p) - The data + X : ndarray of shape (n_samples, n_features) + The centered data matrix (all features). - observed_feature_idxs: np.array-like,(p,) - The observation locations + observed_feature_idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 indicates an observed + feature and 0 indicates a missing feature. - sherman_woodbury : bool,default=False - Whether to use the sherman_woodbury matrix identity + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury matrix identity for + efficient computation under low-rank-plus-diagonal covariance. Returns ------- - scores : float - Log likelihood for each sample + scores : ndarray of shape (n_samples,) + Per-sample conditional log-likelihood values. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -411,23 +433,34 @@ def marginal_score( sherman_woodbury: bool = False, ) -> np.float64: """ - Returns the marginal log-likelihood of a subset of data + Compute the weighted average marginal log-likelihood of a feature subset. + + Evaluates ``mean(log p(X[:, idx==1], covariance_sub))`` over all + samples, where the marginal distribution is derived from the model + covariance restricted to the observed features. Parameters ---------- - X : np.array-like,(N,sum(idxs)) - The data + X : ndarray of shape (n_samples, n_observed) + Data for the observed features only, where + ``n_observed = observed_feature_idxs.sum()``. - observed_feature_idxs: np.array-like,(sum(p),) - The observation locations + observed_feature_idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 indicates an observed + feature. - weights : np.array-like,(N,),default=None - The optional weights on the samples + weights : ndarray of shape (n_samples,), default=None + Optional sample weights. If ``None``, all samples are weighted + equally. + + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity for efficient + computation under low-rank-plus-diagonal covariance structure. Returns ------- - avg_score : float - Average log likelihood + avg_score : np.float64 + Weighted average marginal log-likelihood. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -453,20 +486,26 @@ def marginal_score_samples( sherman_woodbury: bool = False, ): """ - Returns the marginal log-likelihood of a subset of data + Compute the per-sample marginal log-likelihood of a feature subset. Parameters ---------- - X : np.array-like,(N,sum(observed_feature_idxs)) - The data + X : ndarray of shape (n_samples, n_observed) + Data for the observed features only, where + ``n_observed = observed_feature_idxs.sum()``. + + observed_feature_idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 indicates an observed + feature. - observed_feature_idxs: np.array-like,(sum(p),) - The observation locations + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity for efficient + computation under low-rank-plus-diagonal covariance structure. Returns ------- - scores : float - Average log likelihood + scores : ndarray of shape (n_samples,) + Per-sample marginal log-likelihood values. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -486,20 +525,25 @@ def score( sherman_woodbury: bool = False, ) -> np.float64: """ - Returns the average log liklihood of data. + Compute the weighted average log-likelihood of the data. Parameters ---------- - X : np.array-like,(N,sum(p)) - The data + X : ndarray of shape (n_samples, n_features) + The centered data matrix. - weights : np.array-like,(N,),default=None - The optional weights on the samples + weights : ndarray of shape (n_samples,), default=None + Optional sample weights. If ``None``, all samples are weighted + equally. Weights are normalized to have mean 1. + + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity for efficient + computation under low-rank-plus-diagonal covariance structure. Returns ------- - avg_score : float - Average log likelihood + avg_score : np.float64 + Weighted average log-likelihood. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -515,17 +559,21 @@ def score_samples( self, X: NDArray[np.float64], sherman_woodbury: bool = False ) -> NDArray[np.float64]: """ - Return the log likelihood of each sample + Compute the per-sample log-likelihood. Parameters ---------- - X : NDArray[np.float64],(N,sum(p)) - The data + X : ndarray of shape (n_samples, n_features) + The centered data matrix. + + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity for efficient + computation under low-rank-plus-diagonal covariance structure. Returns ------- - scores : np.float64 - Log likelihood for each sample + scores : ndarray of shape (n_samples,) + Per-sample log-likelihood values. """ if self.W_ is None: raise ValueError("Model has not been fit yet") @@ -537,17 +585,12 @@ def score_samples( def get_entropy(self) -> np.float64: """ - Computes the entropy of a Gaussian distribution parameterized by - covariance. - - Parameters - ---------- - None + Compute the differential entropy of the model's Gaussian distribution. Returns ------- entropy : np.float64 - The differential entropy of the distribution + Differential entropy of the distribution in nats. """ return _entropy(self.get_covariance()) @@ -555,18 +598,18 @@ 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. + Compute the differential entropy of a marginal feature subset. Parameters ---------- - observed_feature_idxs: NDArray[np.float64],(sum(p),) - The observation locations + observed_feature_idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 indicates the features + to include in the entropy computation. Returns ------- entropy : np.float64 - The differential entropy of the distribution + Differential entropy of the marginal distribution in nats. """ return _entropy_subset(self.get_covariance(), observed_feature_idxs) @@ -576,21 +619,25 @@ def mutual_information( observed_feature_idxs2: NDArray[np.float64], ) -> np.float64: """ - This computes the mutual information bewteen the two sets of - covariates based on the model. + Compute the mutual information between two groups of features. + + Uses the fitted model covariance to compute the mutual information + ``I(X1 ; X2)`` based on the Gaussian entropy formula. Parameters ---------- - observed_feature_idxs1 : np.array-like,(p,) - First group of variables + observed_feature_idxs1 : ndarray of shape (n_features,) + Boolean or integer index array selecting the first group of + features. - observed_feature_idxs2 : np.array-like,(p,) - Second group of variables + observed_feature_idxs2 : ndarray of shape (n_features,) + Boolean or integer index array selecting the second group of + features. Returns ------- mutual_information : np.float64 - The mutual information between the two variables + Mutual information between the two feature groups in nats. """ covariance = self.get_covariance() return _mutual_information( @@ -606,25 +653,23 @@ def __init__( 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. + Base class for factor models fitted via stochastic gradient descent. + + Provides shared infrastructure (loss tracking, device handling, + optimizer setup) for SGD-based Gaussian factor models. Parameters ---------- - n_components : int,default=2 - The latent dimensionality + n_components : int, default=2 + The dimensionality of the latent space. - training_options : dict,default={} - The options for gradient descent + training_options : dict, default=None + Options for stochastic gradient descent. Keys and defaults are + set by ``_fill_training_options``. - prior_options : dict,default={} - The options for priors on model parameters - - Sets - ---- - creationDate : datetime - The date/time that the object was created + prior_options : dict, default=None + Options for the prior on model parameters. Keys and defaults + are set by ``_fill_prior_options`` in the subclass. """ super().__init__(n_components=n_components) @@ -640,31 +685,33 @@ 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. + Fill in default training options for stochastic gradient descent. + + Merges user-supplied options with defaults. Raises ``ValueError`` + if any unrecognized keys are present. Parameters ---------- training_options : dict - The original options set by the user passed as a dictionary + User-supplied training options. Recognized keys: + + - ``n_iterations`` : int, default=3000 + Number of SGD iterations. + - ``learning_rate`` : float, default=1e-2 + Step size for the optimizer. + - ``use_gpu`` : bool, default=True + Whether to use a GPU if available. + - ``method`` : str, default='Nadam' + Optimization algorithm name. + - ``batch_size`` : int, default=100 + Mini-batch size per iteration. + - ``momentum`` : float, default=0.9 + Momentum coefficient for the SGD optimizer. - Options + Returns ------- - 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 + tops : dict + Merged dictionary of training options with defaults applied. """ default_options = { "n_iterations": 3000, @@ -688,19 +735,20 @@ def _fill_training_options( def _initialize_save_losses(self) -> None: """ - This method initializes the arrays to track relevant variables - during training at each iteration + Initialize arrays to record training losses at each iteration. - Sets - ---- - losses_likelihood : np.array(n_iterations) - The log likelihood + Sets the following attributes: + + Attributes + ---------- + losses_likelihood : ndarray of shape (n_iterations,) + Per-iteration log-likelihood values. - losses_prior : np.array(n_iterations) - The log prior + losses_prior : ndarray of shape (n_iterations,) + Per-iteration log-prior values. - losses_posterior : np.array(n_iterations) - The log posterior + losses_posterior : ndarray of shape (n_iterations,) + Per-iteration log-posterior values. """ n_iterations = self.training_options["n_iterations"] self.losses_likelihood = np.empty(n_iterations) @@ -716,24 +764,24 @@ def _save_losses( log_posterior: Tensor, ) -> None: """ - Saves the values of the losses at each iteration + Record training loss values for the current iteration. Parameters - ----------- - device ; pytorch.device - The device used for trainging (gpu or cpu) - + ---------- i : int - Current training iteration + Current training iteration index. + + device : Any + PyTorch device (CPU or CUDA) used during training. - losses_likelihood : Tensor - The log likelihood + log_likelihood : Tensor + Log-likelihood value at iteration ``i``. - losses_prior : Tensor | NDArray[np.float64] - The log prior + log_prior : Tensor or ndarray of shape () + Log-prior value at iteration ``i``. - losses_posterior : Tensor - The log posterior + log_posterior : Tensor + Log-posterior value at iteration ``i``. """ if device.type == "cuda": self.losses_likelihood[i] = log_likelihood.detach().cpu().numpy() @@ -754,7 +802,20 @@ def _transform_training_data( self, device: Any, *args: NDArray ) -> list[Tensor]: """ - Convert a list of numpy arrays to tensors + Convert numpy arrays to float32 PyTorch tensors on the target device. + + Parameters + ---------- + device : Any + PyTorch device (CPU or CUDA) to place the tensors on. + + *args : ndarray + One or more numpy arrays to convert. + + Returns + ------- + out : list of Tensor + Tensors corresponding to each input array, cast to float32. """ out = [] for arg in args: @@ -764,15 +825,16 @@ def _transform_training_data( @abstractmethod def _create_prior(self, device: Any): """ - Creates a prior on the parameters taking your trainable variable - dictionary as input + Construct the log-prior function for the model parameters. Parameters ---------- - None + device : Any + PyTorch device (CPU or CUDA) on which to place prior tensors. Returns ------- - log_prior : function - The function representing the negative log density of the prior + log_prior : callable + A function that accepts the list of trainable variable tensors + and returns a scalar Tensor representing the log-prior density. """ diff --git a/gpcr/_base_covariance.py b/gpcr/_base_covariance.py index d2780e2..911ae02 100644 --- a/gpcr/_base_covariance.py +++ b/gpcr/_base_covariance.py @@ -78,18 +78,70 @@ def __init__(self) -> None: self.covariance: NDArray | None = None def get_precision(self) -> NDArray[np.float64]: + """ + Compute the precision matrix (inverse of the covariance matrix). + + Returns + ------- + precision : ndarray of shape (n_features, n_features) + The precision matrix. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ 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") + """ + Compute the stable rank of the covariance matrix. + + The stable rank is defined as ``||A||_F^2 / ||A||_2^2``. + + Returns + ------- + srank : np.float64 + The stable rank of the covariance matrix. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ return _get_stable_rank(self.covariance) def predict(self, Xobs: NDArray, idxs: NDArray): + """ + Predict missing features from observed features. + + Uses the conditional Gaussian formula to predict the unobserved + features (where ``idxs == 0``) from the observed features + (where ``idxs == 1``). + + Parameters + ---------- + Xobs : ndarray of shape (n_samples, n_observed) + Observed feature values, where ``n_observed = idxs.sum()``. + + idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 marks observed features + and 0 marks features to predict. + + Returns + ------- + preds : ndarray of shape (n_samples, n_missing) + Predicted values for the unobserved features. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") @@ -98,6 +150,33 @@ def predict(self, Xobs: NDArray, idxs: NDArray): def conditional_score( self, X: NDArray, idxs: NDArray, weights=None ) -> np.float64: + """ + Compute the weighted average conditional log-likelihood. + + Evaluates ``mean(log p(X[idxs==0] | X[idxs==1], covariance))``. + + Parameters + ---------- + X : ndarray of shape (n_samples, n_features) + The centered data matrix. + + idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 marks observed features + and 0 marks unobserved features. + + weights : ndarray of shape (n_samples,), default=None + Optional sample weights. Normalized to have mean 1. + + Returns + ------- + avg_score : np.float64 + Weighted average conditional log-likelihood. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") @@ -106,6 +185,31 @@ def conditional_score( def conditional_score_samples( self, X: NDArray, idxs: NDArray ) -> NDArray[np.float64]: + """ + Compute the per-sample conditional log-likelihood. + + Evaluates ``log p(X[idxs==0] | X[idxs==1], covariance)`` for each + sample. + + Parameters + ---------- + X : ndarray of shape (n_samples, n_features) + The centered data matrix. + + idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 marks observed features + and 0 marks unobserved features. + + Returns + ------- + scores : ndarray of shape (n_samples,) + Per-sample conditional log-likelihood values. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") @@ -114,6 +218,34 @@ def conditional_score_samples( def marginal_score( self, X: NDArray, idxs: NDArray, weights=None ) -> np.float64: + """ + Compute the weighted average marginal log-likelihood. + + Evaluates ``mean(log p(X[:, idxs==1], covariance_sub))``. + + Parameters + ---------- + X : ndarray of shape (n_samples, n_observed) + Data for the observed features, where + ``n_observed = idxs.sum()``. + + idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 marks the included + features. + + weights : ndarray of shape (n_samples,), default=None + Optional sample weights. Normalized to have mean 1. + + Returns + ------- + avg_score : np.float64 + Weighted average marginal log-likelihood. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") @@ -122,30 +254,124 @@ def marginal_score( def marginal_score_samples( self, X: NDArray, idxs: NDArray ) -> NDArray[np.float64]: + """ + Compute the per-sample marginal log-likelihood. + + Parameters + ---------- + X : ndarray of shape (n_samples, n_observed) + Data for the observed features, where + ``n_observed = idxs.sum()``. + + idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 marks the included + features. + + Returns + ------- + scores : ndarray of shape (n_samples,) + Per-sample marginal log-likelihood values. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ 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: + """ + Compute the weighted average log-likelihood. + + Parameters + ---------- + X : ndarray of shape (n_samples, n_features) + The centered data matrix. + + weights : ndarray of shape (n_samples,), default=None + Optional sample weights. Normalized to have mean 1. + + Returns + ------- + avg_score : np.float64 + Weighted average log-likelihood. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ 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]: + """ + Compute the per-sample log-likelihood. + + Parameters + ---------- + X : ndarray of shape (n_samples, n_features) + The centered data matrix. + + Returns + ------- + scores : ndarray of shape (n_samples,) + Per-sample log-likelihood values. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") return _score_samples(self.covariance, X) def entropy(self) -> np.float64: + """ + Compute the differential entropy of the fitted Gaussian distribution. + + Returns + ------- + entropy : np.float64 + Differential entropy in nats. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ 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: + """ + Compute the differential entropy of a marginal feature subset. + + Parameters + ---------- + idxs : ndarray of shape (n_features,) + Boolean or integer index array where 1 marks the features to + include in the entropy computation. + + Returns + ------- + entropy : np.float64 + Differential entropy of the marginal distribution in nats. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") @@ -154,6 +380,29 @@ def entropy_subset(self, idxs: NDArray[np.float64]) -> np.float64: def mutual_information( self, idxs1: NDArray[np.float64], idxs2: NDArray[np.float64] ) -> np.float64: + """ + Compute the mutual information between two groups of features. + + Parameters + ---------- + idxs1 : ndarray of shape (n_features,) + Boolean or integer index array selecting the first group of + features. + + idxs2 : ndarray of shape (n_features,) + Boolean or integer index array selecting the second group of + features. + + Returns + ------- + mutual_information : np.float64 + Mutual information between the two feature groups in nats. + + Raises + ------ + ValueError + If the covariance matrix has not been set. + """ if self.covariance is None: raise ValueError("Covariance matrix has not been fit") @@ -161,7 +410,17 @@ def mutual_information( def _test_inputs(self, X: NDArray[np.float64]) -> None: """ - Just tests to make sure data is numpy array + Validate that ``X`` is a numpy array with no missing values. + + Parameters + ---------- + X : ndarray + Input data to validate. + + Raises + ------ + ValueError + If ``X`` is not a numpy array or contains NaN values. """ if not isinstance(X, np.ndarray): raise ValueError("Data is not a numpy array") @@ -854,21 +1113,25 @@ def _mutual_information( idxs2: NDArray[np.float64], ) -> np.float64: """ - This computes the mutual information bewteen the two sets of - covariates based on the model. + Compute the mutual information between two groups of features. + + Uses the Gaussian entropy formula: ``I(X1; X2) = H(X1) - H(X1|X2)``. Parameters ---------- - idxs1 : NDArray[np.float64],(p,) - First group of variables + covariance : ndarray of shape (n_features, n_features) + The covariance matrix. + + idxs1 : ndarray of shape (n_features,) + Boolean or integer index array selecting the first group of features. - idxs2 : NDArray[np.float64],(p,) - Second group of variables + idxs2 : ndarray of shape (n_features,) + Boolean or integer index array selecting the second group of features. Returns ------- mutual_information : np.float64 - The mutual information between the two variables + Mutual information between the two feature groups in nats. """ idxs = np.logical_or(idxs1, idxs2).astype(int) cov_sub = covariance[np.ix_(idxs == 1, idxs == 1)] diff --git a/gpcr/_covariance_np.py b/gpcr/_covariance_np.py index 248a4b2..83787f1 100644 --- a/gpcr/_covariance_np.py +++ b/gpcr/_covariance_np.py @@ -31,30 +31,30 @@ class EmpiricalCovariance(BaseCovariance): def __init__(self): """ - This object just fits the covariance matrix as the standard sample - covariance matrix. Does not handle missing values + Fit the standard sample covariance matrix. + + Does not handle missing values. Computes ``X.T @ X / N``. """ super().__init__() def fit(self, X: NDArray) -> "EmpiricalCovariance": """ - This fits a covariance matrix using samples X. + Fit the empirical covariance matrix from data. Parameters ---------- - X : np.array-like,(n_samples,n_covariates) - The centered data + X : ndarray of shape (n_samples, n_features) + The centered data matrix. Must not contain NaN values. Returns ------- self : EmpiricalCovariance - The model + Fitted estimator. 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 + ValueError + If ``X`` is not a numpy array or contains NaN values. """ self._test_inputs(X) self.N, self.p = X.shape @@ -68,8 +68,8 @@ def fit(self, X: NDArray) -> "EmpiricalCovariance": 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 + Fit the covariance matrix as a MAP estimator with an + inverse-Wishart prior. Does not handle missing values. """ super().__init__() if prior_options is None: @@ -79,23 +79,22 @@ def __init__(self, prior_options=None): def fit(self, X: NDArray[np.float_]) -> "BayesianCovariance": """ - This fits a covariance matrix using samples X with MAP estimation. + Fit the MAP covariance estimate using an inverse-Wishart prior. Parameters ---------- - X : np.array-like,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The data matrix. Must not contain NaN values. Returns ------- self : BayesianCovariance - The model + Fitted estimator. 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 + ValueError + If ``X`` is not a numpy array or contains NaN values. """ self._test_inputs(X) self.N, self.p = X.shape @@ -116,18 +115,25 @@ def _fill_prior_options( self, prior_options: dict[str, Any] ) -> dict[str, Any]: """ - This sets the prior options for our inference scheme + Fill in default prior options for MAP estimation. Parameters ---------- prior_options : dict - The original prior options passed as a dictionary + User-supplied prior options. Recognized keys: + + - ``iw_params`` : dict, default ``{'pnu': 2, 'sigma': 1.0}`` + Parameters for the inverse-Wishart prior: - Options + - ``pnu`` : int — excess degrees of freedom above ``p`` + (total ``nu = p + pnu``). + - ``sigma`` : float > 0 — scale of the prior covariance + ``cov = sigma * I_p``. + + Returns ------- - iw_params : dict,default={'pnu':2,'sigma':1.0} - pnu : int - nu = p + pnu - sigma : float>0 - cov = sigma*I_p + options : dict + Merged dictionary of prior options with defaults applied. """ default_options = { "iw_params": {"pnu": 2, "sigma": 1.0}, diff --git a/gpcr/_misc_np.py b/gpcr/_misc_np.py index 6a03069..774bfd3 100644 --- a/gpcr/_misc_np.py +++ b/gpcr/_misc_np.py @@ -3,21 +3,23 @@ def softplus_inverse_np(y): """ - Computes the inverse of the softplus activation of x in a - numerically stable way + Compute the inverse of the softplus function in a numerically stable way. - Softplus: y = log(exp(x) + 1) - Softplus^{-1}: y = np.log(np.exp(x) - 1) + Softplus is defined as ``y = log(exp(x) + 1)``, so its inverse is + ``x = log(exp(y) - 1)``. For very small or very large ``y``, the naive + formula is numerically unstable; this function handles those edge cases + by clamping and substitution. Parameters ---------- - x : np.array - Original array + y : array-like + Input array; should contain positive values (the range of softplus). Returns ------- - x : np.array - Transformed array + x : ndarray + Array of the same shape as ``y`` containing the inverse softplus + values. """ min_threshold = 10**-15 max_threshold = 500 @@ -34,25 +36,49 @@ def softplus_inverse_np(y): def subset_square_matrix_np(Sigma, idxs): """ - This returns a symmetric subset of a square matrix + Extract a symmetric submatrix from a square matrix. Parameters ---------- - Sigma : np.array-like,(p,p) - Covariance matrix (presumably) + Sigma : ndarray of shape (n, n) + The square matrix to subset (typically a covariance matrix). + + idxs : ndarray of shape (n,) + Boolean or integer index array where 1 marks the rows/columns + to select. - idxs : np.array-like,(p,) - Binary vector, 1 means select Returns ------- - Sigma_sub : np.array-like,(sum(idxs),sum(idxs)) - The subset of matrix + Sigma_sub : ndarray of shape (k, k) + The submatrix indexed by ``idxs == 1``, where ``k = idxs.sum()``. """ Sigma_sub = Sigma[np.ix_(idxs == 1, idxs == 1)] return Sigma_sub def classify_missingness(matrix): + """ + Group rows of a matrix by their NaN-missingness pattern. + + Identifies unique patterns of missing values (NaN positions) across + rows and groups the rows accordingly. + + Parameters + ---------- + matrix : ndarray of shape (n_samples, n_features) + Input matrix potentially containing NaN values. + + Returns + ------- + matrices_list : list of ndarray + List of sub-matrices, one per unique missingness pattern. + Each sub-matrix contains all rows sharing that pattern. + + vectors_list : list of ndarray of bool + List of boolean observation masks, one per unique pattern, + where ``True`` indicates that the feature is observed + (not NaN) in that group. + """ pattern_dict = {} matrices_list = [] vectors_list = [] diff --git a/gpcr/generative_fa.py b/gpcr/generative_fa.py index 5adacf1..4c018c9 100644 --- a/gpcr/generative_fa.py +++ b/gpcr/generative_fa.py @@ -44,28 +44,28 @@ def fit( sherman_woodbury: bool = False, ) -> "PPCA": """ - Fits a model given covariates X as well as option labels y in the - supervised methods + Fit the PPCA model to data using stochastic gradient descent. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The centered data matrix. - progress_bar : bool,default=True - Whether to print the progress bar to monitor time + progress_bar : bool, default=True + Whether to display a tqdm progress bar during training. - seed : int,default=2021 - The seed of the random number generator + seed : int, default=2021 + Seed for the random number generator used for mini-batch sampling. - sherman_woodbury : bool,default=False - Whether to use the Sherman Woodbury identity to calculate - the likelihood. Advantageous in high-p situations + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity to evaluate the + log-likelihood, which is advantageous when ``n_features`` is + large relative to ``n_components``. Returns ------- self : PPCA - The model + Fitted estimator. """ self._test_inputs(X) training_options = self.training_options @@ -138,18 +138,14 @@ def fit( def get_covariance(self): """ - Gets the covariance matrix + Compute the model covariance matrix. - Sigma = W^TW + sigma2*I - - Parameters - ---------- - None + The covariance is given by ``Sigma = W^T W + sigma^2 * I``. Returns ------- - covariance : NDArray(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The PPCA covariance matrix. """ if self.W_ is None or self.sigma2_ is None or self.p is None: raise ValueError("Fit model first") @@ -158,16 +154,14 @@ def get_covariance(self): def get_noise(self): """ - Returns the observational noise as a diagonal matrix + Return the observational noise as a diagonal matrix. - Parameters - ---------- - None + For PPCA the noise is isotropic: ``Lambda = sigma^2 * I``. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + The isotropic diagonal noise matrix. """ if self.sigma2_ is None or self.p is None: raise ValueError("Fit model first") @@ -176,12 +170,22 @@ def get_noise(self): def _create_prior(self, device): """ - This creates the function representing prior on pararmeters + Construct the log-prior function for PPCA parameters. + + Creates a closure over ``prior_options`` that computes the + log-prior on ``W_`` (Gaussian regularization) and ``sigma`` + (Gamma prior on the noise variance). Parameters ---------- - log_prior : function - The function representing the log density of the prior + device : Any + PyTorch device (CPU or CUDA) on which to place prior tensors. + + Returns + ------- + log_prior : callable + Function that accepts the list of trainable variable tensors + and returns a scalar Tensor of the log-prior value. """ prior_options = self.prior_options @@ -207,25 +211,27 @@ 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. + Initialize model parameters from a sklearn PCA warm start. + + Fits a PCA model to ``X``, uses its loadings as ``W_``, and sets + the initial noise ``sigma^2`` to the mean squared reconstruction + error. Parameters ---------- - device ; pytorch.device - The device used for trainging (gpu or cpu) + device : Any + PyTorch device (CPU or CUDA) on which to place the tensors. - X : NDArray,(n_samples,p) - The data + X : ndarray of shape (n_samples, n_features) + The training data. Returns ------- - W_ : torch.tensor,shape=(n_components,p) - The loadings of our latent factor model + W_ : Tensor of shape (n_components, n_features) + Initial loadings tensor with ``requires_grad=True``. - sigmal_ : torch.tensor - The unrectified variance of the model + sigmal_ : Tensor of shape (1,) + Initial unrectified noise variance with ``requires_grad=True``. """ model = PCA(self.n_components) S_hat = model.fit_transform(X) @@ -241,20 +247,23 @@ def _store_instance_variables( # type: ignore[override] self, device: Any, trainable_variables: list[Tensor] ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list[Tensor] - List of tensorflow variables saved + device : Any + PyTorch device used during training. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + trainable_variables : list of Tensor + Tensors ``[W_, sigmal_]`` after optimization. + + Attributes + ---------- + W_ : ndarray of shape (n_components, n_features) + Fitted factor loadings. sigma2_ : np.float_ - The isotropic variance + Fitted isotropic noise variance. """ if device.type == "cuda": self.W_ = trainable_variables[0].detach().cpu().numpy() @@ -269,7 +278,18 @@ def _store_instance_variables( # type: ignore[override] def _test_inputs(self, X: NDArray[np.float_]) -> None: """ - Just tests to make sure data is numpy array + Validate inputs before fitting. + + Parameters + ---------- + X : ndarray + The training data. + + Raises + ------ + ValueError + If ``X`` is not a numpy array or the batch size exceeds the + number of samples. """ if not isinstance(X, np.ndarray): raise ValueError("Data is numpy array") @@ -280,12 +300,24 @@ def _fill_prior_options( self, prior_options: dict[str, Any] ) -> dict[str, Any]: """ - Fills in options for prior parameters + Fill in default prior options for PPCA parameters. + + Parameters + ---------- + prior_options : dict + User-supplied prior options. Recognized keys: - Paramters - --------- - new_dict : dictionary - The prior parameters used to specify the prior + - ``weight_W`` : float, default=0.01 + L2 regularization weight on the loadings ``W``. + - ``alpha`` : float, default=3.0 + Shape parameter of the Gamma prior on the noise variance. + - ``beta`` : float, default=3.0 + Rate parameter of the Gamma prior on the noise variance. + + Returns + ------- + options : dict + Merged dictionary of prior options with defaults applied. """ default_dict = {"weight_W": 0.01, "alpha": 3.0, "beta": 3.0} @@ -300,20 +332,22 @@ def __init__( 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. + Factor analysis with heterogeneous diagonal noise per feature. + + Each feature is given its own noise variance, unlike PPCA which + uses a single isotropic noise parameter. Parameters ---------- - n_components : int,default=2 - The latent dimensionality + n_components : int, default=2 + The dimensionality of the latent space. - training_options : dict,default={} - The options for gradient descent + training_options : dict, default=None + Options for gradient descent passed to ``_fill_training_options``. - prior_options : dict,default={} - The options for priors on model parameters + prior_options : dict, default=None + Options for priors on model parameters passed to + ``_fill_prior_options``. """ super().__init__( n_components=n_components, @@ -336,17 +370,27 @@ def fit( sherman_woodbury: bool = False, ) -> "FactorAnalysis": """ - Fits a model given covariates X + Fit the factor analysis model to data using stochastic gradient descent. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The centered data matrix. + + progress_bar : bool, default=True + Whether to display a tqdm progress bar during training. + + seed : int, default=2021 + Seed for the random number generator used for mini-batch sampling. + + sherman_woodbury : bool, default=False + If ``True``, uses the Sherman-Woodbury identity to evaluate the + log-likelihood. Returns ------- self : FactorAnalysis - The model + Fitted estimator. """ self._test_inputs(X) training_options = self.training_options @@ -413,18 +457,14 @@ def fit( def get_covariance(self) -> NDArray[np.float_]: """ - Gets the covariance matrix - - Sigma = W^TW + sigma2*I + Compute the model covariance matrix. - Parameters - ---------- - None + The covariance is given by ``Sigma = W^T W + diag(sigmas)``. Returns ------- - covariance : NDArray(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The factor analysis covariance matrix. """ if self.W_ is None or self.sigmas_ is None: raise ValueError("Fit model first") @@ -433,16 +473,12 @@ def get_covariance(self) -> NDArray[np.float_]: def get_noise(self) -> NDArray[np.float_]: """ - Returns the observational noise as a diagonal matrix - - Parameters - ---------- - None + Return the observational noise as a diagonal matrix. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + Diagonal noise matrix ``diag(sigmas_)``. """ if self.sigmas_ is None: raise ValueError("Fit model first") @@ -451,12 +487,21 @@ def get_noise(self) -> NDArray[np.float_]: def _create_prior(self, device) -> Callable[[list[Tensor]], Tensor]: """ - This creates the function representing prior on pararmeters + Construct the log-prior function for factor analysis parameters. + + Creates a closure that computes a Gamma log-prior on the + (softmax-transformed) per-feature noise variances. Parameters ---------- - log_prior : function - The function representing the negative log density of the prior + device : Any + PyTorch device (CPU or CUDA) on which to place prior tensors. + + Returns + ------- + log_prior : callable + Function that accepts the list of trainable variable tensors + and returns a scalar Tensor of the log-prior value. """ def log_prior(trainable_variables: list[Tensor]) -> Tensor: @@ -473,33 +518,50 @@ def _fill_prior_options( self, prior_options: dict[str, Any] ) -> dict[str, Any]: """ - Fills in options for prior parameters + Fill in default prior options for factor analysis parameters. + + Parameters + ---------- + prior_options : dict + User-supplied prior options. Recognized keys: + + - ``alpha`` : float, default=3.0 + Shape parameter of the Gamma prior on noise variances. + - ``beta`` : float, default=3.0 + Rate parameter of the Gamma prior on noise variances. - Paramters - --------- - new_dict : dictionary - The prior parameters used to specify the prior + Returns + ------- + options : dict + Merged dictionary of prior options with defaults applied. """ 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 + Initialize model parameters from a sklearn PCA warm start. + + Fits a PCA model to ``X``, uses its loadings as ``W_``, and + initializes per-feature noise variances to the mean squared + reconstruction error. Parameters ---------- - X : NDArray,(n_samples,p) - The data + device : Any + PyTorch device (CPU or CUDA) on which to place the tensors. + + X : ndarray of shape (n_samples, n_features) + The training data. Returns ------- - W_ : torch.tensor,(n_components,p) - The loadings of our latent factor model + W_ : Tensor of shape (n_components, n_features) + Initial loadings tensor with ``requires_grad=True``. - sigmal_ : torch.tensor,(p,) - The noise of each covariate, unrectified + sigmal_ : Tensor of shape (n_features,) + Initial unrectified per-feature noise variances with + ``requires_grad=True``. """ if self.p is None: raise ValueError("Fit model first") @@ -523,20 +585,23 @@ def _store_instance_variables( # type: ignore[override] self, device: Any, trainable_variables: list[Tensor] ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of tensorflow variables saved + device : Any + PyTorch device used during training. + + trainable_variables : list of Tensor + Tensors ``[W_, sigmal_]`` after optimization. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_components, n_features) + Fitted factor loadings. - sigmas_ : NDArray,(n_components,p) - The diagonal variances + sigmas_ : ndarray of shape (n_features,) + Fitted per-feature noise standard deviations (after softplus). """ if device.type == "cuda": self.W_ = trainable_variables[0].detach().cpu().numpy() @@ -551,7 +616,18 @@ def _store_instance_variables( # type: ignore[override] def _test_inputs(self, X: NDArray[np.float_]): """ - Just tests to make sure data is numpy array + Validate inputs before fitting. + + Parameters + ---------- + X : ndarray + The training data. + + Raises + ------ + ValueError + If ``X`` is not a numpy array or the batch size exceeds the + number of samples. """ if not isinstance(X, np.ndarray): raise ValueError("Data is numpy array") diff --git a/gpcr/gpcr.py b/gpcr/gpcr.py index e124ec4..6d751dd 100644 --- a/gpcr/gpcr.py +++ b/gpcr/gpcr.py @@ -40,28 +40,27 @@ 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. + Compute the posterior parameters of the latent distribution p(S|X). - Computes the parameters for p(S|X) - - Description in future to be released paper + Currently implemented only for PPCA due to its closed-form expression. + Computes the projection matrix and posterior covariance via the matrix + M = W W^T + sigma * I. Parameters ---------- - W_ : Tensor(n_components,p) - The loadings + W_ : Tensor of shape (n_components, p) + The factor loadings matrix. sigma_ : Tensor - Isotropic noise + Scalar isotropic noise variance. Returns ------- - Proj_X : Tensor(n_components,p) - Beta such that np.dot(Proj_X, X) = E[S|X] + Proj_X : Tensor of shape (n_components, p) + Projection matrix such that ``Proj_X @ X`` gives ``E[S|X]``. - Cov : Tensor(n_components,n_components) - Var(S|X) + Cov : Tensor of shape (n_components, n_components) + Posterior covariance ``Var(S|X)``. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") n_components = int(W_.shape[0]) @@ -190,31 +189,33 @@ def fit( # type: ignore[override] seed: int = 2021, ) -> "GPCRvi": """ - Fits a model given covariates X as well as option labels y in the - supervised methods + Fit the GPCRvi model to data with a supervised auxiliary task. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The centered covariate matrix. + + y : ndarray of shape (n_samples,) or (n_samples, n_outputs) + Supervised labels or continuous targets to predict from the + first ``n_supervised`` latent dimensions. - y : NDArray,(n_samples,n_prediction) - Covariates we wish to predict. For now lazy and assuming - logistic regression. + task : str, default='classification' + The prediction task type. One of: - task : string,default='classification' - Is this prediction, multinomial regression, or classification + - ``'classification'`` — binary classification via logistic loss. + - ``'regression'`` — continuous regression via MSE loss. - progress_bar : bool,default=True - Whether to print the progress bar to monitor time + progress_bar : bool, default=True + Whether to display a tqdm progress bar during training. - seed : int,default=2021 - The random number generator seed used to ensure reproducibility + seed : int, default=2021 + Seed for the random number generator used for mini-batch sampling. Returns ------- self : GPCRvi - The model + Fitted estimator. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -406,27 +407,26 @@ def fit( # type: ignore[override] def _store_instance_variables(self, trainable_variables: list[Tensor]): """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of saved variables of type Tensor + trainable_variables : list of Tensor + Tensors ``[W_, sigmal_, d_weights, d_bias]`` after optimization. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_components, n_features) + Fitted factor loadings. sigma2_ : float - The isotropic variance + Fitted isotropic noise variance. - d_weights_ : NDArray,(n_supervised,) - The learned predictive coefficients - - d_bias_ : NDArray,(1,) - The learned predictive bias + d_weights_ : ndarray of shape (n_supervised,) + Learned predictive weight vector. + d_bias_ : ndarray of shape (1,) + Learned predictive bias term. """ self.W_ = trainable_variables[0].detach().cpu().numpy() self.sigma2_ = ( @@ -437,23 +437,28 @@ def _store_instance_variables(self, trainable_variables: list[Tensor]): 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. + Initialize model parameters using a supervised warm start. + + Fits a ``PPCAsupervisedRandomized`` model with high supervision + strength to obtain a good initial loading for the first latent + dimension, then fills remaining components with PCA loadings. Parameters ---------- - X : NDArray,(n_samples,p) - The data + X : ndarray of shape (n_samples, n_features) + The covariate data. + + Y : ndarray of shape (n_samples, 1) + The supervised targets used to orient the first latent dimension. Returns ------- - W_ : torch.tensor,shape=(n_components,p) - The loadings of our latent factor model - - sigmal_ : torch.tensor - The unrectified variance of the model + W_ : Tensor of shape (n_components, n_features) + Initial loadings tensor with ``requires_grad=True``. + sigmal_ : Tensor of shape (1,) + Initial unrectified isotropic noise variance with + ``requires_grad=True``. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = PPCAsupervisedRandomized(n_components=1, mu=100000.0) @@ -480,7 +485,21 @@ def _initialize_variables(self, X: NDArray, Y: NDArray): def _test_inputs(self, X, y): """ - Just tests to make sure data is numpy array and dimensions match + Validate inputs before fitting. + + Parameters + ---------- + X : ndarray + The covariate data. + + y : array-like + The supervised labels. + + Raises + ------ + ValueError + If ``X`` is not a numpy array, the batch size exceeds the number + of samples, or the lengths of ``X`` and ``y`` do not match. """ if not isinstance(X, np.ndarray): raise ValueError("Data must be numpy array") @@ -492,12 +511,17 @@ def _test_inputs(self, X, y): def _create_prior(self): """ - This creates the function representing prior on pararmeters + Construct the log-prior function for GPCRvi parameters. - Parameters - ---------- - log_prior : function - The function representing the log density of the prior + Creates a closure over ``prior_options`` that computes the + log-prior on ``W_`` (Gaussian regularization) and ``sigma`` + (Gamma prior on the noise variance). + + Returns + ------- + log_prior : callable + Function that accepts the list of trainable variable tensors + and returns a scalar Tensor of the log-prior value. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") prior_options = self.prior_options @@ -528,21 +552,21 @@ def _save_losses( log_posterior, ) -> None: """ - Saves the values of the losses at each iteration + Record training loss values for the current iteration. Parameters - ----------- + ---------- i : int - Current training iteration + Current training iteration index. - losses_likelihood : Tensor - The log likelihood + log_likelihood : Tensor + Log-likelihood value at iteration ``i``. - losses_prior : NDArray[np.float_] | Tensor - The log prior + log_prior : Tensor or ndarray + Log-prior value at iteration ``i``. - losses_posterior : Tensor - The log posterior + log_posterior : Tensor + Log-posterior value at iteration ``i``. """ self.losses_likelihood[i] = log_likelihood.detach().cpu().numpy() if isinstance(log_prior, Tensor): diff --git a/gpcr/ppca_augmented.py b/gpcr/ppca_augmented.py index 7b6dfea..90af720 100644 --- a/gpcr/ppca_augmented.py +++ b/gpcr/ppca_augmented.py @@ -81,20 +81,20 @@ def fit( self, X: NDArray[np.float_], Y: NDArray[np.float_] ) -> "PPCAadversarial": """ - Fits a model given covariates X + Fit the adversarial PPCA model. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The covariate data. - Y : NDArray,(n_samples,n_confounders) - The concommitant data we want to remove + Y : ndarray of shape (n_samples, n_confounders) + The confounder data to decorrelate from the latent representation. Returns ------- self : PPCAadversarial - The model + Fitted estimator. """ self.mean_x = np.mean(X, axis=0) self.mean_y = np.mean(Y, axis=0) @@ -145,18 +145,14 @@ def fit( def get_covariance(self) -> NDArray[np.float_]: """ - Gets the covariance matrix + Compute the model covariance matrix. - Sigma = W^TW + sigma2*I - - Parameters - ---------- - None + The covariance is given by ``Sigma = W^T W + sigma^2 * I``. Returns ------- - covariance : NDArray,(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The PPCA 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") @@ -165,16 +161,14 @@ def get_covariance(self) -> NDArray[np.float_]: def get_noise(self): """ - Returns the observational noise as a diagonal matrix + Return the observational noise as a diagonal matrix. - Parameters - ---------- - None + For PPCA-based models the noise is isotropic: ``Lambda = sigma^2 * I``. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + The isotropic diagonal noise matrix. """ if self.sigma2_ is None or self.p is None: raise ValueError("Model has not been fit yet") @@ -188,20 +182,23 @@ def _store_instance_variables( ], ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of variables saved + trainable_variables : tuple + A 3-tuple ``(W, sigma2, D)`` of fitted variables. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_features, n_components) + Fitted factor loadings (transposed for storage). sigma2_ : float - The isotropic variance + Fitted isotropic noise variance. + + D_ : ndarray of shape (n_confounders, n_components) + Fitted adversarial weight matrix (transposed for storage). """ self.W_ = trainable_variables[0].T self.sigma2_ = trainable_variables[1] @@ -277,20 +274,20 @@ def fit( self, X: NDArray[np.float_], Y: NDArray[np.float_] ) -> "PPCAsupervised": """ - Fits a model given covariates X + Fit the supervised PPCA model. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The covariate data. - Y : NDArray,(n_samples,n_predicted) - The data we want to predict + Y : ndarray of shape (n_samples, n_targets) + The target data to predict from the latent representation. Returns ------- - self : PPCAadversarial - The model + self : PPCAsupervised + Fitted estimator. """ self.mean_x = np.mean(X, axis=0) self.mean_y = np.mean(Y, axis=0) @@ -341,18 +338,14 @@ def fit( def get_covariance(self) -> NDArray[np.float_]: """ - Gets the covariance matrix - - Sigma = W^TW + sigma2*I + Compute the model covariance matrix. - Parameters - ---------- - None + The covariance is given by ``Sigma = W^T W + sigma^2 * I``. Returns ------- - covariance : NDArray,(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The PPCA 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") @@ -361,16 +354,14 @@ def get_covariance(self) -> NDArray[np.float_]: def get_noise(self): """ - Returns the observational noise as a diagonal matrix + Return the observational noise as a diagonal matrix. - Parameters - ---------- - None + For PPCA-based models the noise is isotropic: ``Lambda = sigma^2 * I``. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + The isotropic diagonal noise matrix. """ if self.sigma2_ is None or self.p is None: raise ValueError("Model has not been fit yet") @@ -384,20 +375,23 @@ def _store_instance_variables( ], ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of variables saved + trainable_variables : tuple + A 3-tuple ``(W, sigma2, Phi)`` of fitted variables. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_features, n_components) + Fitted factor loadings (transposed for storage). sigma2_ : float - The isotropic variance + Fitted isotropic noise variance. + + Phi_ : ndarray of shape (n_targets, n_components) + Fitted predictive weight matrix (transposed for storage). """ self.W_ = trainable_variables[0].T self.sigma2_ = trainable_variables[1] @@ -478,23 +472,23 @@ def fit( Z: NDArray[np.float_], ) -> "PPCASupAdversarial": """ - Fits a model given covariates X + Fit the simultaneous supervised and adversarial PPCA model. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The covariate data. - Y : NDArray,(n_samples,n_predictors) - The data we want to predict + Y : ndarray of shape (n_samples, n_targets) + The target data to predict from the latent representation. - Z : NDArray,(n_samples,n_confounders) - The concommitant data we want to remove + Z : ndarray of shape (n_samples, n_confounders) + The confounder data to decorrelate from the latent representation. Returns ------- - self : PPCA_sup_adversarial - The model + self : PPCASupAdversarial + Fitted estimator. """ self.mean_x = np.mean(X, axis=0) self.mean_y = np.mean(Y, axis=0) @@ -567,18 +561,14 @@ def fit( def get_covariance(self) -> NDArray[np.float_]: """ - Gets the covariance matrix - - Sigma = W^TW + sigma2*I + Compute the model covariance matrix. - Parameters - ---------- - None + The covariance is given by ``Sigma = W^T W + sigma^2 * I``. Returns ------- - covariance : NDArray,(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The PPCA 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") @@ -587,16 +577,14 @@ def get_covariance(self) -> NDArray[np.float_]: def get_noise(self): """ - Returns the observational noise as a diagonal matrix + Return the observational noise as a diagonal matrix. - Parameters - ---------- - None + For PPCA-based models the noise is isotropic: ``Lambda = sigma^2 * I``. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + The isotropic diagonal noise matrix. """ if self.sigma2_ is None or self.p is None: raise ValueError("Model has not been fit yet") @@ -613,20 +601,23 @@ def _store_instance_variables( ], ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of variables saved + trainable_variables : tuple + A 4-tuple ``(W, sigma2, D1, D2)`` of fitted variables. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_features, n_components) + Fitted factor loadings (transposed for storage). sigma2_ : float - The isotropic variance + Fitted isotropic noise variance. + + D_ : ndarray of shape (n_targets, n_components) + Fitted predictive weight matrix (transposed for storage). """ self.W_ = trainable_variables[0].T self.sigma2_ = trainable_variables[1] @@ -650,8 +641,25 @@ def _save_variables(self): def _select_covariance_estimator(regularization): """ - This is a tiny method for selecting the estimator for the covariance - matrix. + Instantiate a covariance estimator by name. + + Parameters + ---------- + regularization : str + Name of the covariance estimation method. One of: + ``'Empirical'``, ``'Linear'``, ``'LinearInverse'``, + ``'QuadraticInverse'``, ``'GeometricInverse'``, ``'NonLinear'``. + + Returns + ------- + model_cov : BaseCovariance + An unfitted covariance estimator instance. + + Raises + ------ + ValueError + If ``'Bayesian'`` is requested (currently unsupported) or if + ``regularization`` is not a recognized string. """ model_cov: Union[ EmpiricalCovariance, @@ -705,20 +713,24 @@ def fit( self, X: NDArray[np.float_], Y: NDArray[np.float_] ) -> "PPCAadversarialRandomized": """ - Fits a model given covariates X + Fit the randomized adversarial PPCA model. + + Uses a randomized range finder to project to a low-dimensional + subspace before fitting, reducing computation for high-dimensional + data. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The covariate data. - Y : NDArray,(n_samples,n_confounders) - The concommitant data we want to remove + Y : ndarray of shape (n_samples, n_confounders) + The confounder data to decorrelate from the latent representation. Returns ------- - self : PPCAadversarial - The model + self : PPCAadversarialRandomized + Fitted estimator. """ self.mean_x = np.mean(X, axis=0) self.mean_y = np.mean(Y, axis=0) @@ -786,18 +798,17 @@ def fit( def get_covariance(self) -> NDArray[np.float_]: """ - Gets the covariance matrix + Compute the model covariance matrix. - Sigma = W^TW + sigma2*I - - Parameters - ---------- - None + Uses the explained variance to reconstruct the PPCA-style + covariance: ``Sigma = W_mod W_mod^T + sigma^2 * I``, where + ``W_mod`` scales the loadings by the square root of the + explained-minus-noise variance. Returns ------- - covariance : NDArray,(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The model 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") @@ -808,16 +819,14 @@ def get_covariance(self) -> NDArray[np.float_]: def get_noise(self): """ - Returns the observational noise as a diagonal matrix + Return the observational noise as a diagonal matrix. - Parameters - ---------- - None + For PPCA-based models the noise is isotropic: ``Lambda = sigma^2 * I``. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + The isotropic diagonal noise matrix. """ if self.sigma2_ is None or self.p is None: raise ValueError("Model has not been fit yet") @@ -832,20 +841,20 @@ def _store_instance_variables( ], ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of variables saved + trainable_variables : tuple + A 2-tuple ``(W, sigma2)`` of fitted variables. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_features, n_components) + Fitted factor loadings (transposed for storage). sigma2_ : float - The isotropic variance + Fitted isotropic noise variance. """ self.W_ = trainable_variables[0].T self.sigma2_ = trainable_variables[1] @@ -891,20 +900,24 @@ def fit( self, X: NDArray[np.float_], Y: NDArray[np.float_] ) -> "PPCAsupervisedRandomized": """ - Fits a model given covariates X + Fit the randomized supervised PPCA model. + + Uses a randomized range finder to project to a low-dimensional + subspace before fitting, reducing computation for high-dimensional + data. Parameters ---------- - X : NDArray,(n_samples,n_covariates) - The data + X : ndarray of shape (n_samples, n_features) + The covariate data. - Y : NDArray,(n_samples,n_confounders) - The concommitant data we want to remove + Y : ndarray of shape (n_samples, n_targets) + The target data to predict from the latent representation. Returns ------- - self : PPCAadversarial - The model + self : PPCAsupervisedRandomized + Fitted estimator. """ self.mean_x = np.mean(X, axis=0) self.mean_y = np.mean(Y, axis=0) @@ -972,18 +985,17 @@ def fit( def get_covariance(self) -> NDArray[np.float_]: """ - Gets the covariance matrix - - Sigma = W^TW + sigma2*I + Compute the model covariance matrix. - Parameters - ---------- - None + Uses the explained variance to reconstruct the PPCA-style + covariance: ``Sigma = W_mod W_mod^T + sigma^2 * I``, where + ``W_mod`` scales the loadings by the square root of the + explained-minus-noise variance. Returns ------- - covariance : NDArray,(p,p) - The covariance matrix + covariance : ndarray of shape (n_features, n_features) + The model 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") @@ -994,16 +1006,14 @@ def get_covariance(self) -> NDArray[np.float_]: def get_noise(self): """ - Returns the observational noise as a diagonal matrix + Return the observational noise as a diagonal matrix. - Parameters - ---------- - None + For PPCA-based models the noise is isotropic: ``Lambda = sigma^2 * I``. Returns ------- - Lambda : NDArray,(p,p) - The observational noise + Lambda : ndarray of shape (n_features, n_features) + The isotropic diagonal noise matrix. """ if self.sigma2_ is None or self.p is None: raise ValueError("Model has not been fit yet") @@ -1018,20 +1028,20 @@ def _store_instance_variables( ], ) -> None: """ - Saves the learned variables + Store the learned parameters as numpy arrays on the instance. Parameters ---------- - trainable_variables : list - List of variables saved + trainable_variables : tuple + A 2-tuple ``(W, sigma2)`` of fitted variables. - Sets - ---- - W_ : NDArray,(n_components,p) - The loadings + Attributes + ---------- + W_ : ndarray of shape (n_features, n_components) + Fitted factor loadings (transposed for storage). sigma2_ : float - The isotropic variance + Fitted isotropic noise variance. """ self.W_ = trainable_variables[0].T self.sigma2_ = trainable_variables[1]