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

scdesigner/distributions/negbin_irls_funs.py

@@ -0,0 +1,456 @@
+import torch
+from typing import Optional
+import torch.special as spec
+
+# ==============================================================================
+# Weighted Least Squares Solver
+# ==============================================================================
+
+def solve_weighted_least_squares(X, weights, responses):
+    """
+    Solve multiple independent weighted least squares problems in parallel.
+
+    For each column j, solves: (X'W_j X)β_j = X'W_j z_j
+    where W_j is a diagonal matrix with weights[:, j] on the diagonal.
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        Design matrix (n × p)
+    weights : torch.Tensor
+        Weight matrix (n × m), one weight vector per response
+    responses : torch.Tensor
+        Working responses (n × m)
+
+    Returns
+    -------
+    torch.Tensor
+        Coefficient matrix (p × m)
+    """
+    # Precompute outer products X_i ⊗ X_i for each observation
+    X_outer = torch.einsum("ni,nj->nij", X, X)  # (n × p × p)
+
+    # Compute weighted normal equations: (X'WX) for all m responses at once
+    eye = torch.eye(X.shape[1]).unsqueeze(0)
+    weighted_XX = torch.einsum("nm,nij->mij", weights, X_outer)  # (m × p × p)
+    weighted_XX = weighted_XX + 1e-5 * eye
+
+    # Compute X'Wz for all responses
+    weighted_Xy = torch.einsum("ni,nm->mi", X, weights * responses)  # (m × p)
+
+    # Solve all systems at once
+    coefficients = torch.linalg.solve(weighted_XX, weighted_Xy.unsqueeze(-1))
+    return coefficients.squeeze(-1).T  # (p × m)
+
+
+# ==============================================================================
+# Mean Parameter Updates (Beta)
+# ==============================================================================
+
+def update_mean_coefficients(X, counts, beta, dispersion, clip: float = 5.0):
+    """
+    Update mean model coefficients using one Newton-Raphson step.
+
+    Uses IRLS (Iteratively Reweighted Least Squares) with:
+    - Working weights: W = μ/(1 + μ/θ)
+    - Working response: Z = Xβ + (Y - μ)/μ
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        Design matrix (n × p)
+    counts : torch.Tensor
+        Observed counts (n × m)
+    beta : torch.Tensor
+        Current coefficients (p × m)
+    dispersion : torch.Tensor
+        Current dispersion parameters (n × m)
+    clip : float, optional
+        Maximum absolute value for linear predictor, by default 5.0
+
+    Returns
+    -------
+    torch.Tensor
+        Updated coefficients (p × m)
+    """
+    linear_pred = torch.clip(X @ beta, min=-clip, max=clip)
+    mean = torch.exp(linear_pred)
+    weights = mean / (1 + mean / dispersion)
+    working_response = linear_pred + (counts - mean) / mean
+    return solve_weighted_least_squares(X, weights, working_response)
+
+
+# ==============================================================================
+# Dispersion Parameter Updates (Gamma)
+# ==============================================================================
+
+def update_dispersion_coefficients(Z, counts, mean, gamma, clip: float = 5.0):
+    """
+    Update dispersion model coefficients using one Fisher scoring step.
+
+    Uses working response U = η + θ·s/w where:
+    - η = Zγ (linear predictor)
+    - s = score with respect to θ
+    - w = approximate Fisher information
+
+    Parameters
+    ----------
+    Z : torch.Tensor
+        Dispersion design matrix (n × q)
+    counts : torch.Tensor
+        Observed counts (n × m)
+    mean : torch.Tensor
+        Current mean estimates (n × m)
+    gamma : torch.Tensor
+        Current dispersion coefficients (q × m)
+    clip : float, optional
+        Maximum absolute value for linear predictor, by default 5.0
+
+    Returns
+    -------
+    torch.Tensor
+        Updated dispersion coefficients (q × m)
+    """
+    linear_pred = torch.clip(Z @ gamma, min=-clip, max=clip)
+    dispersion = torch.exp(linear_pred)
+
+    # Score: ∂ℓ/∂θ
+    psi_diff = spec.digamma(counts + dispersion) - spec.digamma(dispersion)
+    score = (psi_diff + torch.log(dispersion) - torch.log(mean + dispersion) +
+             (mean - counts) / (mean + dispersion))
+
+    # Approximate Fisher information (replaces exact Hessian)
+    # Approximation: θY/(θ + Y) ≈ θ²[ψ₁(θ) - ψ₁(Y + θ)]
+    weights = ((dispersion * counts) / (dispersion + counts)).clip(min=1e-6)
+    working_response = linear_pred + (dispersion * score) / weights
+    return solve_weighted_least_squares(Z, weights, working_response)
+
+
+# ==============================================================================
+# Initialization
+# ==============================================================================
+
+def estimate_constant_dispersion(X, counts, beta):
+    """
+    Estimate constant dispersion for each response using method of moments.
+
+    Uses Pearson residuals: θ̂ = (Σμ) / max(χ² - df, 0.1)
+    where χ² = Σ(Y - μ)²/μ and df = n - p.
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        Design matrix (n × p)
+    counts : torch.Tensor
+        Observed counts (n × m)
+    beta : torch.Tensor
+        Mean coefficients (p × m)
+
+    Returns
+    -------
+    torch.Tensor
+        Dispersion estimates (m,)
+    """
+    mean = torch.exp(X @ beta)
+    pearson_chi2 = torch.sum((counts - mean)**2 / mean, dim=0)
+    sum_mean = torch.sum(mean, dim=0)
+
+    degrees_freedom = counts.shape[0] - X.shape[1]
+    dispersion = sum_mean / torch.clip(pearson_chi2 - degrees_freedom, min=0.1)
+    return torch.clip(dispersion, min=0.1)
+
+
+def fit_poisson_initial(X, counts, tol: float = 1e-3, max_iter: int = 100, clip: float = 5.0):
+    """
+    Fit Poisson GLM to initialize mean parameters.
+
+    Parameters
+    ----------
+    X : torch.Tensor
+        Design matrix (n × p)
+    counts : torch.Tensor
+        Observed counts (n × m)
+    tol : float, optional
+        Convergence tolerance, by default 1e-3
+    max_iter : int, optional
+        Maximum iterations, by default 100
+    clip : float, optional
+        Maximum absolute value for linear predictor, by default 5.0
+
+    Returns
+    -------
+    torch.Tensor
+        Initial coefficients (p × m)
+    """
+    n_features, n_responses = X.shape[1], counts.shape[1]
+    beta = torch.zeros((n_features, n_responses))
+
+    for _ in range(max_iter):
+        beta_old = beta.clone()
+        linear_pred = torch.clip(X @ beta, min=-clip, max=clip)
+        mean = torch.exp(linear_pred)
+        working_response = linear_pred + (counts - mean) / mean
+
+        beta = solve_weighted_least_squares(X, mean, working_response)
+        if torch.max(torch.abs(beta - beta_old)) < tol:
+            break
+
+    return beta
+
+
+def accumulate_poisson_statistics(loader, beta, n_genes, p_mean, clip = 5):
+    """
+    Accumulate weighted normal equations for Poisson IRLS across batches.
+
+    Parameters
+    ----------
+    loader : DataLoader
+        DataLoader yielding (y_batch, x_dict)
+    beta : torch.Tensor
+        Current coefficients (p_mean × n_genes)
+    n_genes : int
+        Number of genes
+    p_mean : int
+        Number of mean predictors
+    clip : float, optional
+        Maximum absolute value for linear predictor, by default 5
+
+    Returns
+    -------
+    weighted_XX : torch.Tensor
+        Accumulated X'WX (n_genes × p_mean × p_mean)
+    weighted_Xy : torch.Tensor
+        Accumulated X'Wz (p_mean × n_genes)
+    """
+    weighted_XX = torch.zeros((n_genes, p_mean, p_mean))
+    weighted_Xy = torch.zeros((p_mean, n_genes))
+
+    for y_batch, x_dict in loader:
+        X = x_dict['mean'].to("cpu")
+
+        linear_pred = torch.clip(X @ beta, min=-clip, max=clip)
+        mean = torch.exp(linear_pred)
+        working_response = linear_pred + (y_batch.to("cpu") - mean) / mean
+
+        X_outer = torch.einsum("ni,nj->nij", X, X)
+        weighted_XX += torch.einsum("nm,nij->mij", mean, X_outer)
+        weighted_Xy += torch.einsum("ni,nm->im", X, mean * working_response)
+
+    return weighted_XX, weighted_Xy
+
+
+def accumulate_dispersion_statistics(loader, beta, clip = 5):
+    """
+    Accumulate Pearson statistics for method of moments dispersion estimation.
+
+    Parameters
+    ----------
+    loader : DataLoader
+        DataLoader yielding (y_batch, x_dict)
+    beta : torch.Tensor
+        Mean coefficients (p_mean × n_genes)
+    clip : float, optional
+        Maximum absolute value for linear predictor, by default 5
+
+    Returns
+    -------
+    sum_mean : torch.Tensor
+        Total predicted mean (n_genes,)
+    sum_pearson : torch.Tensor
+        Total Pearson chi-squared (n_genes,)
+    n_total : int
+        Total number of observations
+    """
+    sum_mean = torch.zeros(beta.shape[1])
+    sum_pearson = torch.zeros(beta.shape[1])
+    n_total = 0
+
+    for y_batch, x_dict in loader:
+        X = x_dict['mean'].to('cpu')
+        linear_pred = torch.clip(X @ beta, min=-clip, max=clip)
+        mean_batch = torch.exp(linear_pred)
+
+        sum_mean += mean_batch.sum(dim=0)
+        sum_pearson += ((y_batch.to('cpu') - mean_batch)**2 / mean_batch).sum(dim=0)
+        n_total += y_batch.shape[0]
+
+    return sum_mean, sum_pearson, n_total
+
+
+def initialize_parameters(loader, n_genes, p_mean, p_disp, max_iter = 10,
+                          tol = 1e-3, clip = 5):
+    """
+    Initialize parameters using batched Poisson IRLS followed by MoM dispersion.
+
+    Logic:
+        1. Iteratively fit Poisson GLM by accumulating X'WX and X'WZ across batches
+        2. Use fitted Poisson means to estimate dispersion via Method of Moments
+
+    Parameters
+    ----------
+    loader : DataLoader
+        DataLoader yielding (y_batch, x_dict)
+    n_genes : int
+        Number of response columns (genes)
+    p_mean : int
+        Number of predictors in the mean model
+    p_disp : int
+        Number of predictors in the dispersion model
+    max_iter : int, optional
+        Maximum Poisson IRLS iterations, by default 10
+    tol : float, optional
+        Convergence tolerance for beta coefficients, by default 1e-3
+    clip : float, optional
+        Maximum absolute value for linear predictor, by default 10
+
+    Returns
+    -------
+    beta_init : torch.Tensor
+        (p_mean × n_genes) tensor
+    gamma_init : torch.Tensor
+        (p_disp × n_genes) tensor
+    """
+    beta = torch.zeros((p_mean, n_genes))
+    for _ in range(max_iter):
+        weighted_XX, weighted_Xy = accumulate_poisson_statistics(
+            loader, beta, n_genes, p_mean, clip
+        )
+
+        eye = torch.eye(p_mean).unsqueeze(0)
+        weighted_XX_reg = weighted_XX + 1e-6 * eye
+        beta_new = torch.linalg.solve(
+            weighted_XX_reg, weighted_Xy.T.unsqueeze(-1)
+        ).squeeze(-1).T
+
+        if torch.max(torch.abs(beta_new - beta)) < tol:
+            beta = beta_new
+            break
+        beta = beta_new
+
+    sum_mean, sum_pearson, n_total = accumulate_dispersion_statistics(
+        loader, beta, clip
+    )
+
+    degrees_freedom = n_total - p_mean
+    dispersion = sum_mean / torch.clip(sum_pearson - degrees_freedom, min=0.1)
+
+    gamma = torch.zeros((p_disp, n_genes))
+    gamma[0, :] = torch.log(torch.clip(dispersion, min=0.1))
+    return beta, gamma
+
+
+# ==============================================================================
+# Batch Log-Likelihood
+# ==============================================================================
+
+def compute_batch_loglikelihood(y, mu, r):
+    """
+    Compute the negative binomial log-likelihood for a batch.
+
+    Formula:
+        ℓ = Σ [log Γ(Y+θ) - log Γ(θ) - log Γ(Y+1) + θ log θ + Y log μ - (Y+θ)log(μ+θ)]
+
+    Parameters
+    ----------
+    y : torch.Tensor
+        Observed counts (n_batch × m_active)
+    mu : torch.Tensor
+        Predicted means (n_batch × m_active)
+    r : torch.Tensor
+        Dispersion parameters (n_batch × m_active)
+
+    Returns
+    -------
+    torch.Tensor
+        Total log-likelihood per response (m_active,)
+    """
+    ll = (
+        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)
+    )
+    return torch.sum(ll, dim=0)
+
+
+# ==============================================================================
+# Stochastic IRLS Step
+# ==============================================================================
+
+def step_stochastic_irls(
+    y,
+    X,
+    Z,
+    beta,
+    gamma,
+    eta: float = 0.8,
+    tol: float = 1e-4,
+    ll_prev: Optional[torch.Tensor] = None,
+    clip_mean: float = 10.0,
+    clip_disp: float = 10.0
+):
+    """
+    Perform a single damped Newton-Raphson update on a minibatch.
+
+    Logic:
+        1. Compute log-likelihood with current coefficients.
+        2. Perform one IRLS step for Mean (Beta) and Dispersion (Gamma).
+        3. Re-compute log-likelihood to determine convergence.
+        4. Return updated coefficients and boolean convergence mask.
+
+    Parameters
+    ----------
+    y : torch.Tensor
+        Count batch (n × m)
+    X : torch.Tensor
+        Mean design matrix (n × p)
+    Z : torch.Tensor
+        Dispersion design matrix (n × q)
+    beta : torch.Tensor
+        Current mean coefficients (p × m)
+    gamma : torch.Tensor
+        Current dispersion coefficients (q × m)
+    eta : float, optional
+        Damping factor (learning rate), 1.0 is pure Newton step, by default 0.8
+    tol : float, optional
+        Relative log-likelihood change threshold for convergence, by default 1e-4
+    ll_prev : torch.Tensor, optional
+        Previous log-likelihood values, by default None
+    clip_mean : float, optional
+        Maximum absolute value for mean linear predictor, by default 10.0
+    clip_disp : float, optional
+        Maximum absolute value for dispersion linear predictor, by default 10.0
+
+    Returns
+    -------
+    beta_next : torch.Tensor
+        Updated mean coefficients (p × m)
+    gamma_next : torch.Tensor
+        Updated dispersion coefficients (q × m)
+    converged : torch.Tensor
+        Boolean mask of converged responses (m,)
+    ll_next : torch.Tensor
+        Updated log-likelihood values (m,)
+    """
+    # --- 2. Update Mean (Beta) ---
+    # Working weights W = μ/(1 + μ/θ)
+    beta_target = update_mean_coefficients(X, y, beta, torch.exp(Z @ gamma), clip=clip_mean)
+    beta_next = (1 - eta) * beta + eta * beta_target
+
+    # --- 3. Update Dispersion (Gamma) ---
+    # Update depends on the latest mean estimates
+    linear_pred_mu = torch.clip(X @ beta_next, min=-clip_mean, max=clip_mean)
+    mu = torch.exp(linear_pred_mu)
+    gamma_target = update_dispersion_coefficients(Z, y, mu, gamma, clip=clip_disp)
+    gamma_next = (1 - eta) * gamma + eta * gamma_target
+
+    # --- 4. Convergence Check ---
+    linear_pred_r_next = torch.clip(Z @ gamma_next, min=-clip_disp, max=clip_disp)
+    ll_next = compute_batch_loglikelihood(y, mu, torch.exp(linear_pred_r_next))
+
+    # Relative improvement in the objective function
+    if ll_prev is not None:
+        rel_change = torch.abs(ll_next - ll_prev) / (torch.abs(ll_prev) + 1e-10)
+        converged = rel_change <= tol
+    else:
+        converged = False
+    return beta_next, gamma_next, converged, ll_next

scdesigner/distributions/poisson.py

@@ -0,0 +1,88 @@
+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 poisson
+
+class Poisson(Marginal):
+    """Poisson marginal estimator
+
+    This subclass behaves like `Marginal` but assumes each gene follows a
+    Poisson distribution with mean `mu_j(x)` that depends on covariates `x`
+    via the provided `formula` object.
+
+    The allowed formula keys are 'mean', defaulting to a single `mean` term
+    if a string formula is supplied.
+
+    Examples
+    --------
+    >>> from scdesigner.distributions import Poisson
+    >>> from scdesigner.datasets import pancreas
+    >>>
+    >>> sim = Poisson(formula="~ bs(pseudotime, df=5)")
+    >>> 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'])
+        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)")
+
+        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")
+        return y * torch.log(mu) - mu - torch.lgamma(y + 1)
+
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """Invert pseudoobservations."""
+        mu, u = self._local_params(x, u)
+        y = poisson(mu).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, y = self._local_params(x, y)
+        u1 = poisson(mu).cdf(y)
+        u2 = np.where(y > 0, poisson(mu).cdf(y - 1), 0)
+        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) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+        params = self.predict(x)
+        mu = params.get('mean')
+        if y is None:
+            return _to_numpy(mu)
+        return _to_numpy(mu, y)

scdesigner/{minimal → distributions}/zero_inflated_negbin.py

@@ -1,13 +1,41 @@
-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
 
 class ZeroInflatedNegBin(Marginal):
-    """Zero-inflated negative-binomial marginal estimator"""
+    """Zero-inflated negative-binomial marginal estimator
+
+    This subclass models a two-part mixture for counts. For each feature
+    j the observation follows a mixture: with probability `pi_j(x)` the value
+    is an extra zero (inflation), otherwise the count is drawn from a
+    negative-binomial distribution NB(mu_j(x), r_j(x)) parameterized here via
+    a mean `mu_j(x)` and dispersion `r_j(x)`. All parameters may depend on
+    covariates `x` through the `formula` argument.
+
+    The allowed formula keys are 'mean', 'dispersion', and 'zero_inflation'.
+
+    Examples
+    --------
+    >>> from scdesigner.distributions import ZeroInflatedNegBin
+    >>> from scdesigner.datasets import pancreas
+    >>>
+    >>> sim = ZeroInflatedNegBin(formula={"mean": "~ pseudotime", "dispersion": "~ 1", "zero_inflation": "~ pseudotime"})
+    >>> sim.setup_data(pancreas)
+    >>> sim.fit(max_epochs=1, verbose=False)
+    >>>
+    >>> # evaluate p(y | x) and model parameters
+    >>> 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'])
         super().__init__(formula)
@@ -25,7 +53,8 @@ class ZeroInflatedNegBin(Marginal):
             "dispersion": torch.exp,
             "zero_inflation": 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,
@@ -35,7 +64,7 @@ class ZeroInflatedNegBin(Marginal):
             optimizer_kwargs=optimizer_kwargs
         )
 
-    def likelihood(self, batch):
+    def likelihood(self, batch) -> torch.Tensor:
         """Compute the negative log-likelihood"""
         y, x = batch
         params = self.predict(x)
@@ -56,14 +85,14 @@ class ZeroInflatedNegBin(Marginal):
         # 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]):
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]) -> 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):
+    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, pi, y = self._local_params(x, y)
@@ -76,7 +105,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)
         mu = params.get('mean')
         r = params.get('dispersion')