scdesigner 0.0.5__py3-none-any.whl → 0.0.10__py3-none-any.whl

scdesigner/minimal/marginal.py

@@ -1,154 +0,0 @@
-from .kwargs import DEFAULT_ALLOWED_KWARGS, _filter_kwargs
-from .loader import adata_loader, get_device
-from anndata import AnnData
-from typing import Union, Dict, Optional, Tuple
-import pandas as pd
-import torch
-import torch.nn as nn
-from abc import ABC, abstractmethod
-
-
-class Marginal(ABC):
-    def __init__(self, formula: Union[Dict, str]):
-        self.formula = formula
-        self.feature_dims = None
-        self.loader = None
-        self.n_outcomes = None
-        self.predict = None
-        self.predictor_names = None
-        self.parameters = None
-        self.device = get_device()
-
-    def setup_data(self, adata: AnnData, batch_size: int = 1024, **kwargs):
-        """Set up the dataloader for the AnnData object."""
-        # keep a reference to the AnnData for later use (e.g., var_names)
-        self.adata = adata
-        self.loader = adata_loader(adata, self.formula, batch_size=batch_size, **kwargs)
-        X_batch, obs_batch = next(iter(self.loader))
-        self.n_outcomes = X_batch.shape[1]
-        self.feature_dims = {k: v.shape[1] for k, v in obs_batch.items()}
-        self.predictor_names = self.loader.dataset.predictor_names
-
-    def fit(self, max_epochs: int = 100, **kwargs):
-        """Fit the marginal predictor using vanilla PyTorch training loop."""
-        if self.predict is None:
-            self.setup_optimizer(**kwargs)
-
-        for epoch in range(max_epochs):
-            epoch_loss, n_batches = 0.0, 0
-
-            for batch in self.loader:
-                y, x = batch
-                if y.device != self.device:
-                    y = y.to(self.device)
-                    x = {k: v.to(self.device) for k, v in x.items()}
-
-                self.predict.optimizer.zero_grad()
-                loss = self.predict.loss_fn((y, x))
-                loss.backward()
-                self.predict.optimizer.step()
-
-                epoch_loss += loss.item()
-                n_batches += 1
-
-            avg_loss = epoch_loss / n_batches
-            print(f"Epoch {epoch}/{max_epochs}, Loss: {avg_loss:.4f}", end='\r')
-        self.parameters = self.format_parameters()
-
-    def format_parameters(self):
-        """Convert fitted coefficient tensors into pandas DataFrames.
-
-        Returns:
-            dict: mapping from parameter name -> pandas.DataFrame with rows
-                corresponding to predictor column names (from
-                `self.predictor_names[param]`) and columns corresponding to
-                `self.adata.var_names` (gene names). The values are moved to
-                CPU and converted to numpy floats.
-        """
-        var_names = list(self.adata.var_names)
-
-        dfs = {}
-        for param, tensor in self.predict.coefs.items():
-            coef_np = tensor.detach().cpu().numpy()
-            row_names = list(self.predictor_names[param])
-            dfs[param] = pd.DataFrame(coef_np, index=row_names, columns=var_names)
-        return dfs
-
-    def num_params(self):
-        """Return the number of parameters."""
-        if self.predict is None:
-            return 0
-        return sum(p.numel() for p in self.predict.parameters() if p.requires_grad)
-
-    @abstractmethod
-    def setup_optimizer(self, **kwargs):
-        raise NotImplementedError
-
-    @abstractmethod
-    def likelihood(self, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
-        """Compute the (negative) log-likelihood or loss for a batch.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
-        """Invert pseudoobservations."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor]):
-        """Uniformize using learned CDF.
-        """
-        raise NotImplementedError
-
-
-class GLMPredictor(nn.Module):
-    """GLM-style predictor with arbitrary named parameters.
-
-    Args:
-        n_outcomes: number of model outputs (e.g. genes)
-        feature_dims: mapping from param name -> number of covariate features
-        link_fns: optional mapping from param name -> callable(link) applied to linear predictor
-
-    The module will create one coefficient matrix per named parameter with shape
-    (n_features_for_param, n_outcomes) and expose them as Parameters under
-    `self.coefs[param_name]`.
-    """
-    def __init__(
-        self,
-        n_outcomes: int,
-        feature_dims: Dict[str, int],
-        link_fns: Dict[str, callable] = None,
-        loss_fn: Optional[callable] = None,
-        optimizer_class: Optional[callable] = torch.optim.Adam,
-        optimizer_kwargs: Optional[Dict] = None,
-    ):
-        super().__init__()
-        self.n_outcomes = int(n_outcomes)
-        self.feature_dims = dict(feature_dims)
-        self.param_names = list(self.feature_dims.keys())
-
-        self.link_fns = link_fns or {k: torch.exp for k in self.param_names}
-        self.coefs = nn.ParameterDict()
-        for key, dim in self.feature_dims.items():
-            self.coefs[key] = nn.Parameter(torch.zeros(dim, self.n_outcomes))
-        self.reset_parameters()
-
-        self.loss_fn = loss_fn
-        self.to(get_device())
-
-        optimizer_kwargs = optimizer_kwargs or {}
-        filtered_kwargs = _filter_kwargs(optimizer_kwargs, DEFAULT_ALLOWED_KWARGS['optimizer'])
-        self.optimizer = optimizer_class(self.parameters(), **filtered_kwargs)
-
-    def reset_parameters(self):
-        for p in self.coefs.values():
-            nn.init.normal_(p, mean=0.0, std=1e-4)
-
-    def forward(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
-        out = {}
-        for name in self.param_names:
-            x_beta = obs_dict[name] @ self.coefs[name]
-            link = self.link_fns.get(name, torch.exp)
-            out[name] = link(x_beta)
-        return out

scdesigner/minimal/negbin.py

@@ -1,73 +0,0 @@
-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 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

@@ -1,231 +0,0 @@
-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

@@ -1,96 +0,0 @@
-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
-from abc import ABC, abstractmethod
-
-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

@@ -1,50 +0,0 @@
-from .scd3 import SCD3Simulator
-from .negbin import NegBin
-from .zero_inflated_negbin import ZeroInflatedNegBin
-from .gaussian import Gaussian
-from .standard_copula import StandardCopula
-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 = StandardCopula(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 = StandardCopula(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 = StandardCopula(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 = StandardCopula(copula_formula)
-        super().__init__(marginal, covariance)

scdesigner/minimal/simulator.py

@@ -1,25 +0,0 @@
-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