scdesigner 0.0.1__py3-none-any.whl

scdesigner/estimators/negbin.py

@@ -0,0 +1,129 @@
+from . import gaussian_copula_factory as gcf
+from . import glm_factory as factory
+from .. import format
+from .. import data
+from anndata import AnnData
+from scipy.stats import nbinom
+import numpy as np
+import pandas as pd
+import torch
+from typing import Union
+
+###############################################################################
+## Regression functions that operate on numpy arrays
+###############################################################################
+
+
+def negbin_regression_likelihood(params, X_dict, y):    
+    n_mean_features = X_dict["mean"].shape[1]
+    n_dispersion_features = X_dict["dispersion"].shape[1]
+    n_outcomes = y.shape[1]
+
+    # form the mean and dispersion parameters
+    coef_mean = params[: n_mean_features * n_outcomes].\
+        reshape(n_mean_features, n_outcomes)
+    coef_dispersion = params[n_mean_features * n_outcomes :].\
+        reshape(n_dispersion_features, n_outcomes)
+    r = torch.exp(X_dict["dispersion"] @ coef_dispersion)
+    mu = torch.exp(X_dict["mean"] @ coef_mean)
+
+    # compute the negative log likelihood
+    log_likelihood = (
+        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 -torch.sum(log_likelihood)
+
+
+def negbin_initializer(x_dict, y, device):
+    n_mean_features = x_dict["mean"].shape[1]
+    n_outcomes = y.shape[1]
+    n_dispersion_features = x_dict["dispersion"].shape[1]
+    return torch.zeros(
+        n_mean_features * n_outcomes\
+            + n_dispersion_features * n_outcomes, 
+        requires_grad=True, device=device
+    )
+
+
+def negbin_postprocessor(params, x_dict, y):
+    n_mean_features = x_dict["mean"].shape[1]
+    n_outcomes = y.shape[1]
+    n_dispersion_features = x_dict["dispersion"].shape[1]
+    coef_mean = format.to_np(params[:n_mean_features * n_outcomes]).\
+        reshape(n_mean_features, n_outcomes)
+    coef_dispersion = format.to_np(params[n_mean_features * n_outcomes:]).\
+        reshape(n_dispersion_features, n_outcomes)
+    return {"coef_mean": coef_mean, "coef_dispersion": coef_dispersion}
+
+
+negbin_regression_array = factory.multiple_formula_regression_factory(
+    negbin_regression_likelihood, negbin_initializer, negbin_postprocessor
+)
+
+
+###############################################################################
+## Regression functions that operate on AnnData objects
+###############################################################################
+
+def format_negbin_parameters(
+    parameters: dict, var_names: list, mean_coef_index: list, 
+    dispersion_coef_index: list
+) -> dict:
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=mean_coef_index
+    )
+    parameters["coef_dispersion"] = pd.DataFrame(
+        parameters["coef_dispersion"], columns=var_names, index=dispersion_coef_index
+    )
+    return parameters
+
+def format_negbin_parameters_with_loaders(
+    parameters: dict, var_names: list, dls: dict
+) -> dict:
+    # Extract the coefficient indices from the dataloaders
+    mean_coef_index = dls["mean"].dataset.x_names
+    dispersion_coef_index = dls["dispersion"].dataset.x_names
+    
+    return format_negbin_parameters(parameters, var_names, mean_coef_index, dispersion_coef_index)
+
+def negbin_regression(
+    adata: AnnData, formula: Union[str, dict], chunk_size: int = int(1e4), batch_size=512, **kwargs
+) -> dict:
+    formula = data.standardize_formula(formula, allowed_keys=['mean', 'dispersion'])
+    
+    loaders = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+    parameters = negbin_regression_array(loaders, **kwargs)
+    return format_negbin_parameters(
+        parameters, list(adata.var_names), loaders["mean"].dataset.x_names, loaders["dispersion"].dataset.x_names
+    )
+
+###############################################################################
+## Copula versions for negative binomial regression
+###############################################################################
+
+
+def negbin_uniformizer(parameters, X_dict, y, epsilon=1e-3):
+    r = np.exp(X_dict["dispersion"] @ parameters["coef_dispersion"])
+    mu = np.exp(X_dict["mean"] @ parameters["coef_mean"])
+    u1 = nbinom(n=r, p=r / (r + mu)).cdf(y)
+    u2 = np.where(y > 0, nbinom(n=r, p=r / (r + mu)).cdf(y - 1), 0)
+    v = np.random.uniform(size=y.shape)
+    return np.clip(v * u1 + (1 - v) * u2, epsilon, 1 - epsilon)
+
+
+negbin_copula_array = gcf.gaussian_copula_array_factory(
+    negbin_regression_array, negbin_uniformizer
+) # should accept a dictionary of dataloaders
+
+negbin_copula = gcf.gaussian_copula_factory(
+    negbin_copula_array, format_negbin_parameters_with_loaders, 
+    param_name=['mean', 'dispersion']
+)

scdesigner/estimators/pnmf.py

@@ -0,0 +1,160 @@
+from scipy.sparse.linalg import svds
+from sklearn.cluster import KMeans
+import numpy as np
+import pandas as pd
+import torch
+
+
+# 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, batch_size: int = 512, 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 = to_np(a).reshape(n_features, n_outcomes)
+    loc = to_np(loc).reshape(n_features, n_outcomes)
+    beta = to_np(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
+    """
+    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):
+    # 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):
+    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 to_np(x):
+    return x.detach().cpu().numpy()
+
+
+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

scdesigner/estimators/poisson.py

@@ -0,0 +1,100 @@
+from . import gaussian_copula_factory as gcf
+from . import glm_factory as factory
+from .. import data
+from .. import format
+from anndata import AnnData
+from scipy.stats import poisson
+import numpy as np
+import pandas as pd
+import torch
+
+###############################################################################
+## Regression functions that operate on numpy arrays
+###############################################################################
+
+
+def poisson_regression_likelihood(params, X, y, epsilon=1e-6):
+    # get appropriate parameter shape
+    n_features = X['mean'].shape[1]
+    n_outcomes = y.shape[1]
+
+    # compute the negative log likelihood
+    beta = params.reshape(n_features, n_outcomes)
+    mu = torch.exp(X['mean'] @ beta)
+    log_likelihood = y * torch.log(mu+epsilon) - mu - torch.lgamma(y + 1)
+    return -torch.sum(log_likelihood)
+
+
+def poisson_initializer(x, y, device):
+    n_features, n_outcomes = x['mean'].shape[1], y.shape[1]
+    return torch.zeros(n_features * n_outcomes, requires_grad=True, device=device)
+
+
+def poisson_postprocessor(params, x, y):
+    coef_mean = format.to_np(params).reshape(x['mean'].shape[1], y.shape[1])
+    return {"coef_mean": coef_mean}
+
+
+poisson_regression_array = factory.multiple_formula_regression_factory(
+    poisson_regression_likelihood, poisson_initializer, poisson_postprocessor
+)
+
+###############################################################################
+## Regression functions that operate on AnnData objects
+###############################################################################
+
+
+def format_poisson_parameters(
+    parameters: dict, var_names: list, coef_index: list
+) -> dict:
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=coef_index
+    )
+    return parameters
+
+
+def poisson_regression(
+    adata: AnnData,
+    formula: str,
+    chunk_size: int = int(1e4),
+    batch_size: int = 512,
+    **kwargs
+) -> dict:
+    formula = data.standardize_formula(formula, allowed_keys=['mean'])
+    loaders = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+
+    parameters = poisson_regression_array(loaders, **kwargs)
+    return format_poisson_parameters(
+        parameters, list(adata.var_names), loaders["mean"].dataset.x_names
+    )
+
+
+###############################################################################
+## Copula versions for poisson regression
+###############################################################################
+
+
+def poisson_uniformizer(parameters, x, y, epsilon=1e-3):
+    mu = np.exp(x['mean'] @ parameters["coef_mean"])
+    u1 = poisson(mu).cdf(y)
+    u2 = np.where(y > 0, poisson(mu).cdf(y - 1), 0)
+    v = np.random.uniform(size=y.shape)
+    return np.clip(v * u1 + (1 - v) * u2, epsilon, 1 - epsilon)
+
+def format_poisson_parameters_with_loaders(parameters: dict, var_names: list, dls: dict) -> dict:
+    beta_coef_index = dls["mean"].dataset.x_names
+    
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=beta_coef_index
+    )
+    return parameters
+
+poisson_copula_array = gcf.gaussian_copula_array_factory(
+    poisson_regression_array, poisson_uniformizer
+) 
+
+poisson_copula = gcf.gaussian_copula_factory(
+    poisson_copula_array, format_poisson_parameters_with_loaders, ['mean']
+)

scdesigner/estimators/zero_inflated_negbin.py

@@ -0,0 +1,195 @@
+import warnings
+from . import gaussian_copula_factory as gcf
+from .. import format
+from .. import data
+from . import glm_factory as factory
+from anndata import AnnData
+from scipy.stats import nbinom
+import numpy as np
+import pandas as pd
+import torch
+from typing import Union
+from scipy.special import expit
+
+###############################################################################
+## Regression functions that operate on numpy arrays
+###############################################################################
+
+
+def zero_inflated_negbin_regression_likelihood(params, X_dict, y):
+    # get appropriate parameter shape
+    n_mean_features = X_dict["mean"].shape[1]
+    n_dispersion_features = X_dict["dispersion"].shape[1]
+    n_zero_inflation_features = X_dict["zero_inflation"].shape[1]
+    n_outcomes = y.shape[1]
+
+    # define the likelihood parameters
+    coef_mean = params[: n_mean_features * n_outcomes].\
+        reshape(n_mean_features, n_outcomes)
+    coef_dispersion = params[n_mean_features * n_outcomes :\
+        n_mean_features * n_outcomes + n_dispersion_features * n_outcomes].\
+        reshape(n_dispersion_features, n_outcomes)
+    coef_zero_inflation = params[n_mean_features * n_outcomes + \
+        n_dispersion_features * n_outcomes :].\
+        reshape(n_zero_inflation_features, n_outcomes)
+    
+    mu = torch.exp(X_dict["mean"] @ coef_mean)
+    r = torch.exp(X_dict["dispersion"] @ coef_dispersion)
+    pi = torch.sigmoid(X_dict["zero_inflation"] @ coef_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)
+    log_likelihood = torch.log(
+        pi * (y == 0) + (1 - pi) * torch.exp(negbin_loglikelihood) + 1e-10
+    )
+    return -torch.sum(log_likelihood)
+
+
+def zero_inflated_negbin_initializer(X_dict, y, device):
+    n_mean_features = X_dict["mean"].shape[1]
+    n_dispersion_features = X_dict["dispersion"].shape[1]
+    n_zero_inflation_features = X_dict["zero_inflation"].shape[1]
+    n_outcomes = y.shape[1]
+    return torch.zeros(
+        n_mean_features * n_outcomes + n_dispersion_features * n_outcomes + \
+        n_zero_inflation_features * n_outcomes, requires_grad=True, device=device
+    )
+
+
+def zero_inflated_negbin_postprocessor(params, X_dict, y):
+    n_mean_features = X_dict["mean"].shape[1]
+    n_dispersion_features = X_dict["dispersion"].shape[1]
+    n_zero_inflation_features = X_dict["zero_inflation"].shape[1]
+    n_outcomes = y.shape[1]
+    coef_mean = format.to_np(params[:n_mean_features * n_outcomes]).\
+        reshape(n_mean_features, n_outcomes)
+    coef_dispersion = format.to_np(params[n_mean_features * n_outcomes\
+        : n_mean_features * n_outcomes + n_dispersion_features * n_outcomes]).\
+        reshape(n_dispersion_features, n_outcomes)
+    coef_zero_inflation = format.to_np(params[n_mean_features * n_outcomes \
+        + n_dispersion_features * n_outcomes :]).\
+        reshape(n_zero_inflation_features, n_outcomes)
+    return {"coef_mean": coef_mean, "coef_dispersion": coef_dispersion,\
+        "coef_zero_inflation": coef_zero_inflation}
+
+
+zero_inflated_negbin_regression_array = factory.multiple_formula_regression_factory(
+    zero_inflated_negbin_regression_likelihood,
+    zero_inflated_negbin_initializer,
+    zero_inflated_negbin_postprocessor,
+)
+
+###############################################################################
+## Regression functions that operate on AnnData objects
+###############################################################################
+
+
+def format_zero_inflated_negbin_parameters(
+    parameters: dict, var_names: list, mean_coef_index: 
+        list, dispersion_coef_index: list, zero_inflation_coef_index: list
+) -> dict:
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=mean_coef_index
+    )
+    parameters["coef_dispersion"] = pd.DataFrame(
+        parameters["coef_dispersion"], columns=var_names, index=dispersion_coef_index
+    )
+    parameters["coef_zero_inflation"] = pd.DataFrame(
+        parameters["coef_zero_inflation"], columns=var_names, index=zero_inflation_coef_index
+    )
+    return parameters
+
+def format_zero_inflated_negbin_parameters_with_loaders(
+    parameters: dict, var_names: list, dls: dict
+) -> dict:
+    mean_coef_index = dls["mean"].dataset.x_names
+    dispersion_coef_index = dls["dispersion"].dataset.x_names
+    zero_inflation_coef_index = dls["zero_inflation"].dataset.x_names
+    return format_zero_inflated_negbin_parameters(
+        parameters, var_names, mean_coef_index, dispersion_coef_index, zero_inflation_coef_index
+    )
+
+def standardize_zero_inflated_negbin_formula(formula: Union[str, dict]) -> dict:
+    '''
+    Convert string formula to dict and validate type.
+    If formula is a string, it is the formula for the mean parameter.
+    If formula is a dictionary, it is a dictionary of formulas for the mean, dispersion, and zero_inflation parameters.
+    '''
+    # Convert string formula to dict and validate type
+    formula = {'mean': formula, 'dispersion': '~ 1', 'zero_inflation': '~ 1'} \
+        if isinstance(formula, str) else formula
+    if not isinstance(formula, dict):
+        raise ValueError("formula must be a string or a dictionary")
+    
+    # Define allowed keys and set defaults
+    allowed_keys = {'mean', 'dispersion', 'zero_inflation'}
+    formula_keys = set(formula.keys())
+
+    # check for required keys and warn about extras
+    if not formula_keys & allowed_keys:
+        raise ValueError("formula must have at least one of \
+                the following keys: mean, dispersion, zero_inflation")
+    
+    # warn about unused keys
+    if extra_keys := formula_keys - allowed_keys:
+        warnings.warn(
+            f"Invalid formulas in dictionary for zero-inflated \
+                negative binomial regression: {extra_keys}",
+            UserWarning,
+        )
+    
+    # set default values for missing keys
+    formula.update({k: '~ 1' for k in allowed_keys - formula_keys})
+    return formula
+
+
+def zero_inflated_negbin_regression(
+    adata: AnnData, formula: Union[str, dict], chunk_size: int = int(1e4), batch_size=512, **kwargs
+) -> dict:
+    formula = data.standardize_formula(formula, allowed_keys=['mean', 'dispersion', 'zero_inflation'])
+    
+    loaders = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+
+    parameters = zero_inflated_negbin_regression_array(loaders, **kwargs)
+    return format_zero_inflated_negbin_parameters(
+        parameters, list(adata.var_names), loaders["mean"].dataset.x_names, 
+        loaders["dispersion"].dataset.x_names, loaders["zero_inflation"].dataset.x_names
+    )
+
+
+###############################################################################
+## Copula versions for ZINB regression
+###############################################################################
+
+
+def zero_inflated_negbin_uniformizer(parameters, X_dict, y, epsilon=1e-3):
+    r, mu, pi = (
+        np.exp(X_dict["dispersion"] @ parameters["coef_dispersion"]),
+        np.exp(X_dict["mean"] @ parameters["coef_mean"]),
+        expit(X_dict["zero_inflation"] @ parameters["coef_zero_inflation"]),
+    )
+    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)
+    v = np.random.uniform(size=y.shape)
+    return np.clip(v * u1 + (1 - v) * u2, epsilon, 1 - epsilon)
+
+
+zero_inflated_negbin_copula = gcf.gaussian_copula_factory(
+    gcf.gaussian_copula_array_factory(
+        zero_inflated_negbin_regression_array, zero_inflated_negbin_uniformizer
+    ),
+    format_zero_inflated_negbin_parameters_with_loaders,
+    param_name=['mean', 'dispersion', 'zero_inflation'],
+)

scdesigner/estimators/zero_inflated_poisson.py

@@ -0,0 +1,85 @@
+from anndata import AnnData
+from .. import format
+from .. import data
+from . import glm_factory as factory
+import pandas as pd
+import torch
+
+
+def zero_inflated_poisson_regression_likelihood(params, X, y):
+    # get appropriate parameter shape
+    mean_n_features = X['mean'].shape[1]
+    zero_inflation_n_features = X['zero_inflation'].shape[1]
+    n_outcomes = y.shape[1]
+
+    # define the likelihood parameters
+    b_elem = mean_n_features * n_outcomes
+    coef_mean = params[:b_elem].reshape(mean_n_features, n_outcomes)
+    coef_zero_inflation = params[b_elem:].reshape(zero_inflation_n_features, n_outcomes)
+
+    zero_inflation = torch.sigmoid(X['zero_inflation'] @ coef_zero_inflation)
+    mu = torch.exp(X['mean'] @ coef_mean)
+    poisson_loglikelihood = y * torch.log(mu + 1e-10) - mu - torch.lgamma(y + 1)
+
+    # return the mixture, with an offset to prevent log(0)
+    log_likelihood = torch.log(
+        zero_inflation * (y == 0) + (1 - zero_inflation) * torch.exp(poisson_loglikelihood) + 1e-10
+    )
+    return -torch.sum(log_likelihood)
+
+
+def zero_inflated_poisson_initializer(x, y, device):
+    mean_n_features = x['mean'].shape[1]
+    zero_inflation_n_features = x['zero_inflation'].shape[1]
+    n_outcomes = y.shape[1]
+    return torch.zeros(
+        mean_n_features * n_outcomes + zero_inflation_n_features * n_outcomes, requires_grad=True, device=device
+    )
+
+
+def zero_inflated_poisson_postprocessor(params, x, y):
+    mean_n_features = x['mean'].shape[1]
+    zero_inflation_n_features = x['zero_inflation'].shape[1]
+    n_outcomes = y.shape[1]
+    b_elem = mean_n_features * n_outcomes
+    coef_mean = format.to_np(params[:b_elem]).reshape(mean_n_features, n_outcomes)
+    coef_zero_inflation = format.to_np(params[b_elem:]).reshape(zero_inflation_n_features, n_outcomes)
+    return {"coef_mean": coef_mean, "coef_zero_inflation": coef_zero_inflation}
+
+
+zero_inflated_poisson_regression_array = factory.multiple_formula_regression_factory(
+    zero_inflated_poisson_regression_likelihood,
+    zero_inflated_poisson_initializer,
+    zero_inflated_poisson_postprocessor,
+)
+
+
+###############################################################################
+## Regression functions that operate on AnnData objects
+###############################################################################
+
+
+def format_zero_inflated_poisson_parameters(
+    parameters: dict, var_names: list, mean_coef_index: list, 
+    zero_inflation_coef_index: list
+) -> dict:
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=mean_coef_index
+    )
+    parameters["coef_zero_inflation"] = pd.DataFrame(
+        parameters["coef_zero_inflation"], columns=var_names, index=zero_inflation_coef_index
+    )
+    return parameters
+
+
+def zero_inflated_poisson_regression(
+    adata: AnnData, formula: str, chunk_size: int = int(1e4), batch_size=512, **kwargs
+) -> dict:
+    formula = data.standardize_formula(formula, allowed_keys=['mean', 'zero_inflation'])
+    loaders = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+    parameters = zero_inflated_poisson_regression_array(loaders, **kwargs)
+    return format_zero_inflated_poisson_parameters(
+        parameters, list(adata.var_names), loaders["mean"].dataset.x_names, loaders["zero_inflation"].dataset.x_names
+    )

scdesigner/format/__init__.py

@@ -0,0 +1,4 @@
+from .format import to_np, format_input_anndata, format_matrix
+from .print import print_simulator
+
+__all__ = ["to_np", "format_input_anndata", "format_matrix", "print_simulator"]

scdesigner/format/format.py

@@ -0,0 +1,20 @@
+from anndata import AnnData
+from formulaic import model_matrix
+import pandas as pd
+import scipy.sparse
+
+def to_np(x):
+    return x.detach().cpu().numpy()
+
+def format_input_anndata(adata: AnnData) -> AnnData:
+    result = adata.copy()
+    if isinstance(result.X, scipy.sparse._csc.csc_matrix):
+        result.X = result.X.todense()
+    return result
+
+def format_matrix(obs: pd.DataFrame, formula: str):
+    if formula is not None:
+        x = model_matrix(formula, pd.DataFrame(obs))
+    else:
+        x = obs
+    return x