scdesigner 0.0.1__py3-none-any.whl

scdesigner/minimal/negbin.py

@@ -0,0 +1,73 @@
+from .formula import standardize_formula
+from .marginal import GLMPredictor, Marginal
+from .loader import _to_numpy
+from typing import Union, Dict, Optional
+import torch
+import numpy as np
+from scipy.stats import nbinom
+
+class NegBin(Marginal):
+    """Negative-binomial marginal estimator"""
+    def __init__(self, formula: Union[Dict, str]):
+        formula = standardize_formula(formula, allowed_keys=['mean', 'dispersion'])
+        super().__init__(formula)
+
+    def setup_optimizer(
+            self,
+            optimizer_class: Optional[callable] = torch.optim.Adam,
+            **optimizer_kwargs,
+    ):
+        if self.loader is None:
+            raise RuntimeError("self.loader is not set (call setup_data first)")
+
+        nll = lambda batch: -self.likelihood(batch).sum()
+        self.predict = GLMPredictor(
+            n_outcomes=self.n_outcomes,
+            feature_dims=self.feature_dims,
+            loss_fn=nll,
+            optimizer_class=optimizer_class,
+            optimizer_kwargs=optimizer_kwargs
+        )
+
+    def likelihood(self, batch):
+        """Compute the negative log-likelihood"""
+        y, x = batch
+        params = self.predict(x)
+        mu = params.get('mean')
+        r = params.get('dispersion')
+        return (
+            torch.lgamma(y + r)
+            - torch.lgamma(r)
+            - torch.lgamma(y + 1.0)
+            + r * torch.log(r)
+            + y * torch.log(mu)
+            - (r + y) * torch.log(r + mu)
+        )
+
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
+        """Invert pseudoobservations."""
+        mu, r, u = self._local_params(x, u)
+        p = r / (r + mu)
+        y = nbinom(n=r, p=p).ppf(u)
+        return torch.from_numpy(y).float()
+
+    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor], epsilon=1e-6):
+        """Return uniformized pseudo-observations for counts y given covariates x."""
+        # cdf values using scipy's parameterization
+        mu, r, y = self._local_params(x, y)
+        p = r / (r + mu)
+        u1 = nbinom(n=r, p=p).cdf(y)
+        u2 = np.where(y > 0, nbinom(n=r, p=p).cdf(y - 1), 0.0)
+
+        # randomize within discrete mass to get uniform(0,1)
+        v = np.random.uniform(size=y.shape)
+        u = np.clip(v * u1 + (1.0 - v) * u2, epsilon, 1.0 - epsilon)
+        return torch.from_numpy(u).float()
+
+    def _local_params(self, x, y=None):
+        params = self.predict(x)
+        mu = params.get('mean')
+        r = params.get('dispersion')
+        if y is None:
+            return _to_numpy(mu, r)
+        return _to_numpy(mu, r, y)

scdesigner/minimal/positive_nonnegative_matrix_factorization.py

@@ -0,0 +1,231 @@
+from .formula import standardize_formula
+from .loader import _to_numpy
+from .simulator import Simulator
+from anndata import AnnData
+from formulaic import model_matrix
+from scipy.stats import gamma
+from typing import Union, Dict
+import numpy as np
+import pandas as pd
+import torch
+
+################################################################################
+## Functions for estimating PNMF regression
+################################################################################
+
+# computes PNMF weight and score, ncol specify the number of clusters
+def pnmf(log_data, nbase=3, **kwargs):  # data is np array, log transformed read data
+    """
+    Computes PNMF weight and score.
+
+    :log_data: log transformed np array of read data
+    :ncol: specify the number of clusters
+    :return: W (weights, gene x base) and S (scores, base x cell) as numpy arrays
+    """
+    U = left_singular(log_data, nbase)
+    W = pnmf_eucdist(log_data, U, **kwargs)
+    W = W / np.linalg.norm(W, ord=2)
+    S = W.T @ log_data
+    return W, S
+
+
+def gamma_regression_array(
+    x: np.array, y: np.array, lr: float = 0.1, epochs: int = 40
+) -> dict:
+    x = torch.tensor(x, dtype=torch.float32)
+    y = torch.tensor(y, dtype=torch.float32)
+
+    n_features, n_outcomes = x.shape[1], y.shape[1]
+    a = torch.zeros(n_features * n_outcomes, requires_grad=True)
+    loc = torch.zeros(n_features * n_outcomes, requires_grad=True)
+    beta = torch.zeros(n_features * n_outcomes, requires_grad=True)
+    optimizer = torch.optim.Adam([a, loc, beta], lr=lr)
+
+    for i in range(epochs):
+        optimizer.zero_grad()
+        loss = negative_gamma_log_likelihood(a, beta, loc, x, y)
+        loss.backward()
+        optimizer.step()
+
+    a, loc, beta = _to_numpy(a, loc, beta)
+    a = a.reshape(n_features, n_outcomes)
+    loc = loc.reshape(n_features, n_outcomes)
+    beta = beta.reshape(n_features, n_outcomes)
+    return {"a": a, "loc": loc, "beta": beta}
+
+
+def class_generator(score, n_clusters=3):
+    """
+    Generates one-hot encoding for score classes
+    """
+    from sklearn.cluster import KMeans
+    kmeans = KMeans(n_clusters, random_state=0)  # Specify the number of clusters
+    kmeans.fit(score.T)
+    labels = kmeans.labels_
+    num_classes = len(np.unique(labels))
+    one_hot = np.eye(num_classes)[labels].astype(int)
+    return labels
+
+
+###############################################################################
+## Helpers for deriving PNMF
+###############################################################################
+
+
+def pnmf_eucdist(X, W_init, maxIter=500, threshold=1e-4, tol=1e-10, verbose=False, **kwargs):
+    # initialization
+    W = W_init  # initial W is the PCA of X
+    XX = X @ X.T
+
+    # iterations
+    for iter in range(maxIter):
+        if verbose and (iter + 1) % 10 == 0:
+            print("%d iterations used." % (iter + 1))
+        W_old = W
+
+        XXW = XX @ W
+        SclFactor = np.dot(W, W.T @ XXW) + np.dot(XXW, W.T @ W)
+
+        # QuotientLB
+        SclFactor = MatFindlb(SclFactor, tol)
+        SclFactor = XXW / SclFactor
+        W = W * SclFactor  # somehow W *= SclFactor doesn't work?
+
+        norm_W = np.linalg.norm(W)
+        W /= norm_W
+        W = MatFind(W, tol)
+
+        diffW = np.linalg.norm(W_old - W) / np.linalg.norm(W_old)
+        if diffW < threshold:
+            break
+
+    return W
+
+
+# left singular vector of X
+def left_singular(X, k):
+    from scipy.sparse.linalg import svds
+    U, _, _ = svds(X, k=k)
+    return np.abs(U)
+
+
+def MatFindlb(A, lb):
+    B = np.ones(A.shape) * lb
+    Alb = np.where(A < lb, B, A)
+    return Alb
+
+
+def MatFind(A, ZeroThres):
+    B = np.zeros(A.shape)
+    Atrunc = np.where(A < ZeroThres, B, A)
+    return Atrunc
+
+
+###############################################################################
+## Helpers for training PNMF regression
+###############################################################################
+
+
+def shifted_gamma_pdf(x, alpha, beta, loc):
+    if not torch.is_tensor(x):
+        x = torch.tensor(x)
+    mask = x < loc
+    y_clamped = torch.clamp(x - loc, min=1e-12)
+
+    log_pdf = (
+        alpha * torch.log(beta)
+        - torch.lgamma(alpha)
+        + (alpha - 1) * torch.log(y_clamped)
+        - beta * y_clamped
+    )
+    loss = -torch.mean(log_pdf[~mask])
+    n_invalid = mask.sum()
+    if n_invalid > 0:  # force samples to be greater than loc
+        loss = loss + 1e10 * n_invalid.float()
+    return loss
+
+
+def negative_gamma_log_likelihood(log_a, log_beta, loc, X, y):
+    n_features = X.shape[1]
+    n_outcomes = y.shape[1]
+
+    a = torch.exp(log_a.reshape(n_features, n_outcomes))
+    beta = torch.exp(log_beta.reshape(n_features, n_outcomes))
+    loc = loc.reshape(n_features, n_outcomes)
+    return shifted_gamma_pdf(y, X @ a, X @ beta, X @ loc)
+
+def format_gamma_parameters(
+    parameters: dict,
+    W_index: list,
+    coef_index: list,
+) -> dict:
+    parameters["a"] = pd.DataFrame(parameters["a"], index=coef_index)
+    parameters["loc"] = pd.DataFrame(parameters["loc"], index=coef_index)
+    parameters["beta"] = pd.DataFrame(parameters["beta"], index=coef_index)
+    parameters["W"] = pd.DataFrame(parameters["W"], index=W_index)
+    return parameters
+
+
+################################################################################
+## Associated PNMF Objects
+################################################################################
+
+class PositiveNMF(Simulator):
+    """Positive nonnegative matrix factorization marginal estimator"""
+    def __init__(self, formula: Union[Dict, str], **kwargs):
+        self.formula = standardize_formula(formula, allowed_keys=['mean'])
+        self.parameters = None
+        self.hyperparams = kwargs
+
+
+    def setup_data(self, adata: AnnData, **kwargs):
+        self.log_data = np.log1p(adata.X).T
+        self.n_outcomes = self.log_data.shape[1]
+        self.template = adata
+        self.x = model_matrix(self.formula["mean"], adata.obs)
+        self.columns = self.x.columns
+        self.x = np.asarray(self.x)
+
+
+    def fit(self, adata: AnnData, lr: float=0.1):
+        self.setup_data(adata)
+        W, S = pnmf(self.log_data, **self.hyperparams)
+        parameters = gamma_regression_array(self.x, S.T, lr)
+        parameters["W"] = W
+        self.parameters = format_gamma_parameters(
+            parameters, list(self.template.var_names), list(self.columns)
+        )
+
+
+    def predict(self, obs=None, **kwargs):
+        """Predict from an obs dataframe"""
+        if obs is None:
+            obs = self.template.obs
+
+        x = model_matrix(self.formula["mean"], obs)
+        a, loc, beta = (
+            x @ np.exp(self.parameters["a"]),
+            x @ self.parameters["loc"],
+            x @ np.exp(self.parameters["beta"]),
+        )
+        return {"a": a, "loc": loc, "beta": beta}
+
+
+    def sample(self, obs=None):
+        """Generate samples."""
+        if obs is None:
+            obs = self.template.obs
+        W = self.parameters["W"]
+        parameters = self.predict(obs)
+        a, loc, beta = parameters["a"], parameters["loc"], parameters["beta"]
+        sim_score = gamma(a, loc, 1 / beta).rvs()
+        samples = np.exp(W @ sim_score.T).T
+
+        # thresholding samples
+        floor = np.floor(samples)
+        samples = floor + np.where(samples - floor < 0.9, 0, 1) - 1
+        samples = np.where(samples < 0, 0, samples)
+
+        result = AnnData(X=samples, obs=obs)
+        result.var_names = self.template.var_names
+        return result

scdesigner/minimal/scd3.py

@@ -0,0 +1,95 @@
+from .copula import Copula
+from .loader import obs_loader, adata_loader
+from .marginal import Marginal
+from .simulator import Simulator
+from anndata import AnnData
+from tqdm import tqdm
+import torch
+import numpy as np
+
+class SCD3Simulator(Simulator):
+    """Simulation wrapper"""
+
+    def __init__(self, marginal: Marginal, copula: Copula):
+        self.marginal = marginal
+        self.copula = copula
+        self.template = None
+        self.parameters = None
+
+    def fit(
+        self,
+        adata: AnnData,
+        **kwargs):
+        """Fit the simulator"""
+        self.template = adata
+        self.marginal.setup_data(adata, **kwargs)
+        self.marginal.setup_optimizer(**kwargs)
+        self.marginal.fit(**kwargs)
+
+        # copula simulator
+        self.copula.setup_data(adata, self.marginal.formula, **kwargs)
+        self.copula.fit(self.marginal.uniformize, **kwargs)
+        self.parameters = {
+            "marginal": self.marginal.parameters,
+            "copula": self.copula.parameters
+        }
+
+    def predict(self, obs=None, batch_size: int = 1000, **kwargs):
+        """Predict from an obs dataframe"""
+        # prepare an internal data loader for this obs
+        if obs is None:
+            obs = self.template.obs
+        loader = obs_loader(
+            obs,
+            self.marginal.formula,
+            batch_size=batch_size,
+            **kwargs
+        )
+
+        # get predictions across batches
+        local_parameters = []
+        for _, x_dict in loader:
+            l = self.marginal.predict(x_dict)
+            local_parameters.append(l)
+
+        # convert to a merged dictionary
+        keys = list(local_parameters[0].keys())
+        return {
+            k: torch.cat([d[k] for d in local_parameters]).detach().cpu().numpy()
+            for k in keys
+        }
+
+    def sample(self, obs=None, batch_size: int = 1000, **kwargs):
+        """Generate samples."""
+        if obs is None:
+            obs = self.template.obs
+        loader = obs_loader(
+            obs,
+            self.copula.formula | self.marginal.formula,
+            batch_size=batch_size,
+            **kwargs
+        )
+
+        # get samples across batches
+        samples = []
+        for _, x_dict in loader:
+            u = self.copula.pseudo_obs(x_dict)
+            u = torch.from_numpy(u)
+            samples.append(self.marginal.invert(u, x_dict))
+        samples = torch.cat(samples).detach().cpu().numpy()
+        return AnnData(X = samples, obs=obs)
+
+    def complexity(self, adata: AnnData = None, **kwargs):
+        if adata is None:
+            adata = self.template
+
+        N, ll = 0, 0
+        loader = adata_loader(adata, self.marginal.formula | self.copula.formula, **kwargs)
+        for batch in tqdm(loader, desc="Computing log-likelihood..."):
+            ll += self.copula.likelihood(self.marginal.uniformize, batch).sum()
+            N += len(batch[0])
+
+        return {
+            "aic": -2 * ll + 2 * self.copula.num_params(),
+            "bic": -2 * ll + np.log(N) * self.copula.num_params()
+        }

scdesigner/minimal/scd3_instances.py

@@ -0,0 +1,50 @@
+from .scd3 import SCD3Simulator
+from .negbin import NegBin
+from .zero_inflated_negbin import ZeroInflatedNegBin
+from .gaussian import Gaussian
+from .standard_covariance import StandardCovariance
+from typing import Optional
+
+
+class NegBinCopula(SCD3Simulator):
+    def __init__(self,
+                 mean_formula: Optional[str] = None,
+                 dispersion_formula: Optional[str] = None,
+                 copula_formula: Optional[str] = None) -> None:
+        marginal = NegBin({"mean": mean_formula, "dispersion": dispersion_formula})
+        covariance = StandardCovariance(copula_formula)
+        super().__init__(marginal, covariance)
+
+
+class ZeroInflatedNegBinCopula(SCD3Simulator):
+    def __init__(self,
+                 mean_formula: Optional[str] = None,
+                 dispersion_formula: Optional[str] = None,
+                 zero_inflation_formula: Optional[str] = None,
+                 copula_formula: Optional[str] = None) -> None:
+        marginal = ZeroInflatedNegBin({
+            "mean": mean_formula,
+            "dispersion": dispersion_formula,
+            "zero_inflation_formula": zero_inflation_formula
+        })
+        covariance = StandardCovariance(copula_formula)
+        super().__init__(marginal, covariance)
+
+
+class BernoulliCopula(SCD3Simulator):
+    def __init__(self,
+                 mean_formula: Optional[str] = None,
+                 copula_formula: Optional[str] = None) -> None:
+        marginal = NegBin({"mean": mean_formula})
+        covariance = StandardCovariance(copula_formula)
+        super().__init__(marginal, covariance)
+
+
+class GaussianCopula(SCD3Simulator):
+    def __init__(self,
+                 mean_formula: Optional[str] = None,
+                 sdev_formula: Optional[str] = None,
+                 copula_formula: Optional[str] = None) -> None:
+        marginal = Gaussian({"mean": mean_formula, "sdev": sdev_formula})
+        covariance = StandardCovariance(copula_formula)
+        super().__init__(marginal, covariance)

scdesigner/minimal/simulator.py

@@ -0,0 +1,25 @@
+from anndata import AnnData
+from typing import Dict
+from pandas import DataFrame
+from abc import abstractmethod
+
+class Simulator:
+    """Simulation abstract class"""
+
+    def __init__(self):
+        self.parameters = None
+
+    @abstractmethod
+    def fit(self, anndata: AnnData, **kwargs) -> None:
+        """Fit the simulator"""
+        self.template = anndata
+
+    @abstractmethod
+    def predict(self, obs: DataFrame=None, **kwargs) -> Dict:
+        """Predict from an obs dataframe"""
+        pass
+
+    @abstractmethod
+    def sample(self, obs: DataFrame=None, **kwargs) -> AnnData:
+        """Generate samples."""
+        pass

scdesigner/minimal/standard_covariance.py

@@ -0,0 +1,124 @@
+from .copula import Copula
+from .formula import standardize_formula
+from .kwargs import DEFAULT_ALLOWED_KWARGS, _filter_kwargs
+from anndata import AnnData
+from scipy.stats import norm, multivariate_normal
+from tqdm import tqdm
+from typing import Dict, Union, Callable, Tuple
+import numpy as np
+import pandas as pd
+import torch
+
+
+class StandardCovariance(Copula):
+    def __init__(self, formula: str = "~ 1"):
+        formula = standardize_formula(formula, allowed_keys=['group'])
+        super().__init__(formula)
+        self.groups = None
+
+
+    def setup_data(self, adata: AnnData, marginal_formula: Dict[str, str], **kwargs):
+        data_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['data'])
+        super().setup_data(adata, marginal_formula, **data_kwargs)
+        _, obs_batch = next(iter(self.loader))
+        obs_batch_group = obs_batch.get("group")
+
+        # fill in group indexing variables
+        self.groups = self.loader.dataset.predictor_names["group"]
+        self.n_groups = len(self.groups)
+        self.group_col = {g: i for i, g in enumerate(self.groups)}
+
+        # check that obs_batch is a binary grouping matrix
+        unique_vals = torch.unique(obs_batch_group)
+        if (not torch.all((unique_vals == 0) | (unique_vals == 1)).item()):
+            raise ValueError("Only categorical groups are currently supported in copula covariance estimation.")
+
+    def fit(self, uniformizer: Callable, **kwargs):
+        sums = {g: np.zeros(self.n_outcomes) for g in self.groups}
+        second_moments = {g: np.eye(self.n_outcomes) for g in self.groups}
+        Ng = {g: 0 for g in self.groups}
+
+        for y, x_dict in tqdm(self.loader, desc="Estimating copula covariance"):
+            memberships = x_dict.get("group").numpy()
+            u = uniformizer(y, x_dict)
+
+            for g in self.groups:
+                ix = np.where(memberships[:, self.group_col[g]] == 1)
+                z = norm().ppf(u[ix])
+                second_moments[g] += z.T @ z
+                sums[g] += z.sum(axis=0)
+                Ng[g] += len(ix[0])
+
+        covariances = {}
+        for g in self.groups:
+            mean = sums[g] / Ng[g]
+            covariances[g] = second_moments[g] / Ng[g] - np.outer(mean, mean)
+
+        if len(self.groups) == 1:
+            covariances = list(covariances.values())[0]
+        self.parameters = self.format_parameters(covariances)
+
+    def format_parameters(self, covariances: Union[Dict, np.array]):
+        var_names = self.adata.var_names
+        def to_df(mat):
+            return pd.DataFrame(mat, index=var_names, columns=var_names)
+
+        if isinstance(covariances, dict):
+            formatted = {}
+            for k, v in covariances.items():
+                formatted[k] = to_df(v)
+            covariances = formatted
+            return covariances
+
+        if isinstance(covariances, (np.ndarray, list, tuple)):
+            covariances = to_df(covariances)
+        return covariances
+
+    def pseudo_obs(self, x_dict: Dict):
+        # convert one-hot encoding memberships to a map
+        #      {"group1": [indices of group 1], "group2": [indices of group 2]}
+        memberships = x_dict.get("group").numpy()
+        group_ix = {g: np.where(memberships[:, self.group_col[g] == 1])[0] for g in self.groups}
+
+        # initialize the result
+        u = np.zeros((len(memberships), self.n_outcomes))
+        parameters = self.parameters
+        if type(parameters) is not dict:
+            parameters = {group: parameters}
+
+        # loop over groups and sample each part in turn
+        for group, sigma in parameters.items():
+            z = np.random.multivariate_normal(
+                mean=np.zeros(self.n_outcomes),
+                cov=sigma,
+                size=len(group_ix[group])
+            )
+            normal_distn = norm(0, np.diag(sigma) ** 0.5)
+            u[group_ix[group]] = normal_distn.cdf(z)
+        return u
+
+    def likelihood(self, uniformizer: Callable, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
+        # uniformize the observations
+        y, x_dict = batch
+        u = uniformizer(y, x_dict)
+        z = norm().ppf(u)
+
+        # same group manipulation as for pseudobs
+        parameters = self.parameters
+        if type(parameters) is not dict:
+            parameters = {group: parameters}
+
+        memberships = x_dict.get("group").numpy()
+        group_ix = {g: np.where(memberships[:, self.group_col[g] == 1])[0] for g in self.groups}
+        ll = np.zeros(len(z))
+        for group, sigma in parameters.items():
+            ix = group_ix[group]
+            if len(ix) > 0:
+                copula_ll = multivariate_normal.logpdf(z[ix], np.zeros(sigma.shape[0]), sigma)
+                ll[ix] = copula_ll - norm.logpdf(z[ix]).sum(axis=1)
+        return ll
+
+    def num_params(self, **kwargs):
+        S = self.parameters
+        per_group = [(np.sum(S[g].values != 0) - S[g].shape[0]) / 2 for g in self.groups]
+        return sum(per_group)