scdesigner 0.0.1__tar.gz

scdesigner-0.0.1/.gitignore

@@ -0,0 +1,8 @@
+build*
+*egg*
+*cache*
+examples/data/*.h5ad
+examples/data
+*logs*
+*DS_store
+tests_langtian/

scdesigner-0.0.1/PKG-INFO

@@ -0,0 +1,23 @@
+Metadata-Version: 2.4
+Name: scdesigner
+Version: 0.0.1
+Summary: Interactive simulation for rigorous and transparent multi-omics analysis.
+Project-URL: Homepage, https://github.com/krisrs1128/scDesigner/
+Project-URL: Issues, https://github.com/krisrs1128/scDesigner/Issues/
+Author-email: Kris Sankaran <ksankaran@wisc.edu>
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.8
+Requires-Dist: formulaic
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: rich
+Requires-Dist: scanpy
+Requires-Dist: scipy
+Requires-Dist: torch
+Requires-Dist: tqdm
+Description-Content-Type: text/markdown
+
+
+### scDesigner

scdesigner-0.0.1/README.md

@@ -0,0 +1,2 @@
+
+### scDesigner

scdesigner-0.0.1/pyproject.toml

@@ -0,0 +1,35 @@
+[project]
+name = "scdesigner"
+version = "0.0.1"
+authors = [
+  { name="Kris Sankaran", email="ksankaran@wisc.edu" },
+]
+description = "Interactive simulation for rigorous and transparent multi-omics analysis."
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "formulaic",
+    "numpy",
+    "pandas",
+    "rich",
+    "scanpy",
+    "scipy",
+    "torch",
+    "tqdm"
+]
+
+[project.urls]
+Homepage = "https://github.com/krisrs1128/scDesigner/"
+Issues = "https://github.com/krisrs1128/scDesigner/Issues/"
+
+[build-system]
+requires = ["formulaic", "pandas", "numpy", "scipy", "hatchling", "torch", "rich"]
+build-backend = "hatchling.build"
+
+[tool.hatch.metadata]
+allow-direct-references = true

scdesigner-0.0.1/src/scdesigner/data/__init__.py

@@ -0,0 +1,16 @@
+from .formula import FormulaViewDataset, formula_loader, multiple_formula_loader, standardize_formula
+from .group import FormulaGroupViewDataset, formula_group_loader, stack_collate, multiple_formula_group_loader
+from .sparse import SparseMatrixDataset, SparseMatrixLoader
+
+__all__ = [
+    "FormulaViewDataset",
+    "SparseMatrixDataset",
+    "SparseMatrixLoader",
+    "FormulaGroupViewDataset",
+    "formula_loader",
+    "formula_group_loader",
+    "stack_collate",
+    "multiple_formula_loader",
+    "multiple_formula_group_loader",
+    "standardize_formula"
+]

scdesigner-0.0.1/src/scdesigner/data/formula.py

@@ -0,0 +1,137 @@
+from anndata import AnnData
+from formulaic import model_matrix
+from torch.utils.data import DataLoader, TensorDataset
+from typing import Union
+import numpy as np
+import pandas as pd
+import scipy.sparse
+import torch
+import torch.utils.data as td
+import warnings
+
+def formula_loader(
+    adata: AnnData, formula=None, chunk_size=int(1e4), batch_size: int = None
+):
+    device = check_device()
+    if adata.isbacked:
+        ds = FormulaViewDataset(adata, formula, chunk_size, device)
+        dataloader = td.DataLoader(ds, batch_size=batch_size)
+        ds.x_names = model_matrix_names(adata, formula, ds.categories)
+    else:
+        # convert sparse to dense matrix
+        y = adata.X
+        if isinstance(y, scipy.sparse._csc.csc_matrix):
+            y = y.todense()
+
+        # create tensor-based loader
+        x = model_matrix(formula, pd.DataFrame(adata.obs))
+        ds = TensorDataset(
+            torch.tensor(np.array(x), dtype=torch.float32).to(device),
+            torch.tensor(y, dtype=torch.float32).to(device),
+        )
+        ds.x_names = list(x.columns)
+        dataloader = DataLoader(ds, batch_size=batch_size, shuffle=False)
+
+    return dataloader
+
+def multiple_formula_loader(
+    adata: AnnData, formulas: dict, chunk_size=int(1e4), batch_size: int = None
+):
+    dataloaders = {}
+    for key in formulas.keys():
+        dataloaders[key] = formula_loader(adata, formulas[key], chunk_size, batch_size)
+    return dataloaders
+
+class FormulaViewDataset(td.Dataset):
+    def __init__(self, view, formula=None, chunk_size=int(1e4), device=None):
+        super().__init__()
+        self.view = view
+        self.formula = formula
+        self.len = len(view)
+        self.cur_range = range(0, min(self.len, chunk_size))
+        self.categories = column_levels(view.obs)
+        self.x = None
+        self.y = None
+        self.device = device or check_device()
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, ix):
+        if self.x is None or ix not in self.cur_range:
+            self.cur_range = range(ix, min(ix + len(self.cur_range), self.len))
+            view_inmem = self.view[self.cur_range].to_memory()
+            self.x = safe_model_matrix(
+                view_inmem.obs, self.formula, self.categories
+            ).to(self.device)
+            self.y = torch.from_numpy(view_inmem.X.toarray().astype(np.float32)).to(
+                self.device
+            )
+        return self.x[ix - self.cur_range[0]], self.y[ix - self.cur_range[0]]
+
+
+def replace_cols(obs, categories):
+    for k in obs.columns:
+        if str(obs[k].dtype) == "category":
+            obs[k] = obs[k].astype(pd.CategoricalDtype(categories[k]))
+    return obs
+
+
+def model_matrix_names(adata, formula, categories):
+    if adata.isbacked:
+        obs = adata[:1].to_memory().obs
+    else:
+        obs = adata.obs
+
+    if formula is None:
+        return list(obs.columns)
+
+    obs = replace_cols(obs, categories)
+    return list(model_matrix(formula, pd.DataFrame(obs)).columns)
+
+
+def safe_model_matrix(obs, formula, categories):
+    if formula is None:
+        return obs
+
+    obs = replace_cols(obs, categories)
+    x = model_matrix(formula, pd.DataFrame(obs))
+    return torch.from_numpy(np.array(x).astype(np.float32))
+
+
+def column_levels(obs):
+    categories = {}
+    for k in obs.columns:
+        obs_type = str(obs[k].dtype)
+        if obs_type in ["category", "object"]:
+            categories[k] = obs[k].unique()
+    return categories
+
+
+def check_device():
+    return torch.device(
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps" if torch.backends.mps.is_available() else "cpu"
+    )
+
+def standardize_formula(formula: Union[str, dict], allowed_keys = None):
+    # The first element of allowed_keys should be the name of default parameter
+    if allowed_keys is None:
+        raise ValueError("Internal error: allowed_keys must be specified")
+    formula = {allowed_keys[0]: formula} if isinstance(formula, str) else formula
+
+    formula_keys = set(formula.keys())
+    allowed_keys = set(allowed_keys)
+
+    if not formula_keys & allowed_keys:
+        raise ValueError(f"formula must have at least one of the following keys: {allowed_keys}")
+    
+    if extra_keys := formula_keys - allowed_keys:
+        warnings.warn(
+            f"Invalid formulas in dictionary will not be used: {extra_keys}",
+            UserWarning,
+        )
+    
+    formula.update({k: '~ 1' for k in allowed_keys - formula_keys})
+    return formula

scdesigner-0.0.1/src/scdesigner/data/group.py

@@ -0,0 +1,123 @@
+from . import formula as fl
+from anndata import AnnData
+from formulaic import model_matrix
+import numpy as np
+import pandas as pd
+import scipy
+import torch
+import torch.utils.data as td
+
+
+def formula_group_loader(
+    adata: AnnData,
+    formula=None,
+    grouping_variable=None,
+    chunk_size=int(1e4),
+    batch_size: int = None,
+):
+    device = fl.check_device()
+    if grouping_variable is None:
+        adata.obs["_copula_group"] = "shared_group"
+        grouping_variable = "_copula_group"
+        adata.obs["_copula_group"] = adata.obs["_copula_group"].astype("category")
+
+    if adata.isbacked:
+        ds = FormulaGroupViewDataset(
+            adata, formula, grouping_variable, chunk_size, device
+        )
+        dataloader = td.DataLoader(ds, batch_size=batch_size, collate_fn=stack_collate())
+        ds.x_names = fl.model_matrix_names(adata, formula, ds.categories)
+    else:
+        # convert sparse to dense matrix
+        y = adata.X
+        if isinstance(y, scipy.sparse._csc.csc_matrix):
+            y = y.todense()
+
+        # wrap the entire data into a dataset
+        x = model_matrix(formula, pd.DataFrame(adata.obs))
+        ds = td.StackDataset(
+            x=td.TensorDataset(
+                torch.tensor(np.array(x), dtype=torch.float32).to(device)
+            ),
+            y=td.TensorDataset(torch.tensor(y, dtype=torch.float32).to(device)),
+            groups=ListDataset(adata.obs[grouping_variable]),
+        )
+        ds.groups = list(adata.obs[grouping_variable].dtype.categories)
+        ds.x_names = list(x.columns)
+        dataloader = td.DataLoader(ds, batch_size=batch_size, collate_fn=stack_collate(pop=True))
+
+    return dataloader
+
+def multiple_formula_group_loader(adata: AnnData, formulas: dict, grouping_variable=None,
+                                  chunk_size=int(1e4), batch_size: int = None):
+    dataloaders = {}
+    for key in formulas.keys():
+        dataloaders[key] = formula_group_loader(adata, formulas[key], grouping_variable, chunk_size, batch_size)
+    return dataloaders 
+
+class FormulaGroupViewDataset(td.Dataset):
+    def __init__(
+        self,
+        view,
+        formula=None,
+        grouping_variable=None,
+        chunk_size=int(1e4),
+        device=None,
+    ):
+        super().__init__()
+        self.device = device or fl.check_device()
+        self.formula = formula
+        self.categories = fl.column_levels(view.obs)
+        self.grouping_variable = grouping_variable
+        self.groups = list(self.categories[grouping_variable].dtype.categories)
+        self.len = len(view)
+        self.memberships = None
+        self.view = view
+        self.x = None
+        self.y = None
+        self.cur_range = range(0, min(self.len, chunk_size))
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, ix):
+        if self.x is None or ix not in self.cur_range:
+            self.cur_range = range(ix, min(ix + len(self.cur_range), self.len))
+            view_inmem = self.view[self.cur_range].to_memory()
+            self.memberships = view_inmem.obs[self.grouping_variable]
+            self.x = fl.safe_model_matrix(
+                view_inmem.obs, self.formula, self.categories
+            ).to(self.device)
+            self.y = torch.from_numpy(view_inmem.X.toarray().astype(np.float32)).to(
+                self.device
+            )
+        return {
+            "x": self.x[ix - self.cur_range[0]],
+            "y": self.y[ix - self.cur_range[0]],
+            "groups": self.memberships[ix - self.cur_range[0]],
+        }
+
+
+class ListDataset(td.Dataset):
+    """
+    Simple DS to store groups
+    """
+
+    def __init__(self, list):
+        self.list = list
+
+    def __len__(self):
+        return len(self.list)
+
+    def __getitem__(self, idx):
+        return self.list[idx]
+
+def stack_collate(pop=False, groups=True):
+    def f(batch):
+        x = torch.stack([sample["x"][0] if pop else sample["x"] for sample in batch])
+        y = torch.stack([sample["y"][0] if pop else sample["y"] for sample in batch])
+        if groups:
+            G = tuple([sample["groups"] for sample in batch])
+            return [x, y, G]
+        return [x, y]
+    return f

scdesigner-0.0.1/src/scdesigner/data/sparse.py

@@ -0,0 +1,39 @@
+import anndata
+import torch
+import torch.utils.data as td
+
+class SparseMatrixLoader:
+    def __init__(self, adata: anndata.AnnData, batch_size: int = None):
+        ds = SparseMatrixDataset(adata, batch_size)
+        self.loader = td.DataLoader(ds, batch_size=None)
+
+
+class SparseMatrixDataset(td.IterableDataset):
+    def __init__(self, anndata: anndata.AnnData, batch_size: int = None):
+        self.n_rows = anndata.X.shape[0]
+        if batch_size is None:
+            batch_size = self.n_rows
+
+        self.sparse_matrix = anndata.X
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        for i in range(0, self.n_rows, self.batch_size):
+            batch_indices = range(i, min(i + self.batch_size, self.n_rows))
+            batch_rows = self.sparse_matrix[batch_indices, :]
+
+            # Convert to sparse CSR tensor
+            batch_indices_rows, batch_indices_cols = batch_rows.nonzero()
+            batch_values = batch_rows.data
+
+            batch_sparse_tensor = torch.sparse_coo_tensor(
+                torch.tensor([batch_indices_rows, batch_indices_cols]),
+                torch.tensor(batch_values, dtype=torch.float32),
+                (len(batch_indices), self.sparse_matrix.shape[1]),
+            ).to_sparse_csr()
+
+            yield batch_sparse_tensor
+
+    def __len__(self):
+        return (self.n_rows + self.batch_size - 1) // self.batch_size
+

scdesigner-0.0.1/src/scdesigner/diagnose/__init__.py

@@ -0,0 +1,65 @@
+from .plot import (
+    plot_umap,
+    plot_hist,
+    compare_means,
+    compare_variances,
+    compare_standard_deviation,
+    compare_umap,
+    compare_pca,
+)
+from .aic_bic import compose_marginal_diagnose, compose_gcopula_diagnose
+from .. import estimators as est
+
+__all__ = [
+    "bernoulli_gcopula_diagnose",
+    "bernoulli_regression_diagnose",
+    "compare_means",
+    "compare_pca",
+    "compare_standard_deviation",
+    "compare_umap",
+    "compare_variances",
+    "gaussian_regression_diagnose",
+    "negbin_gcopula_diagnose",
+    "negbin_regression_diagnose",
+    "plot_hist",
+    "plot_pca",
+    "plot_umap",
+    "poisson_gcopula_diagnose",
+    "poisson_regression_diagnose",
+    "zinb_gcopula_diagnose",
+    "zinb_regression_diagnose",
+    "zip_regression_diagnose"
+]
+
+
+###############################################################################
+## Methods for calculating marginal/gaussian copula AIC/BIC
+###############################################################################
+
+negbin_regression_diagnose = compose_marginal_diagnose(est.negbin.negbin_regression_likelihood,
+                                                       allowed_keys=['mean', 'dispersion'])
+negbin_gcopula_diagnose = compose_gcopula_diagnose(est.negbin.negbin_regression_likelihood,
+                                                   est.negbin.negbin_uniformizer,
+                                                   allowed_keys=['mean', 'dispersion'])
+poisson_regression_diagnose = compose_marginal_diagnose(est.poisson.poisson_regression_likelihood,
+                                                        allowed_keys=['mean'])
+poisson_gcopula_diagnose = compose_gcopula_diagnose(est.poisson.poisson_regression_likelihood,
+                                                    est.poisson.poisson_uniformizer,
+                                                    allowed_keys=['mean'])
+bernoulli_regression_diagnose = compose_marginal_diagnose(est.bernoulli.bernoulli_regression_likelihood,
+                                                          allowed_keys=['mean'])
+bernoulli_gcopula_diagnose = compose_gcopula_diagnose(est.bernoulli.bernoulli_regression_likelihood,
+                                                      est.bernoulli.bernoulli_uniformizer,
+                                                      allowed_keys=['mean'])
+zinb_regression_diagnose = compose_marginal_diagnose(est.zero_inflated_negbin.zero_inflated_negbin_regression_likelihood,
+                                                     allowed_keys=['mean', 'dispersion', 'zero_inflation'])
+zinb_gcopula_diagnose = compose_gcopula_diagnose(est.zero_inflated_negbin.zero_inflated_negbin_regression_likelihood,
+                                                 est.zero_inflated_negbin.zero_inflated_negbin_uniformizer,
+                                                 allowed_keys=['mean', 'dispersion', 'zero_inflation'])
+zip_regression_diagnose = compose_marginal_diagnose(est.zero_inflated_poisson.zero_inflated_poisson_regression_likelihood,
+                                                    allowed_keys=['mean', 'zero_inflation'])
+gaussian_regression_diagnose = compose_marginal_diagnose(est.gaussian.gaussian_regression_likelihood,
+                                                         allowed_keys=['mean', 'sdev'])
+gaussian_gcopula_diagnose = compose_gcopula_diagnose(est.gaussian.gaussian_regression_likelihood, 
+                                                     est.gaussian.gaussian_uniformizer,
+                                                     allowed_keys=['mean', 'sdev'])

scdesigner-0.0.1/src/scdesigner/diagnose/aic_bic.py

@@ -0,0 +1,119 @@
+from .. import data
+from anndata import AnnData
+from formulaic import model_matrix
+from scipy.stats import norm, multivariate_normal
+from typing import Union
+import numpy as np
+import pandas as pd
+import torch, scipy
+
+
+def marginal_aic_bic(likelihood, params: dict, adata: AnnData, 
+                     formula: Union[str, dict], allowed_keys: set,
+                     param_order: list = None, transform: list = None, 
+                     chunk_size: int = int(1e4), batch_size=512):
+    device = data.formula.check_device()
+    nsample = len(adata)
+    params = likelihood_unwrapper(params, param_order, transform).to(device)
+    nparam = len(params)
+
+    # create batches for likelihood calculation
+    formula = data.standardize_formula(formula, allowed_keys)
+    loader = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+    keys = list(loader.keys())
+    loaders = list(loader.values())
+    num_keys = len(keys)
+    
+    ll = 0
+    with torch.no_grad():
+        for batches in zip(*loaders):
+            x = {keys[i]: batches[i][0].to(device) for i in range(num_keys)}
+            y = batches[0][1].to(device)
+            ll += -likelihood(params, x, y)
+    aic = 2 * nparam - 2 * ll
+    bic = np.log(nsample) * nparam - 2 * ll
+    return aic.cpu().item(), bic.cpu().item()
+
+
+def gaussian_copula_aic_bic(uniformizer, params: dict, adata: AnnData, 
+                            formula: Union[str, dict], allowed_keys: set, copula_groups=None):
+    params = uniformizer_unwrapper(params)
+    covariance = params['covariance']
+    y = adata.X
+    if isinstance(y, scipy.sparse._csc.csc_matrix):
+        y = y.todense()
+    formula = data.standardize_formula(formula, allowed_keys)
+    X = {key: model_matrix(formula[key], pd.DataFrame(adata.obs)) for key in formula}
+    if copula_groups is not None:
+        memberships = adata.obs[copula_groups]
+    else:
+        copula_groups = "shared_group"
+        memberships = np.array(["shared_group"] * y.shape[0])
+        
+    u = uniformizer(params, X, y)
+    groups = covariance.keys()
+    nparam = {
+        g: (np.sum(covariance[g] != 0) - covariance[g].shape[0]) / 2 for g in groups
+    }
+    aic = 0 # in the future may add group-wise AIC/BIC
+    bic = 0
+    for g in groups:
+        ix = np.where(memberships == g)[0]
+        z = norm().ppf(u[ix])
+        copula_ll = multivariate_normal.logpdf(
+            z, np.zeros(covariance[g].shape[0]), covariance[g]
+        )
+        marginal_ll = norm.logpdf(z)
+        aic += -2 * (np.sum(copula_ll) - np.sum(marginal_ll)) + 2 * nparam[g]
+        bic += (
+            -2 * (np.sum(copula_ll) - np.sum(marginal_ll))
+            + np.log(z.shape[0]) * nparam[g]
+        )
+    return aic, bic
+
+
+def compose_marginal_diagnose(likelihood, allowed_keys: set, param_order: list = None, transform: list = None):
+    def diagnose(params: dict, adata: AnnData, formula: Union[str, dict],
+                chunk_size: int = int(1e4), batch_size=512):
+        return marginal_aic_bic(likelihood, params, adata, formula, allowed_keys,
+                                param_order, transform, chunk_size, batch_size)
+    return diagnose
+
+
+def compose_gcopula_diagnose(likelihood, uniformizer, allowed_keys: set, 
+                             param_order: list = None, transform: list = None):
+    def diagnose(params: dict, adata: AnnData, formula: str, copula_groups=None,
+                 chunk_size: int = int(1e4), batch_size=512):
+        marginal_aic, marginal_bic = marginal_aic_bic(likelihood, params, adata, formula, allowed_keys,
+                                                    param_order, transform, chunk_size, batch_size)
+        copula_aic, copula_bic = gaussian_copula_aic_bic(uniformizer, params, adata, formula, allowed_keys, copula_groups)
+        return marginal_aic, marginal_bic, copula_aic, copula_bic
+    return diagnose
+
+
+###############################################################################
+## Helper for converting params to match likelihood/uniformizer input format
+###############################################################################
+
+def likelihood_unwrapper(params: dict, param_order: list = None, transform: list = None):
+    l = []
+    keys_to_process = [k for k in params if k != "covariance"] if param_order is None else param_order
+
+    for idx, k in enumerate(keys_to_process):
+        feature = params[k]
+        v = torch.Tensor(feature.values).reshape(1, feature.shape[0] * feature.shape[1])[0]
+        if transform is not None:
+            v = transform[idx](v)
+        l.append(v)
+
+    return torch.cat(l, dim=0)
+
+def uniformizer_unwrapper(params):
+    params = params = {key: params[key].values if key!='covariance' else params[key] for key in params}
+    if not isinstance(params['covariance'], dict):
+        params['covariance'] = {'shared_group': params['covariance'].values}
+    else:
+        params['covariance'] = {key: params['covariance'][key].values for key in params['covariance']}
+    return params