scdesigner 0.0.1__py3-none-any.whl

scdesigner/minimal/transform.py

@@ -0,0 +1,145 @@
+from typing import Union, Sequence
+import numpy as np
+import re
+import torch
+import copy
+
+
+def nullify(sim, row_pattern: str, col_pattern: str, param: str):
+    """Nullify marginal parameters
+
+    Zero out entries matching (row_pattern, col_pattern) for the marginal `param`.
+    """
+    sim = copy.deepcopy(sim)
+    df = sim.parameters["marginal"][param]
+    matches = data_frame_mask(df, row_pattern, col_pattern)
+    mask = (~matches).astype(float)
+    mat = sim.marginal.predict.coefs[param].detach().cpu().numpy() * mask
+    _update_marginal_param(sim, param, mat)
+    return sim
+
+
+def amplify(sim, factor: float, row_pattern: str, col_pattern: str, param: str):
+    """Multiply selected marginal entries by factor."""
+    sim = copy.deepcopy(sim)
+    df = sim.parameters["marginal"][param]
+    matches = data_frame_mask(df, row_pattern, col_pattern).astype(float)
+    mask = factor * matches + np.ones_like(matches)
+    mat = sim.marginal.predict.coefs[param].detach().cpu().numpy() * mask
+    _update_marginal_param(sim, param, mat)
+    return sim
+
+
+def decorrelate(sim, row_pattern: str, col_pattern: str, group: Union[str, None] = None):
+    """Zero out selected off-diagonal entries of a covariance."""
+    sim = copy.deepcopy(sim)
+    def _apply_to_df(df):
+        m1 = data_frame_mask(df, ".", col_pattern)
+        m2 = data_frame_mask(df, row_pattern, ".")
+        mask = (m1 | m2)
+        np.fill_diagonal(mask, False)
+        df.values[mask] = 0
+
+    cov = sim.parameters["copula"]
+    _apply_to_groups(cov, group, _apply_to_df)
+    return sim
+
+
+def correlate(sim, factor: float, row_pattern: str, col_pattern: str, group: Union[str, None] = None):
+    """Multiply selected off-diagonal entries by factor."""
+    sim = copy.deepcopy(sim)
+    def _apply_to_df(df):
+        m1 = data_frame_mask(df, ".", col_pattern)
+        m2 = data_frame_mask(df, row_pattern, ".")
+        mask = (m1 | m2)
+        np.fill_diagonal(mask, False)
+        df.values[mask] = df.values[mask] * factor
+
+    cov = sim.parameters["copula"]
+    _apply_to_groups(cov, group, _apply_to_df)
+    return sim
+
+
+def replace_param(sim, path: Sequence[str], new_param):
+    """Substitute a new parameter for an old one.
+
+    Use the path to the parameter starting from sim.parameters to identify the
+    parameter to transform.  Examples: ['marginal','mean'] or
+    ['copula','group_name']
+    """
+    sim = copy.deepcopy(sim)
+    if path[0] == "marginal":
+        param = path[1]
+        mat = np.asarray(new_param)
+        _update_marginal_param(sim, param, mat)
+
+    if path[0] == "copula":
+        key = path[1]
+        cov = sim.parameters["copula"]
+        if isinstance(cov, dict):
+            cov[key] = new_param
+        else:
+            sim.parameters["copula"] = new_param
+
+    return sim
+
+
+###############################################################################
+## Helper functions used throughout
+###############################################################################
+
+def str_match(string: str, string_list: list) -> bool:
+    for l in string_list:
+        if l in string:
+            return True
+    return False
+
+
+def data_frame_mask(df, row_pattern=".", col_pattern=".") -> np.array:
+    """Return a boolean mask for a pandas DataFrame where rows/cols match regex patterns.
+
+    The returned mask has shape (n_rows, n_cols) and is True where the
+    intersection of matched rows and matched columns occurs.
+    """
+    mask = np.zeros(df.shape, dtype=bool)
+    if isinstance(col_pattern, str):
+        col_pattern = [col_pattern]
+    if isinstance(row_pattern, str):
+        row_pattern = [row_pattern]
+
+    # check for columns that match at least one pattern
+    col_matches = np.zeros(df.shape[1], dtype=bool)
+    for col in df.columns:
+        if any(re.search(pat, col) for pat in col_pattern):
+            col_matches[df.columns.get_loc(col)] = True
+
+    # same with rows
+    row_matches = np.zeros(df.shape[0], dtype=bool)
+    for idx in df.index:
+        if any(re.search(pat, str(idx)) for pat in row_pattern):
+            row_matches[df.index.get_loc(idx)] = True
+
+    # set mask to the intersection of row and column matches
+    mask = np.outer(row_matches, col_matches)
+    return mask
+
+
+def _apply_to_groups(cov_obj, grp, f):
+    """Apply f to either all group arrays or a single group's array."""
+    if isinstance(cov_obj, dict):
+        if grp is None:
+            for k in list(cov_obj.keys()):
+                f(cov_obj[k])
+        else:
+            f(cov_obj[grp])
+    else:
+        f(cov_obj)
+
+
+def _update_marginal_param(sim, param: str, mat: np.ndarray):
+    """Update the torch Parameter for a marginal `param`"""
+    tensor = sim.marginal.predict.coefs[param]
+    with torch.no_grad():
+        t = torch.from_numpy(np.asarray(mat)).to(dtype=tensor.dtype, device=tensor.device)
+        tensor.copy_(t)
+    sim.parameters["marginal"][param].values[:] = tensor.detach().cpu().numpy()

scdesigner/minimal/zero_inflated_negbin.py

@@ -0,0 +1,86 @@
+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, bernoulli
+
+class ZeroInflatedNegBin(Marginal):
+    """Zero-inflated negative-binomial marginal estimator"""
+    def __init__(self, formula: Union[Dict, str]):
+        formula = standardize_formula(formula, allowed_keys=['mean', 'dispersion', 'zero_inflation'])
+        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)")
+
+        link_funs = {
+            "mean": torch.exp,
+            "dispersion": torch.exp,
+            "zero_inflation": torch.sigmoid,
+        }
+        nll = lambda batch: -self.likelihood(batch).sum()
+        self.predict = GLMPredictor(
+            n_outcomes=self.n_outcomes,
+            feature_dims=self.feature_dims,
+            link_fns=link_funs,
+            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")
+        pi = params.get("zero_inflation")
+
+        # negative binomial component
+        negbin_loglikelihood = (
+            torch.lgamma(y + r)
+            - torch.lgamma(r)
+            - torch.lgamma(y + 1)
+            + r * torch.log(r)
+            + y * torch.log(mu)
+            - (r + y) * torch.log(r + mu)
+        )
+
+        # return the mixture, with an offset to prevent log(0)
+        return torch.log(pi * (y == 0) + (1 - pi) * torch.exp(negbin_loglikelihood) + 1e-10)
+
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
+        """Invert pseudoobservations."""
+        mu, r, pi, u = self._local_params(x, u)
+        y = nbinom(n=r, p=r / (r + mu)).ppf(u)
+        delta = bernoulli(1 - pi).ppf(u)
+        return torch.from_numpy(y * delta).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, pi, y = self._local_params(x, y)
+        nb_distn = nbinom(n=r, p=r / (r + mu))
+        u1 = pi + (1 - pi) * nb_distn.cdf(y)
+        u2 = np.where(y > 0, pi + (1 - pi) * nb_distn.cdf(y-1), 0)
+
+        # randomize within discrete mass to get uniform(0,1)
+        v = np.random.uniform(size=y.shape)
+        u = np.clip(v * u1 + (1 - v) * u2, epsilon, 1 - 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')
+        pi = params.get('zero_inflation')
+        if y is None:
+            return _to_numpy(mu, pi, r)
+        return _to_numpy(mu, r, pi, y)

scdesigner/predictors/__init__.py

@@ -0,0 +1,15 @@
+from .bernoulli import bernoulli_predict
+from .negbin import negbin_predict
+from .poisson import poisson_predict
+from .gaussian import gaussian_predict
+from .zero_inflated_negbin import zero_inflated_negbin_predict
+from .zero_inflated_poisson import zero_inflated_poisson_predict
+
+__all__ = [
+    "bernoulli_predict",
+    "gaussian_predict",
+    "negbin_predict",
+    "poisson_predict",
+    "zero_inflated_negbin_predict",
+    "zero_inflated_poisson_predict",
+]

scdesigner/predictors/bernoulli.py

@@ -0,0 +1,9 @@
+import numpy as np
+import pandas as pd
+from ..format import format_matrix
+
+
+def bernoulli_predict(parameters: dict, obs: pd.DataFrame, formula: str):
+    x = format_matrix(obs, formula)
+    theta = np.exp(x @ parameters["coef_mean"])
+    return {"mean": theta}

scdesigner/predictors/gaussian.py

@@ -0,0 +1,16 @@
+import numpy as np
+import pandas as pd
+from ..format import format_matrix
+from typing import Union
+
+def gaussian_predict(parameters: dict, obs: pd.DataFrame, formula: Union[str, dict]):
+    # Standardize formula to dictionary format
+    if isinstance(formula, str):
+        formula = {'mean': formula, 'sdev': '~ 1'}
+    
+    x_mean = format_matrix(obs, formula["mean"]) 
+    x_dispersion = format_matrix(obs, formula["sdev"])
+    
+    sigma = np.exp(x_dispersion @ parameters["coef_sdev"])
+    mu = x_mean @ parameters["coef_mean"]
+    return {"mean": mu, "sdev": sigma}

scdesigner/predictors/negbin.py

@@ -0,0 +1,17 @@
+import numpy as np
+import pandas as pd
+from ..format import format_matrix
+from typing import Union
+
+def negbin_predict(parameters: dict, obs: pd.DataFrame, formula: Union[str, dict]):
+    # Standardize formula to dictionary format
+    if isinstance(formula, str):
+        formula = {'mean': formula, 'dispersion': '~ 1'}
+    
+    x_mean = format_matrix(obs, formula["mean"]) 
+    x_dispersion = format_matrix(obs, formula["dispersion"])
+    
+    r = np.exp(x_dispersion @ parameters["coef_dispersion"])
+    mu = np.exp(x_mean @ parameters["coef_mean"])
+    # r and mu are still dataframes with column names being the gene names
+    return {"mean": mu, "dispersion": r}

scdesigner/predictors/poisson.py

@@ -0,0 +1,12 @@
+import numpy as np
+import pandas as pd
+from ..format import format_matrix
+from typing import Union
+
+def poisson_predict(parameters: dict, obs: pd.DataFrame, formula: Union[str, dict]):
+    if isinstance(formula, dict):
+        formula = formula['mean']
+    x = format_matrix(obs, formula)
+    mu = np.exp(x @ parameters["coef_mean"])
+    return {"mean": mu}
+

scdesigner/predictors/zero_inflated_negbin.py

@@ -0,0 +1,18 @@
+import numpy as np
+import pandas as pd
+from ..format import format_matrix
+from scipy.special import expit
+from typing import Union
+
+def zero_inflated_negbin_predict(parameters: dict, obs: pd.DataFrame, formula: Union[str, dict]):
+    if isinstance(formula, str):
+        formula = {"mean": formula, "dispersion": "~ 1", "zero_inflation": "~ 1"}
+    x_mean = format_matrix(obs, formula["mean"])
+    x_dispersion = format_matrix(obs, formula["dispersion"])
+    x_zero_inflation = format_matrix(obs, formula["zero_inflation"])
+    r, mu, pi = (
+        np.exp(x_dispersion @ parameters["coef_dispersion"]),
+        np.exp(x_mean @ parameters["coef_mean"]),
+        expit(x_zero_inflation @ parameters["coef_zero_inflation"]),
+    )
+    return {"mean": mu, "dispersion": r, "zero_inflation": pi}

scdesigner/predictors/zero_inflated_poisson.py

@@ -0,0 +1,18 @@
+from ..format import format_matrix
+from typing import Union
+import numpy as np
+import pandas as pd
+
+
+def zero_inflated_poisson_predict(parameters: dict, obs: pd.DataFrame, formula: Union[str, dict]):
+    if isinstance(formula, str):
+        formula = {'beta': formula, 'pi': '~ 1'}
+    mu, pi = (
+        np.exp(format_matrix(obs, formula['mean']) @ parameters["coef_mean"]),
+        sigmoid(format_matrix(obs, formula['zero_inflation']) @ parameters["coef_zero_inflation"]),
+    )
+    return {"mean": mu, "zero_inflation": pi}
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))

scdesigner/samplers/__init__.py

@@ -0,0 +1,23 @@
+from .negbin import negbin_sample, negbin_copula_sample
+from .poisson import poisson_sample, poisson_copula_sample
+from .bernoulli import bernoulli_sample, bernoulli_copula_sample
+from .gaussian import gaussian_regression_sample, gaussian_copula_sample
+from .zero_inflated_negbin import (
+    zero_inflated_negbin_sample,
+    zero_inflated_negbin_copula_sample,
+)
+from .zero_inflated_poisson import zero_inflated_poisson_sample
+
+__all__ = [
+    "negbin_sample",
+    "negbin_copula_sample",
+    "poisson_sample",
+    "poisson_copula_sample",
+    "bernoulli_sample",
+    "bernoulli_copula_sample",
+    "gaussian_regression_sample",
+    "gaussian_copula_sample",
+    "zero_inflated_negbin_sample",
+    "zero_inflated_negbin_copula_sample",
+    "zero_inflated_poisson_sample",
+]

scdesigner/samplers/bernoulli.py

@@ -0,0 +1,27 @@
+from . import glm_factory as glm
+from scipy.stats import bernoulli
+from typing import Union
+import numpy as np
+
+
+def bernoulli_regression_sample_array(local_parameters: dict) -> np.array:
+    theta = local_parameters["mean"]
+    return bernoulli(theta).rvs()
+
+
+def bernoulli_copula_sample_array(
+    local_parameters: dict, covariance: Union[dict, np.array], groups: dict
+) -> np.array:
+    # initialize uniformized gaussian samples
+    N, G = local_parameters["mean"].shape
+    u = glm.gaussian_copula_pseudo_obs(N, G, covariance, groups)
+
+    theta = local_parameters["mean"]
+    return bernoulli(theta).ppf(u)
+
+
+bernoulli_sample = glm.glm_sample_factory(bernoulli_regression_sample_array)
+
+bernoulli_copula_sample = glm.gaussian_copula_sample_factory(
+    bernoulli_copula_sample_array
+)

scdesigner/samplers/gaussian.py

@@ -0,0 +1,25 @@
+from scipy.stats import norm
+from . import glm_factory as glm
+from typing import Union
+import numpy as np
+
+
+def gaussian_regression_sample_array(local_parameters: dict) -> np.array:
+    sigma, mu = local_parameters["sdev"], local_parameters["mean"] # dataframes of shape (n, g)
+    return norm(loc=mu, scale=sigma).rvs()
+
+
+def gaussian_copula_sample_array(
+    local_parameters: dict, covariance: Union[dict, np.array], groups: dict
+) -> np.array:
+    # initialize uniformized gaussian samples
+    N, G = local_parameters["mean"].shape
+    u = glm.gaussian_copula_pseudo_obs(N, G, covariance, groups)
+
+    # transform the correlated uniforms to NB space
+    sigma, mu = local_parameters["sdev"], local_parameters["mean"]
+    return norm(loc=mu, scale=sigma).ppf(u)
+
+
+gaussian_regression_sample = glm.glm_sample_factory(gaussian_regression_sample_array)
+gaussian_copula_sample = glm.gaussian_copula_sample_factory(gaussian_copula_sample_array)

scdesigner/samplers/glm_factory.py

@@ -0,0 +1,41 @@
+import numpy as np
+import pandas as pd
+import anndata as ad
+from typing import Union
+from scipy.stats import norm
+
+
+def glm_sample_factory(sample_array):
+    def sampler(local_parameters: dict, obs: pd.DataFrame) -> ad.AnnData:
+        samples = sample_array(local_parameters)
+        result = ad.AnnData(X=samples, obs=obs)
+        result.var_names = local_parameters["mean"].columns
+        return result
+    return sampler
+
+def gaussian_copula_pseudo_obs(N, G, sigma, groups):
+    u = np.zeros((N, G))
+
+    # cycle across groups
+    for group, ix in groups.items():
+        if type(sigma) is not dict:
+            sigma = {group: sigma}
+
+        z = np.random.multivariate_normal(
+            mean=np.zeros(G), cov=sigma[group], size=len(ix)
+        )
+        normal_distn = norm(0, np.diag(sigma[group] ** 0.5))
+        u[ix] = normal_distn.cdf(z)
+    return u
+
+
+def gaussian_copula_sample_factory(copula_sample_array):
+    def sampler(
+        local_parameters: dict, covariance: Union[dict, np.array], groups: dict, obs: pd.DataFrame
+    ) -> ad.AnnData:
+        samples = copula_sample_array(local_parameters, covariance, groups)
+        result = ad.AnnData(X=samples, obs=obs)
+        result.var_names = local_parameters["mean"].columns
+        return result
+    return sampler
+

scdesigner/samplers/negbin.py

@@ -0,0 +1,25 @@
+from scipy.stats import nbinom
+from . import glm_factory as glm
+from typing import Union
+import numpy as np
+
+
+def negbin_regression_sample_array(local_parameters: dict) -> np.array:
+    r, mu = local_parameters["dispersion"], local_parameters["mean"] # dataframes of shape (n, g)
+    return nbinom(n=r, p=r / (r + mu)).rvs()
+
+
+def negbin_copula_sample_array(
+    local_parameters: dict, covariance: Union[dict, np.array], groups: dict
+) -> np.array:
+    # initialize uniformized gaussian samples
+    N, G = local_parameters["mean"].shape
+    u = glm.gaussian_copula_pseudo_obs(N, G, covariance, groups)
+
+    # transform the correlated uniforms to NB space
+    r, mu = local_parameters["dispersion"], local_parameters["mean"]
+    return nbinom(n=r, p=r / (r + mu)).ppf(u)
+
+
+negbin_sample = glm.glm_sample_factory(negbin_regression_sample_array)
+negbin_copula_sample = glm.gaussian_copula_sample_factory(negbin_copula_sample_array)

scdesigner/samplers/poisson.py

@@ -0,0 +1,25 @@
+from . import glm_factory as glm
+from scipy.stats import poisson
+from typing import Union
+import numpy as np
+
+
+def poisson_regression_sample_array(local_parameters: dict) -> np.array:
+    mu = local_parameters["mean"]
+    return poisson(mu).rvs()
+
+
+def poisson_copula_sample_array(
+    local_parameters: dict, covariance: Union[dict, np.array], groups: dict
+) -> np.array:
+    # initialize uniformized gaussian samples
+    N, G = local_parameters["mean"].shape
+    u = glm.gaussian_copula_pseudo_obs(N, G, covariance, groups)
+
+    # invert using poisson margins
+    mu = local_parameters["mean"]
+    return poisson(mu).ppf(u)
+
+
+poisson_sample = glm.glm_sample_factory(poisson_regression_sample_array)
+poisson_copula_sample = glm.gaussian_copula_sample_factory(poisson_copula_sample_array)

scdesigner/samplers/zero_inflated_negbin.py

@@ -0,0 +1,40 @@
+from scipy.stats import nbinom, bernoulli
+from . import glm_factory as glm
+from typing import Union
+import numpy as np
+
+
+def zero_inflated_negbin_sample_array(local_parameters: dict) -> np.array:
+    mu, r, pi = (
+        local_parameters["mean"],
+        local_parameters["dispersion"],
+        local_parameters["zero_inflation"],
+    )
+    return nbinom(n=r, p=r / (r + mu)).rvs() * bernoulli(1 - pi).rvs()
+
+
+def zero_inflated_negbin_copula_sample_array(
+    local_parameters: dict, covariance: Union[dict, np.array], groups: dict
+) -> np.array:
+    # initialize uniformized gaussian samples
+    N, G = local_parameters["mean"].shape
+    u = glm.gaussian_copula_pseudo_obs(N, G, covariance, groups)
+
+    # get zero inflated NB parameters
+    mu, r, pi = (
+        local_parameters["mean"],
+        local_parameters["dispersion"],
+        local_parameters["zero_inflation"],
+    )
+
+    # zero inflate after first simulating from NB
+    positive_part = nbinom(n=r, p=r / (r + mu)).ppf(u)
+    zero_inflation = bernoulli(1 - pi).ppf(u)
+    return zero_inflation * positive_part
+
+
+zero_inflated_negbin_sample = glm.glm_sample_factory(zero_inflated_negbin_sample_array)
+
+zero_inflated_negbin_copula_sample = glm.gaussian_copula_sample_factory(
+    zero_inflated_negbin_copula_sample_array
+)

scdesigner/samplers/zero_inflated_poisson.py

@@ -0,0 +1,16 @@
+from scipy.stats import poisson, bernoulli
+from . import glm_factory as glm
+import numpy as np
+
+
+def zero_inflated_poisson_sample_array(local_parameters: dict) -> np.array:
+    mu, pi = (
+        local_parameters["mean"],
+        local_parameters["zero_inflation"],
+    )
+    return poisson(mu).rvs() * bernoulli(1 - pi).rvs()
+
+
+zero_inflated_poisson_sample = glm.glm_sample_factory(
+    zero_inflated_poisson_sample_array
+)

scdesigner/simulators/__init__.py

@@ -0,0 +1,31 @@
+from .composite_regressor import CompositeGLMSimulator
+from .glm_simulator import (
+    BernoulliCopulaSimulator,
+    BernoulliRegressionSimulator,
+    NegBinCopulaSimulator,
+    NegBinRegressionSimulator,
+    PoissonCopulaSimulator,
+    PoissonRegressionSimulator,
+    GaussianRegressionSimulator,
+    GaussianCopulaSimulator,
+    ZeroInflatedNegBinCopulaSimulator,
+    ZeroInflatedNegBinRegressionSimulator,
+    ZeroInflatedPoissonRegressionSimulator,
+)
+from .pnmf_regression import PNMFRegressionSimulator
+
+__all__ = [
+    "BernoulliCopulaSimulator",
+    "BernoulliRegressionSimulator",
+    "CompositeGLMSimulator",
+    "GaussianRegressionSimulator",
+    "GaussianCopulaSimulator",
+    "NegBinCopulaSimulator",
+    "NegBinRegressionSimulator",
+    "PNMFRegressionSimulator",
+    "PoissonCopulaSimulator",
+    "PoissonRegressionSimulator",
+    "ZeroInflatedNegBinCopulaSimulator",
+    "ZeroInflatedNegBinRegressionSimulator",
+    "ZeroInflatedPoissonRegressionSimulator",
+]