scdesigner 0.0.1__py3-none-any.whl

scdesigner/diagnose/plot.py

@@ -0,0 +1,242 @@
+import scanpy as sc
+import numpy as np
+import pandas as pd
+import altair as alt
+import matplotlib.pyplot as plt
+
+
+def adata_df(adata):
+    return (
+        pd.DataFrame(adata.X, columns=adata.var_names)
+        .melt(id_vars=[], value_vars=adata.var_names)
+        .reset_index(drop=True)
+    )
+
+
+def merge_samples(adata, sim):
+    source = adata_df(adata)
+    simulated = adata_df(sim)
+    return pd.concat(
+        {"real": source, "simulated": simulated}, names=["source"]
+    ).reset_index(level="source")
+
+
+def plot_umap(
+    adata,
+    color=None,
+    shape=None,
+    facet=None,
+    opacity=0.6,
+    n_comps=20,
+    n_neighbors=15,
+    transform=lambda x: np.log1p(x),
+    **kwargs
+):
+    mapping = {"x": "UMAP1", "y": "UMAP2", "color": color, "shape": shape}
+    mapping = {k: v for k, v in mapping.items() if v is not None}
+
+    adata_ = adata.copy()
+    adata_.X = check_sparse(adata_.X)
+    Z = transform(adata_.X)
+    if Z.shape[1] == adata_.X.shape[1]:
+        adata_.X = transform(adata_.X)
+    else:
+        adata_ = adata_[:, : Z.shape[1]]
+        adata_.X = Z
+        adata_.var_names = [f"transform_{k}" for k in range(Z.shape[1])]
+
+    # umap on the top PCA dimensions
+    sc.pp.pca(adata_, n_comps=n_comps)
+    sc.pp.neighbors(adata_, n_neighbors=n_neighbors, n_pcs=n_comps)
+    sc.tl.umap(adata_, **kwargs)
+
+    # get umap embeddings
+    umap_df = pd.DataFrame(adata_.obsm["X_umap"], columns=["UMAP1", "UMAP2"])
+    umap_df = pd.concat([umap_df, adata_.obs.reset_index(drop=True)], axis=1)
+
+    # encode and visualize
+    chart = alt.Chart(umap_df).mark_point(opacity=opacity).encode(**mapping)
+    if facet is not None:
+        chart = chart.facet(column=alt.Facet(facet))
+    return chart
+
+
+def plot_pca(
+    adata,
+    color=None,
+    shape=None,
+    facet=None,
+    opacity=0.6,
+    plot_dims=[0, 1],
+    transform=lambda x: np.log1p(x),
+    **kwargs
+):
+    mapping = {"x": "PCA1", "y": "PCA2", "color": color, "shape": shape}
+    mapping = {k: v for k, v in mapping.items() if v is not None}
+
+    adata_ = adata.copy()
+    adata_.X = check_sparse(adata_.X)
+    adata_.X = transform(adata_.X)
+
+    # get PCA scores
+    sc.pp.pca(adata_, **kwargs)
+    pca_df = pd.DataFrame(adata_.obsm["X_pca"][:, plot_dims], columns=["PCA1", "PCA2"])
+    pca_df = pd.concat([pca_df, adata_.obs.reset_index(drop=True)], axis=1)
+
+    # plot
+    chart = alt.Chart(pca_df).mark_point(opacity=opacity).encode(**mapping)
+    if facet is not None:
+        chart = chart.facet(column=alt.Facet(facet))
+    return chart
+
+
+def compare_summary(real, simulated, summary_fun):
+    df = pd.DataFrame({"real": summary_fun(real), "simulated": summary_fun(simulated)})
+
+    identity = pd.DataFrame(
+        {
+            "real": [df["real"].min(), df["real"].max()],
+            "simulated": [df["real"].min(), df["real"].max()],
+        }
+    )
+    return alt.Chart(identity).mark_line(color="#dedede").encode(
+        x="real", y="simulated"
+    ) + alt.Chart(df).mark_circle().encode(x="real", y="simulated")
+
+
+def check_sparse(X):
+    if not isinstance(X, np.ndarray):
+        X = X.todense()
+    return X
+
+
+def compare_means(real, simulated, transform=lambda x: x):
+    real_, simulated_ = prepare_dense(real, simulated)
+    summary = lambda a: np.asarray(transform(a.X).mean(axis=0)).flatten()
+    return compare_summary(real_, simulated_, summary)
+
+
+def prepare_dense(real, simulated):
+    real_ = real.copy()
+    simulated_ = simulated.copy()
+    real_.X = check_sparse(real_.X)
+    simulated_.X = check_sparse(simulated_.X)
+    return real_, simulated_
+
+
+def compare_variances(real, simulated, transform=lambda x: x):
+    real_, simulated_ = prepare_dense(real, simulated)
+    summary = lambda a: np.asarray(np.var(transform(a.X), axis=0)).flatten()
+    return compare_summary(real_, simulated_, summary)
+
+
+def compare_standard_deviation(real, simulated, transform=lambda x: x):
+    real_, simulated_ = prepare_dense(real, simulated)
+    summary = lambda a: np.asarray(np.std(transform(a.X), axis=0)).flatten()
+    return compare_summary(real_, simulated_, summary)
+
+
+def concat_real_sim(real, simulated):
+    real_, simulated_ = prepare_dense(real, simulated)
+    real_.obs["source"] = "real"
+    simulated_.obs["source"] = "simulated"
+    return real_.concatenate(simulated_, join="outer", batch_key=None)
+
+
+def compare_umap(real, simulated, transform=lambda x: x, **kwargs):
+    adata = concat_real_sim(real, simulated)
+    return plot_umap(adata, facet="source", transform=transform, **kwargs)
+
+
+def compare_pca(real, simulated, transform=lambda x: x, **kwargs):
+    adata = concat_real_sim(real, simulated)
+    return plot_pca(adata, facet="source", transform=transform, **kwargs)
+
+
+def plot_hist(sim_data, real_data, idx):
+    sim = sim_data[:, idx]
+    real = real_data[:, idx]
+    b = np.linspace(min(min(sim), min(real)), max(max(sim), max(real)), 50)
+
+    plt.hist([real, sim], b, label=["Real", "Simulated"], histtype="bar")
+    plt.xlabel("x")
+    plt.ylabel("Density")
+    plt.legend()
+    plt.show()
+
+
+def compare_ecdf(adata, sim, var_names=None, max_plot=10, n_cols=5, **kwargs):
+    if var_names is None:
+        var_names = adata.var_names[:max_plot]
+
+    combined = merge_samples(adata[:, var_names], sim.sample()[:, var_names])
+    alt.data_transformers.enable("vegafusion")
+
+    plot = (
+        alt.Chart(combined)
+        .transform_window(
+            ecdf="cume_dist()", sort=[{"field": "value"}], groupby=["variable"]
+        )
+        .mark_line(
+            interpolate="step-after",
+        )
+        .encode(
+            x="value:Q",
+            y="ecdf:Q",
+            color="source:N",
+            facet=alt.Facet(
+                "variable", sort=alt.EncodingSortField("value"), columns=n_cols
+            ),
+        )
+        .properties(**kwargs)
+    )
+    plot.show()
+    return plot, combined
+
+
+def compare_boxplot(adata, sim, var_names=None, max_plot=20, **kwargs):
+    if var_names is None:
+        var_names = adata.var_names[:max_plot]
+
+    combined = merge_samples(adata[:, var_names], sim.sample()[:, var_names])
+    alt.data_transformers.enable("vegafusion")
+
+    plot = (
+        alt.Chart(combined)
+        .mark_boxplot(extent="min-max")
+        .encode(
+            x=alt.X("value:Q").scale(zero=False),
+            y=alt.Y(
+                "variable:N",
+                sort=alt.EncodingSortField("mid_box_value", order="descending"),
+            ),
+            facet="source:N",
+        )
+        .properties(**kwargs)
+    )
+    plot.show()
+    return plot, combined
+
+
+def compare_histogram(adata, sim, var_names=None, max_plot=20, n_cols=5, **kwargs):
+    if var_names is None:
+        var_names = adata.var_names[:max_plot]
+
+    combined = merge_samples(adata[:, var_names], sim.sample()[:, var_names])
+    alt.data_transformers.enable("vegafusion")
+
+    plot = (
+        alt.Chart(combined)
+        .mark_bar(opacity=0.7)
+        .encode(
+            x=alt.X("value:Q").bin(maxbins=20),
+            y=alt.Y("count()").stack(None),
+            color="source:N",
+            facet=alt.Facet(
+                "variable", sort=alt.EncodingSortField("bin_maxbins_20_value")
+            ),
+        )
+        .properties(**kwargs)
+    )
+    plot.show()
+    return plot, combined

scdesigner/estimators/__init__.py

@@ -0,0 +1,27 @@
+from .negbin import negbin_regression, negbin_copula
+from .gaussian_copula_factory import group_indices
+from .poisson import poisson_regression, poisson_copula
+from .bernoulli import bernoulli_regression, bernoulli_copula
+from .gaussian import gaussian_regression, gaussian_copula
+from .zero_inflated_negbin import (
+    zero_inflated_negbin_regression,
+    zero_inflated_negbin_copula,
+)
+from .zero_inflated_poisson import zero_inflated_poisson_regression
+from .glm_factory import multiple_formula_regression_factory
+
+__all__ = [
+    "bernoulli_regression",
+    "bernoulli_copula",
+    "negbin_copula",
+    "negbin_regression",
+    "gaussian_regression",
+    "gaussian_copula",
+    "group_indices",
+    "poisson_copula",
+    "poisson_regression",
+    "zero_inflated_negbin_copula",
+    "zero_inflated_negbin_regression",
+    "zero_inflated_poisson_regression",
+    "multiple_formula_regression_factory",
+]

scdesigner/estimators/bernoulli.py

@@ -0,0 +1,85 @@
+import pandas as pd
+from . import gaussian_copula_factory as gcf
+from . import glm_factory as factory
+from .. import data
+from .. import format
+from . import poisson as poi
+from anndata import AnnData
+from scipy.stats import bernoulli
+import numpy as np
+import torch
+from typing import Union
+
+###############################################################################
+## Regression functions that operate on numpy arrays
+###############################################################################
+
+
+def bernoulli_regression_likelihood(params, X_dict, y):
+    # get appropriate parameter shape
+    n_features = X_dict["mean"].shape[1]
+    n_outcomes = y.shape[1]
+
+    # compute the negative log likelihood
+    coef_mean = params.reshape(n_features, n_outcomes)
+    theta = torch.sigmoid(X_dict["mean"] @ coef_mean)
+    log_likelihood = y * torch.log(theta) + (1 - y) * torch.log(1 - theta)
+    return -torch.sum(log_likelihood)
+
+def bernoulli_initializer(X_dict, y, device):
+    n_features = X_dict["mean"].shape[1]
+    n_outcomes = y.shape[1]
+    return torch.zeros(n_features * n_outcomes, requires_grad=True, device=device)
+
+def bernoulli_postprocessor(params, X_dict, y):
+    coef_mean = format.to_np(params).reshape(X_dict["mean"].shape[1], y.shape[1])
+    return {"coef_mean": coef_mean}
+
+bernoulli_regression_array = factory.multiple_formula_regression_factory(
+    bernoulli_regression_likelihood, bernoulli_initializer, bernoulli_postprocessor
+)
+
+###############################################################################
+## Regression functions that operate on AnnData objects
+###############################################################################
+
+
+def bernoulli_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'])
+    loaders = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+    parameters = bernoulli_regression_array(loaders, **kwargs)
+    return poi.format_poisson_parameters(
+        parameters, list(adata.var_names), list(loaders["mean"].dataset.x_names)
+    )
+
+
+###############################################################################
+## Copula versions for bernoulli regression
+###############################################################################
+
+
+def bernoulli_uniformizer(parameters, X_dict, y, epsilon=1e-3, random_seed=42):
+    np.random.seed(random_seed)
+    theta = torch.sigmoid(torch.from_numpy(X_dict["mean"] @ parameters["coef_mean"])).numpy()
+    u1 =  bernoulli(theta).cdf(y)
+    u2 = np.where(y > 0,  bernoulli(theta).cdf(y - 1), 0)
+    v = np.random.uniform(size=y.shape)
+    return np.clip(v * u1 + (1 - v) * u2, epsilon, 1 - epsilon)
+
+def format_bernoulli_parameters_with_loaders(parameters: dict, var_names: list, dls: dict) -> dict:
+    coef_mean_index = dls["mean"].dataset.x_names
+    
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=coef_mean_index
+    )
+    return parameters
+
+bernoulli_copula = gcf.gaussian_copula_factory(
+    gcf.gaussian_copula_array_factory(bernoulli_regression_array, bernoulli_uniformizer),
+    format_bernoulli_parameters_with_loaders,
+    param_name=['mean']
+)

scdesigner/estimators/gaussian.py

@@ -0,0 +1,121 @@
+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 norm
+import numpy as np
+import pandas as pd
+import torch
+from typing import Union
+
+###############################################################################
+## Regression functions that operate on numpy arrays
+###############################################################################
+
+
+
+# Gaussian regression likelihood: regression onto mean and sdev
+def gaussian_regression_likelihood(params, X_dict, y):
+    n_mean_features = X_dict["mean"].shape[1]
+    n_sdev_features = X_dict["sdev"].shape[1]
+    n_outcomes = y.shape[1]
+
+    coef_mean = params[: n_mean_features * n_outcomes].reshape(n_mean_features, n_outcomes)
+    coef_sdev = params[n_mean_features * n_outcomes :].reshape(n_sdev_features, n_outcomes)
+    mu = X_dict["mean"] @ coef_mean
+    sigma = torch.exp(X_dict["sdev"] @ coef_sdev)
+
+    # Negative log likelihood for Gaussian
+    log_likelihood = -0.5 * (torch.log(2 * torch.pi * sigma ** 2) + ((y - mu) ** 2) / (sigma ** 2))
+    return torch.sum(log_likelihood)
+
+
+
+def gaussian_initializer(x_dict, y, device):
+    n_mean_features = x_dict["mean"].shape[1]
+    n_outcomes = y.shape[1]
+    n_sdev_features = x_dict["sdev"].shape[1]
+    return torch.zeros(
+        n_mean_features * n_outcomes + n_sdev_features * n_outcomes,
+        requires_grad=True, device=device
+    )
+
+
+
+def gaussian_postprocessor(params, x_dict, y):
+    n_mean_features = x_dict["mean"].shape[1]
+    n_outcomes = y.shape[1]
+    n_sdev_features = x_dict["sdev"].shape[1]
+    coef_mean = format.to_np(params[:n_mean_features * n_outcomes]).reshape(n_mean_features, n_outcomes)
+    coef_sdev = format.to_np(params[n_mean_features * n_outcomes:]).reshape(n_sdev_features, n_outcomes)
+    return {"coef_mean": coef_mean, "coef_sdev": coef_sdev}
+
+
+
+gaussian_regression_array = factory.multiple_formula_regression_factory(
+    gaussian_regression_likelihood, gaussian_initializer, gaussian_postprocessor
+)
+
+
+###############################################################################
+## Regression functions that operate on AnnData objects
+###############################################################################
+
+
+def format_gaussian_parameters(
+    parameters: dict, var_names: list, mean_coef_index: list, sdev_coef_index: list
+) -> dict:
+    parameters["coef_mean"] = pd.DataFrame(
+        parameters["coef_mean"], columns=var_names, index=mean_coef_index
+    )
+    parameters["coef_sdev"] = pd.DataFrame(
+        parameters["coef_sdev"], columns=var_names, index=sdev_coef_index
+    )
+    return parameters
+
+
+def format_gaussian_parameters_with_loaders(
+    parameters: dict, var_names: list, dls: dict
+) -> dict:
+    mean_coef_index = dls["mean"].dataset.x_names
+    sdev_coef_index = dls["sdev"].dataset.x_names
+    return format_gaussian_parameters(
+        parameters, var_names, mean_coef_index, sdev_coef_index
+    )
+
+
+def gaussian_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', 'sdev'])
+    loaders = data.multiple_formula_loader(
+        adata, formula, chunk_size=chunk_size, batch_size=batch_size
+    )
+    parameters = gaussian_regression_array(loaders, **kwargs)
+    return format_gaussian_parameters(
+        parameters, list(adata.var_names), loaders["mean"].dataset.x_names, 
+        loaders["sdev"].dataset.x_names
+    )
+
+
+###############################################################################
+## Copula versions for gaussian regression
+###############################################################################
+
+def gaussian_uniformizer(parameters, X_dict, y, epsilon=1e-3):
+    mu = X_dict["mean"] @ parameters["coef_mean"]
+    sigma = np.exp(X_dict["sdev"] @ parameters["coef_sdev"])
+    u = norm.cdf(y, loc=mu, scale=sigma)
+    u = np.clip(u, epsilon, 1 - epsilon)
+    return u
+
+gaussian_copula_array = gcf.gaussian_copula_array_factory(
+    gaussian_regression_array, gaussian_uniformizer
+)
+
+gaussian_copula = gcf.gaussian_copula_factory(
+    gaussian_copula_array, format_gaussian_parameters_with_loaders,
+    param_name=['mean', 'sdev']
+)

scdesigner/estimators/gaussian_copula_factory.py

@@ -0,0 +1,152 @@
+from ..data import stack_collate, multiple_formula_group_loader
+from .. import data
+from anndata import AnnData
+from collections.abc import Callable
+from typing import Union
+from scipy.stats import norm
+from torch.utils.data import DataLoader
+import numpy as np
+import pandas as pd
+
+###############################################################################
+## General copula factory functions
+###############################################################################
+
+
+def gaussian_copula_array_factory(marginal_model: Callable, uniformizer: Callable):
+    def copula_fun(loaders: dict[str, DataLoader], lr: float = 0.1, epochs: int = 40, **kwargs):
+        # for the marginal model, ignore the groupings
+        # Strip all dataloaders and create a dictionary to pass to marginal_model
+        formula_loaders = {}
+        for key in loaders.keys():
+            formula_loaders[key] = strip_dataloader(loaders[key], pop="Stack" in type(loaders[key].dataset).__name__)
+        
+        # Call marginal_model with the dictionary of stripped dataloaders
+        parameters = marginal_model(formula_loaders, lr=lr, epochs=epochs, **kwargs)
+
+        # estimate covariance, allowing for different groups
+        parameters["covariance"] = copula_covariance(parameters, loaders, uniformizer)
+        return parameters
+
+    return copula_fun
+
+
+def gaussian_copula_factory(copula_array_fun: Callable, 
+                            parameter_formatter: Callable, 
+                            param_name: list = None):
+    def copula_fun(
+        adata: AnnData,
+        formula: Union[str, dict] = "~ 1",
+        grouping_var: str = None,
+        chunk_size: int = int(1e4),
+        batch_size: int = 512,
+        **kwargs
+    ) -> dict:  
+        
+        if param_name is not None:
+            formula = data.standardize_formula(formula, param_name)
+        
+        dls = multiple_formula_group_loader(
+            adata,
+            formula,
+            grouping_var,
+            chunk_size=chunk_size,
+            batch_size=batch_size,
+        ) # returns a dictionary of dataloaders
+        parameters = copula_array_fun(dls, **kwargs)
+        
+        # Pass the full dls to parameter_formatter so it can extract what it needs
+        parameters = parameter_formatter(
+            parameters, adata.var_names, dls
+        )
+        parameters["covariance"] = format_copula_parameters(parameters, adata.var_names)
+        return parameters
+
+    return copula_fun
+
+
+def copula_covariance(parameters: dict, loaders: dict[str, DataLoader], uniformizer: Callable):
+    first_loader = next(iter(loaders.values()))
+    D = next(iter(first_loader))[1].shape[1] #dimension of y
+    groups = first_loader.dataset.groups # a list of strings of group names
+    sums = {g: np.zeros(D) for g in groups}
+    second_moments = {g: np.eye(D) for g in groups}
+    Ng = {g: 0 for g in groups}
+    keys = list(loaders.keys())
+    loaders = list(loaders.values())
+    num_keys = len(keys)
+    
+    for batches in zip(*loaders):
+        x_batch_dict = {
+            keys[i]: batches[i][0].cpu().numpy() for i in range(num_keys)
+        }
+        y_batch = batches[0][1].cpu().numpy()
+        memberships = batches[0][2] # should be identical for all keys
+        
+        u = uniformizer(parameters, x_batch_dict, y_batch)
+        for g in groups:
+            ix = np.where(np.array(memberships) == g)
+            z = norm().ppf(u[ix])
+            second_moments[g] += z.T @ z
+            sums[g] += z.sum(axis=0)
+            Ng[g] += len(ix[0])
+
+    result = {}
+    for g in groups:
+        mean = sums[g] / Ng[g]
+        result[g] = second_moments[g] / Ng[g] - np.outer(mean, mean)
+
+    if len(groups) == 1:
+        return list(result.values())[0] 
+    return result 
+
+
+###############################################################################
+## Helpers to prepare and postprocess copula parameters
+###############################################################################
+
+
+def group_indices(grouping_var: str, obs: pd.DataFrame) -> dict:
+    """
+    Returns a dictionary of group indices for each group in the grouping variable.
+    """
+    if grouping_var is None:
+        grouping_var = "_copula_group"
+        if "copula_group" not in obs.columns:
+            obs["_copula_group"] = pd.Categorical(["shared_group"] * len(obs))
+    result = {}
+
+    for group in list(obs[grouping_var].dtype.categories):
+        result[group] = np.where(obs[grouping_var].values == group)[0]
+    return result
+
+
+def clip(u: np.array, min: float = 1e-5, max: float = 1 - 1e-5) -> np.array:
+    u[u < min] = min
+    u[u > max] = max
+    return u
+
+
+def format_copula_parameters(parameters: dict, var_names: list):
+    covariance = parameters["covariance"]
+    if type(covariance) is not dict:
+        covariance = pd.DataFrame(
+            parameters["covariance"], columns=list(var_names), index=list(var_names)
+        )
+    else:
+        for group in covariance.keys():
+            covariance[group] = pd.DataFrame(
+                parameters["covariance"][group],
+                columns=list(var_names),
+                index=list(var_names),
+            )
+    return covariance
+
+
+def strip_dataloader(dataloader, pop=False):
+    return DataLoader(
+        dataset=dataloader.dataset,
+        batch_sampler=dataloader.batch_sampler,
+        collate_fn=stack_collate(pop=pop, groups=False),
+    )
+    

scdesigner/estimators/glm_factory.py

@@ -0,0 +1,75 @@
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+import torch
+
+
+def glm_regression_factory(likelihood, initializer, postprocessor) -> dict:
+    def estimator(
+        dataloader: DataLoader,
+        lr: float = 0.1,
+        epochs: int = 40,
+    ):
+        device = check_device()
+        x, y = next(iter(dataloader))
+        params = initializer(x, y, device)
+        optimizer = torch.optim.Adam([params], lr=lr)
+
+        for epoch in range(epochs):
+            for x_batch, y_batch in (pbar := tqdm(dataloader, desc=f"Epoch {epoch + 1}/{epochs}", leave=False)):
+                optimizer.zero_grad()
+                loss = likelihood(params, x_batch, y_batch)
+                loss.backward()
+                optimizer.step()
+                pbar.set_postfix_str(f"loss: {loss.item()}")
+
+
+        return postprocessor(params, x.shape[1], y.shape[1])
+
+    return estimator
+
+def multiple_formula_regression_factory(likelihood, initializer, postprocessor) -> dict:
+    def estimator(
+        dataloaders: dict[str, DataLoader],
+        lr: float = 0.1,
+        epochs: int = 40,
+    ):  
+        device = check_device()
+        x_dict = {}
+        y_dict = {}
+        for key in dataloaders.keys():
+            x_dict[key], y_dict[key] = next(iter(dataloaders[key]))
+        # check if all ys are the same
+        y_ref = y_dict[list(dataloaders.keys())[0]]
+        for key in dataloaders.keys():
+            if not torch.equal(y_dict[key], y_ref):
+                raise ValueError(f"Ys are not the same for {key}")
+        params = initializer(x_dict, y_ref, device) # x is a dictionary of tensors, y is a tensor
+        optimizer = torch.optim.Adam([params], lr=lr)
+        
+        keys = list(dataloaders.keys())
+        loaders = list(dataloaders.values())
+        
+        for epoch in range(epochs):
+            num_keys = len(keys)
+            for batches in (pbar := tqdm(zip(*loaders), desc=f"Epoch {epoch + 1}/{epochs}", leave=False)):
+                x_batch_dict = {
+                    keys[i]: batches[i][0].to(device) for i in range(num_keys)
+                }
+                y_batch = batches[0][1].to(device)
+                optimizer.zero_grad()
+                loss = likelihood(params, x_batch_dict, y_batch) 
+                loss.backward()
+                optimizer.step()
+                pbar.set_postfix_str(f"loss: {loss.item()}")
+
+        return postprocessor(params, x_dict, y_ref)
+
+    return estimator
+
+
+def check_device():
+    return torch.device(
+        "cuda"
+        if torch.cuda.is_available()
+        else "mps" if torch.backends.mps.is_available() else "cpu"
+    )