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

scdesigner/distributions/__init__.py

@@ -0,0 +1,19 @@
+"""Marginal distribution implementations."""
+
+from .negbin import NegBin
+from .negbin_irls import NegBinIRLS
+from .zero_inflated_negbin import ZeroInflatedNegBin
+from .gaussian import Gaussian
+from .bernoulli import Bernoulli
+from .poisson import Poisson
+from .zero_inflated_poisson import ZeroInflatedPoisson
+
+__all__ = [
+    "NegBin",
+    "NegBinIRLS",
+    "ZeroInflatedNegBin",
+    "Gaussian",
+    "Bernoulli",
+    "Poisson",
+    "ZeroInflatedPoisson",
+]

scdesigner/{minimal → distributions}/bernoulli.py

@@ -1,15 +1,42 @@
-from .formula import standardize_formula
-from .marginal import GLMPredictor, Marginal
-from .loader import _to_numpy
-from typing import Union, Dict, Optional
+from ..data.formula import standardize_formula
+from ..base.marginal import GLMPredictor, Marginal
+from ..data.loader import _to_numpy
+from typing import Union, Dict, Optional, Tuple
 import torch
 import numpy as np
-from scipy.stats import nbinom, bernoulli
+from scipy.stats import bernoulli
 
-class ZeroInflatedNegBin(Marginal):
-    """Zero-inflated negative-binomial marginal estimator"""
+class Bernoulli(Marginal):
+    """Bernoulli marginal estimator
+
+    This subclass behaves like `Marginal` but assumes each feature follows a
+    Bernoulli distribution with success probability `theta_j(x)` that depends
+    on covariates `x` through the `formula` argument.
+
+    The allowed formula keys are 'mean' (interpreted as the logit of the
+    success probability when used with a GLM link). If a string formula is
+    provided, it is taken to specify the `mean` model.
+
+    Examples
+    --------
+    >>> from scdesigner.distributions import Bernoulli
+    >>> from scdesigner.datasets import pancreas
+    >>>
+    >>> sim = Bernoulli(formula="~ pseudotime")
+    >>> sim.setup_data(pancreas)
+    >>> sim.fit(max_epochs=1, verbose=False)
+    >>>
+    >>> # evaluate p(y | x) and theta(x)
+    >>> y, x = next(iter(sim.loader))
+    >>> l = sim.likelihood((y, x))
+    >>> y_hat = sim.predict(x)
+    >>>
+    >>> # convert to quantiles and back
+    >>> u = sim.uniformize(y, x)
+    >>> x_star = sim.invert(u, x)
+    """
     def __init__(self, formula: Union[Dict, str]):
-        formula = standardize_formula(formula, allowed_keys=['mean', 'dispersion', 'zero_inflation'])
+        formula = standardize_formula(formula, allowed_keys=['mean'])
         super().__init__(formula)
 
     def setup_optimizer(
@@ -21,7 +48,8 @@ class ZeroInflatedNegBin(Marginal):
             raise RuntimeError("self.loader is not set (call setup_data first)")
 
         link_fns = {"mean": torch.sigmoid}
-        nll = lambda batch: -self.likelihood(batch).sum()
+        def nll(batch):
+            return -self.likelihood(batch).sum()
         self.predict = GLMPredictor(
             n_outcomes=self.n_outcomes,
             feature_dims=self.feature_dims,
@@ -31,20 +59,20 @@ class ZeroInflatedNegBin(Marginal):
             optimizer_kwargs=optimizer_kwargs
         )
 
-    def likelihood(self, batch):
-        """Compute the negative log-likelihood"""
+    def likelihood(self, batch) -> torch.Tensor:
+        """Compute the log-likelihood"""
         y, x = batch
         params = self.predict(x)
         theta = params.get("mean")
         return y * torch.log(theta) + (1 - y) * torch.log(1 - theta)
 
-    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]) -> torch.Tensor:
         """Invert pseudoobservations."""
         theta, u = self._local_params(x, u)
         y = bernoulli(theta).ppf(u)
         return torch.from_numpy(y).float()
 
-    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor], epsilon=1e-6):
+    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor], epsilon=1e-6) -> torch.Tensor:
         """Return uniformized pseudo-observations for counts y given covariates x."""
         theta, y = self._local_params(x, y)
         u1 =  bernoulli(theta).cdf(y)
@@ -53,7 +81,7 @@ class ZeroInflatedNegBin(Marginal):
         u = np.clip(v * u1 + (1 - v) * u2, epsilon, 1 - epsilon)
         return torch.from_numpy(u).float()
 
-    def _local_params(self, x, y=None):
+    def _local_params(self, x, y=None) -> Tuple:
         params = self.predict(x)
         theta = params.get('mean')
         if y is None:

scdesigner/distributions/gaussian.py

@@ -0,0 +1,114 @@
+from ..data.formula import standardize_formula
+from ..base.marginal import GLMPredictor, Marginal
+from ..data.loader import _to_numpy
+from typing import Union, Dict, Optional, Tuple
+import torch
+import numpy as np
+from scipy.stats import norm
+
+class Gaussian(Marginal):
+    """Gaussian marginal estimator
+
+    This subclass behaves like `Marginal` but assuming that each gene follows a
+    normal N(mu[j](x), sigma[j]^2(x)) distribution. The parameters mu[j](x) and
+    sigma[j]^2(x) depend on experimental or biological features x through the
+    formula object.
+
+    The allowed formula keys are 'mean' and 'sdev', defaulting to 'mean' with a
+    fixed standard deviation if only a string formula is passed in.
+
+    Examples
+    --------
+    >>> from scdesigner.distributions import Gaussian
+    >>> from scdesigner.datasets import pancreas
+    >>>
+    >>> sim = Gaussian(formula={"mean": "~ bs(pseudotime, df=5)", "sdev": "~ pseudotime"})
+    >>> sim.setup_data(pancreas)
+    >>> sim.fit(max_epochs=1, verbose=False)
+    >>>
+    >>> # evaluate p(y | x) and mu(x)
+    >>> y, x = next(iter(sim.loader))
+    >>> l = sim.likelihood((y, x))
+    >>> y_hat = sim.predict(x)
+    >>>
+    >>> # convert to quantiles and back
+    >>> u = sim.uniformize(y, x)
+    >>> x_star = sim.invert(u, x)
+    """
+    def __init__(self, formula: Union[Dict, str]):
+        formula = standardize_formula(formula, allowed_keys=['mean', 'sdev'])
+        super().__init__(formula)
+
+    def setup_optimizer(
+            self,
+            optimizer_class: Optional[callable] = torch.optim.Adam,
+            **optimizer_kwargs,
+    ):
+        """
+        Gaussian Model Optimizer
+
+        By default optimization is done using Adam. This can be customized using
+        the `optimizer_class` argument. The link function for the mean is an
+        identity link.
+
+        Parameters
+        ----------
+        optimizer_class : Optional[callable]
+           We optimize the negative log likelihood using the Adam optimizer by
+           default. Alternative torch.optim.* optimizer can be passed in
+           through this argument.
+        **optimizer_kwargs :
+            Arguments that are passed to the optimizer during estimation.
+
+        Returns
+        -------
+            Does not return anything, but modifies the self.predict attribute to
+            refer to the new optimizer object.
+        """
+        if self.loader is None:
+            raise RuntimeError("self.loader is not set (call setup_data first)")
+
+        def nll(batch):
+            return -self.likelihood(batch).sum()
+        link_fns = {"mean": lambda x: x}
+        self.predict = GLMPredictor(
+            n_outcomes=self.n_outcomes,
+            feature_dims=self.feature_dims,
+            link_fns=link_fns,
+            loss_fn=nll,
+            optimizer_class=optimizer_class,
+            optimizer_kwargs=optimizer_kwargs
+        )
+
+    def likelihood(self, batch) -> torch.Tensor:
+        """Compute the log-likelihood"""
+        y, x = batch
+        params = self.predict(x)
+        mu = params.get("mean")
+        sigma = params.get("sdev")
+
+        # log likelihood for Gaussian
+        log_likelihood = -0.5 * (torch.log(2 * torch.pi * sigma ** 2) + ((y - mu) ** 2) / (sigma ** 2))
+        return log_likelihood
+
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """Invert pseudoobservations."""
+        mu, sdev, u = self._local_params(x, u)
+        y = norm(loc=mu, scale=sdev).ppf(u)
+        return torch.from_numpy(y).float()
+
+    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor], epsilon=1e-6) -> torch.Tensor:
+        """Return uniformized pseudo-observations for counts y given covariates x."""
+        # cdf values using scipy's parameterization
+        mu, sdev, y = self._local_params(x, y)
+        u = norm.cdf(y, loc=mu, scale=sdev)
+        u = np.clip(u, epsilon, 1 - epsilon)
+        return torch.from_numpy(u).float()
+
+    def _local_params(self, x, y=None) -> Tuple:
+        params = self.predict(x)
+        mu = params.get('mean')
+        sdev = params.get('sdev')
+        if y is None:
+            return _to_numpy(mu, sdev)
+        return _to_numpy(mu, sdev, y)

scdesigner/distributions/negbin.py

@@ -0,0 +1,121 @@
+from ..base.marginal import GLMPredictor, Marginal
+from ..data.formula import standardize_formula
+from ..data.loader import _to_numpy
+from ..utils.kwargs import _filter_kwargs, DEFAULT_ALLOWED_KWARGS
+from .negbin_irls_funs import initialize_parameters
+from scipy.stats import nbinom
+from typing import Union, Dict, Optional, Tuple
+import numpy as np
+import torch
+
+
+class NegBin(Marginal):
+    """Negative-binomial marginal estimator with poisson initialization
+
+    This subclass behaves like `Marginal` but assumes each gene follows a
+    negative binomial distribution NB(mu_j(x), r_j(x)) parameterized via a mean
+    `mu_j(x)` and dispersion `r_j(x)` that depend on covariates `x` through the
+    provided `formula` object.
+
+    The allowed formula keys are 'mean' and 'dispersion', defaulting to
+    'mean' with a fixed dispersion if only a string formula is passed in.
+
+    Examples
+    --------
+    >>> from scdesigner.distributions import NegBin
+    >>> from scdesigner.datasets import pancreas
+    >>>
+    >>> sim = NegBin(formula={"mean": "~ bs(pseudotime, df=5)", "dispersion": "~ pseudotime"})
+    >>> sim.setup_data(pancreas)
+    >>> sim.fit(max_epochs=1, verbose=False)
+    >>>
+    >>> # evaluate p(y | x) and mu(x)
+    >>> y, x = next(iter(sim.loader))
+    >>> l = sim.likelihood((y, x))
+    >>> y_hat = sim.predict(x)
+    >>>
+    >>> # convert to quantiles and back
+    >>> u = sim.uniformize(y, x)
+    >>> x_star = sim.invert(u, x)
+    """
+    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.AdamW,
+            **optimizer_kwargs,
+    ):
+        if self.loader is None:
+            raise RuntimeError("self.loader is not set (call setup_data first)")
+
+        def nll(batch):
+            return -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) -> torch.Tensor:
+        """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]) -> 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) -> torch.Tensor:
+        """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) -> Tuple:
+        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)
+
+    def fit(self, max_epochs: int = 100, verbose: bool = True, **kwargs):
+        if self.predict is None:
+                self.setup_optimizer(**kwargs)
+
+        # initialize using a poisson fit
+        initialize_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['initialize'])
+        beta_init, gamma_init = initialize_parameters(
+            self.loader, self.n_outcomes, self.feature_dims['mean'],
+            self.feature_dims['dispersion'],
+            **initialize_kwargs
+        )
+        with torch.no_grad():
+            self.predict.coefs['mean'].copy_(beta_init)
+            self.predict.coefs['dispersion'].copy_(gamma_init)
+
+        return Marginal.fit(self, max_epochs, verbose, **kwargs)

scdesigner/distributions/negbin_irls.py

@@ -0,0 +1,72 @@
+import torch
+from .negbin import NegBin
+from .negbin_irls_funs import initialize_parameters, step_stochastic_irls
+from ..data.formula import standardize_formula
+from ..utils.kwargs import _filter_kwargs, DEFAULT_ALLOWED_KWARGS
+from typing import Union, Dict
+
+
+class NegBinIRLS(NegBin):
+    """
+    Negative-Binomial Marginal using Stochastic IRLS with
+    active response tracking and log-likelihood convergence.
+    """
+    def __init__(self, formula: Union[Dict, str]):
+        formula = standardize_formula(formula, allowed_keys=['mean', 'dispersion'])
+        super().__init__(formula, device="cpu")
+
+
+    def fit(self, max_epochs=10, tol=1e-4, eta=0.1, verbose=True, **kwargs):
+        if self.predict is None:
+                self.setup_optimizer(**kwargs)
+
+        # 1. Initialization using poisson fit
+        initialize_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['initialize'])
+        beta_init, gamma_init = initialize_parameters(
+            self.loader, self.n_outcomes, self.feature_dims['mean'],
+            self.feature_dims['dispersion'],
+            **initialize_kwargs
+        )
+
+        with torch.no_grad():
+            self.predict.coefs['mean'].copy_(beta_init)
+            self.predict.coefs['dispersion'].copy_(gamma_init)
+
+        # 2. All genes are active at the start
+        active_mask = torch.ones(self.n_outcomes, dtype=torch.bool)
+        ll_ = - 1e9 * torch.ones(self.n_outcomes, dtype=torch.float32)
+
+        for epoch in range(max_epochs):
+            if not active_mask.any(): break
+            ll, n_batches = 0.0, 0
+
+            with torch.no_grad():
+                for y_batch, x_dict in self.loader:
+
+                    # Slice active genes
+                    idx = torch.where(active_mask)[0]
+                    y_act = y_batch[:, active_mask]
+                    X = x_dict['mean']
+                    Z = x_dict['dispersion']
+
+                    # Fetch current coefficients and update
+                    b_curr = self.predict.coefs['mean'][:, active_mask]
+                    g_curr = self.predict.coefs['dispersion'][:, active_mask]
+                    b_next, g_next, conv_mask, ll_cur = step_stochastic_irls(y_act, X, Z, b_curr, g_curr, eta, tol, ll_[active_mask])
+                    ll_[active_mask] = ll_cur
+
+                    # Update Parameters and de-activate converged genes
+                    with torch.no_grad():
+                        self.predict.coefs['mean'][:, active_mask] = b_next
+                        self.predict.coefs['dispersion'][:, active_mask] = g_next
+                        active_mask[idx[conv_mask]] = False
+
+                    # Accumulate batch log-likelihood using `ll` from the IRLS step
+                    ll += ll_.sum().item()
+                    n_batches += 1
+
+                if verbose and ((epoch + 1) % 10) == 0:
+                    print(f"Epoch {epoch+1}/{max_epochs} | Genes remaining: {active_mask.sum().item()} | Loss: {-ll / n_batches:.4f}", end='\r')
+                    if not active_mask.any(): break
+
+        self.parameters = self.format_parameters()