scdesigner 0.0.1__py3-none-any.whl

scdesigner/format/print.py

@@ -0,0 +1,30 @@
+import rich
+import rich.table
+
+
+def shorten_names(features, max_features=5):
+    if len(features) > max_features:
+        features = features[: (int(max_features - 1))] + ["..."] + features[-1:]
+    return f"""[{', '.join(features)}]"""
+
+
+def print_simulator(margins, copula):
+    table = rich.table.Table(
+        title="[bold magenta]Simulation Plan[/bold magenta]", title_justify="left"
+    )
+    table.add_column("formula")
+    table.add_column("distribution")
+    table.add_column("features")
+
+    i = 1
+    for m in margins:
+        features, margin = m
+        tup = tuple(margin.to_df().iloc[0, :]) + (shorten_names(features),)
+        table.add_row(*tup)
+        i += 1
+
+    rich.print(table)
+    if copula is None:
+        rich.print("Marginal models without copula.")
+
+    return ""

scdesigner/minimal/__init__.py

@@ -0,0 +1,17 @@
+from .scd3_instances import (
+    BernoulliCopula,
+    GaussianCopula,
+    NegBinCopula,
+    ZeroInflatedNegBinCopula
+)
+from .composite import CompositeCopula
+from .positive_nonnegative_matrix_factorization import PositiveNMF
+
+__all__ = [
+    "BernoulliCopula",
+    "CompositeCopula",
+    "GaussianCopula",
+    "NegBinCopula",
+    "PositiveNMF",
+    "ZeroInflatedNegBinCopula"
+]

scdesigner/minimal/bernoulli.py

@@ -0,0 +1,61 @@
+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_fns = {"mean": 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_fns,
+            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)
+        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]):
+        """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):
+        """Return uniformized pseudo-observations for counts y given covariates x."""
+        theta, y = self._local_params(x, y)
+        u1 =  bernoulli(theta).cdf(y)
+        u2 = np.where(y > 0,  bernoulli(theta).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):
+        params = self.predict(x)
+        theta = params.get('mean')
+        if y is None:
+            return _to_numpy(theta)
+        return _to_numpy(theta, y)

scdesigner/minimal/composite.py

@@ -0,0 +1,119 @@
+from .loader import obs_loader
+from .scd3 import SCD3Simulator
+from .standard_covariance import StandardCovariance
+from anndata import AnnData
+from typing import Dict, Optional, List
+import numpy as np
+import torch
+
+class CompositeCopula(SCD3Simulator):
+    def __init__(self, marginals: List,
+                 copula_formula: Optional[str] = None) -> None:
+        self.marginals = marginals
+        self.copula = StandardCovariance(copula_formula)
+        self.template = None
+        self.parameters = None
+        self.merged_formula = None
+
+    def fit(
+        self,
+        adata: AnnData,
+        **kwargs):
+        """Fit the simulator"""
+        self.template = adata
+        merged_formula = {}
+
+        # fit each marginal model
+        for m in range(len(self.marginals)):
+            self.marginals[m][1].setup_data(adata[:, self.marginals[m][0]], **kwargs)
+            self.marginals[m][1].setup_optimizer(**kwargs)
+            self.marginals[m][1].fit(**kwargs)
+
+            # prepare formula for copula loader
+            f = self.marginals[m][1].formula
+            prefixed_f = {f"group{m}_{k}": v for k, v in f.items()}
+            merged_formula = merged_formula | prefixed_f
+
+        # copula simulator
+        self.merged_formula = merged_formula
+        self.copula.setup_data(adata, merged_formula, **kwargs)
+        self.copula.fit(self.merged_uniformize, **kwargs)
+        self.parameters = {
+            "marginal": [m[1].parameters for m in self.marginals],
+            "copula": self.copula.parameters
+        }
+
+    def merged_uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """Produce a merged uniformized matrix for all marginals.
+
+        Delegates to each marginal's `uniformize` method and places the
+        result into the columns of a full matrix according to the variable
+        selection given in `self.marginals[m][0]`.
+        """
+        y_np = y.detach().cpu().numpy()
+        u = np.empty_like(y_np, dtype=float)
+
+        for m in range(len(self.marginals)):
+            sel = self.marginals[m][0]
+            ix = _var_indices(sel, self.template)
+
+            # remove the `group{m}_` prefix we used to distinguish the marginals
+            prefix = f"group{m}_"
+            cur_x = {k.removeprefix(prefix): v if k.startswith(prefix) else v for k, v in x.items()}
+
+            # slice the subset of y for this marginal and call its uniformize
+            y_sub = torch.from_numpy(y_np[:, ix])
+            u[:, ix] = self.marginals[m][1].uniformize(y_sub, cur_x)
+        return torch.from_numpy(u)
+
+    def predict(self, obs=None, batch_size: int = 1000, **kwargs):
+        """Predict from an obs dataframe"""
+        # prepare an internal data loader for this obs
+        if obs is None:
+            obs = self.template.obs
+        loader = obs_loader(
+            obs,
+            self.merged_formula,
+            batch_size=batch_size,
+            **kwargs
+        )
+
+        # prepare per-marginal collectors
+        n_marginals = len(self.marginals)
+        local_pred = [[] for _ in range(n_marginals)]
+
+        # for each batch, call each marginal's predict on its subset of x
+        for _, x_dict in loader:
+            for m in range(n_marginals):
+                prefix = f"group{m}_"
+                # build cur_x where prefixed keys are unprefixed for the marginal
+                cur_x = {k.removeprefix(prefix): v for k, v in x_dict.items()}
+                params = self.marginals[m][1].predict(cur_x)
+                local_pred[m].append(params)
+
+        # merge batch-wise parameter dicts for each marginal and return
+        results = []
+        for m in range(n_marginals):
+            parts = local_pred[m]
+            keys = list(parts[0].keys())
+            results.append({k: torch.cat([d[k] for d in parts]).detach().cpu().numpy() for k in keys})
+
+        return results
+
+
+def _var_indices(sel, adata: AnnData) -> np.ndarray:
+    """Return integer indices of `sel` within `adata.var_names`.
+
+    Expected use: `sel` is a list (or tuple) of variable names (strings).
+    """
+    # If sel is a single string, make it a list so we return consistent shape
+    single_string = False
+    if isinstance(sel, str):
+        sel = [sel]
+        single_string = True
+
+    idx = np.asarray(adata.var_names.get_indexer(sel), dtype=int)
+    if (idx < 0).any():
+        missing = [s for s, i in zip(sel, idx) if i < 0]
+        raise KeyError(f"Variables not found in adata.var_names: {missing}")
+    return idx if not single_string else idx.reshape(-1)

scdesigner/minimal/copula.py

@@ -0,0 +1,33 @@
+from typing import Dict, Callable, Tuple
+import torch
+from anndata import AnnData
+from .loader import adata_loader
+
+class Copula:
+    def __init__(self, formula: str, **kwargs):
+        self.formula = formula
+        self.loader = None
+        self.n_outcomes = None
+        self.parameters = None
+
+    def setup_data(self, adata: AnnData, marginal_formula: Dict[str, str], batch_size: int = 1024, **kwargs):
+        self.adata = adata
+        self.formula = self.formula | marginal_formula
+        self.loader = adata_loader(adata, self.formula, batch_size=batch_size, **kwargs)
+        X_batch, _ = next(iter(self.loader))
+        self.n_outcomes = X_batch.shape[1]
+
+    def fit(self, uniformizer: Callable, **kwargs):
+        raise NotImplementedError
+
+    def pseudo_obs(self, x_dict: Dict):
+        raise NotImplementedError
+
+    def likelihood(self, uniformizer: Callable, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
+        raise NotImplementedError
+
+    def num_params(self, **kwargs):
+        raise NotImplementedError
+
+    def format_parameters(self):
+        raise NotImplementedError

scdesigner/minimal/formula.py

@@ -0,0 +1,23 @@
+from typing import Union
+import warnings
+
+def standardize_formula(formula: Union[str, dict], allowed_keys = None):
+    # The first element of allowed_keys should be the name of default parameter
+    if allowed_keys is None:
+        raise ValueError("Internal error: allowed_keys must be specified")
+    formula = {allowed_keys[0]: formula} if isinstance(formula, str) else formula
+
+    formula_keys = set(formula.keys())
+    allowed_keys = set(allowed_keys)
+
+    if not formula_keys & allowed_keys:
+        raise ValueError(f"formula must have at least one of the following keys: {allowed_keys}")
+    
+    if extra_keys := formula_keys - allowed_keys:
+        warnings.warn(
+            f"Invalid formulas in dictionary will not be used: {extra_keys}",
+            UserWarning,
+        )
+    
+    formula.update({k: '~ 1' for k in allowed_keys - formula_keys})
+    return formula

scdesigner/minimal/gaussian.py

@@ -0,0 +1,65 @@
+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 norm
+
+class Gaussian(Marginal):
+    """Gaussian marginal estimator"""
+    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,
+    ):
+        if self.loader is None:
+            raise RuntimeError("self.loader is not set (call setup_data first)")
+
+        nll = lambda batch: -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):
+        """Compute the negative 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]):
+        """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):
+        """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):
+        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/minimal/kwargs.py

@@ -0,0 +1,24 @@
+def _filter_kwargs(kwargs: dict, allowed: set) -> dict:
+    """Return a new dict containing only keys from kwargs that are in allowed."""
+    return {k: v for k, v in kwargs.items() if k in allowed}
+
+
+# allowed kwargs for common components
+DEFAULT_ALLOWED_KWARGS = {
+    'trainer': {
+        'max_epochs', 'gpus', 'devices', 'accelerator', 'precision',
+        'logger', 'callbacks', 'strategy', 'num_nodes', 'limit_train_batches',
+        'log_every_n_steps', 'accumulate_grad_batches'
+    },
+    'data': {
+        'chunk_size', 'batch_size', 'shuffle', 'num_workers'
+    },
+    'optimizer': {
+        'lr', 'learning_rate', 'momentum', 'weight_decay', 'eps', 'betas',
+        'amsgrad', 'dampening', 'nesterov', 'alpha',
+        'T_max', 'eta_min', 'step_size', 'gamma', 'milestones', 'last_epoch',
+        'verbose', 'patience', 'threshold', 'cooldown',
+        'optimizer_class', 'optimizer', 'scheduler_class', 'scheduler',
+        'monitor', 'interval', 'frequency'
+    }
+}

scdesigner/minimal/loader.py

@@ -0,0 +1,166 @@
+from .kwargs import DEFAULT_ALLOWED_KWARGS, _filter_kwargs
+from anndata import AnnData
+from formulaic import model_matrix
+from torch.utils.data import Dataset, DataLoader
+from typing import Dict
+import numpy as np
+import pandas as pd
+import torch
+
+class AnnDataDataset(Dataset):
+    """Simple PyTorch Dataset for AnnData objects.
+
+    Supports optional chunked loading for backed AnnData objects. When
+    `chunk_size` is provided, the dataset will load contiguous slices
+    of rows (of size `chunk_size`) into memory once and serve individual
+    rows from that cached chunk. This avoids calling `to_memory()` on
+    a per-row basis which is expensive for large backed files.
+    """
+    def __init__(self, adata: AnnData, formula: Dict[str, str], chunk_size: int):
+        self.adata = adata
+        self.formula = formula
+        self.chunk_size = chunk_size
+
+        # keeping track of covariate-related information
+        self.obs_levels = categories(self.adata.obs)
+        self.obs_matrices = {}
+        self.predictor_names = None
+
+        # Internal cache for the currently loaded chunk
+        self._chunk: AnnData | None = None
+        self._chunk_start = 0
+
+    def __len__(self):
+        return len(self.adata)
+
+    def __getitem__(self, idx):
+        """Returns (X, obs) for the given index.
+
+        If `chunk_size` was specified the dataset will load a chunk
+        containing `idx` into memory (if not already cached) and
+        index into that chunk.
+        """
+        self._ensure_chunk_loaded(idx)
+        local_idx = idx - self._chunk_start
+        adata_slice = self._chunk[local_idx]
+
+        # Get X data, accounting for potential sparse matrices
+        X = adata_slice.X
+        if hasattr(X, 'toarray'):
+            X = X.toarray()
+
+        # Get obs data
+        obs_dict = {}
+        for key in self.formula.keys():
+            mat = self.obs_matrices.get(key)
+            obs_dict[key] = to_tensor(mat.values[local_idx: local_idx + 1])
+        return to_tensor(X), obs_dict
+
+    def _ensure_chunk_loaded(self, idx: int) -> None:
+        """Load the chunk that contains `idx` into the internal cache."""
+        start = (idx // self.chunk_size) * self.chunk_size
+        end = min(start + self.chunk_size, len(self.adata))
+
+        if (self._chunk is None) or not (self._chunk_start <= idx < self._chunk_start + len(self._chunk)):
+            # load the next chunk into memory
+            chunk = self.adata[start:end]
+            if getattr(chunk, 'isbacked', False):
+                chunk = chunk.to_memory()
+            self._chunk = chunk
+            self._chunk_start = start
+
+            # Compute model matrices for this chunk's `obs` so we don't need
+            # to keep the full obs data model matrices in memory.
+            obs_coded_chunk = code_levels(self._chunk.obs.copy(), self.obs_levels)
+            self.obs_matrices = {}
+            for key, f in self.formula.items():
+                self.obs_matrices[key] = model_matrix(f, obs_coded_chunk)
+
+            # Capture predictor (column) names from the model matrices once.
+            if self.predictor_names is None:
+                self.predictor_names = {k: list(v.columns) for k, v in self.obs_matrices.items()}
+
+
+def adata_loader(adata: AnnData,
+                 formula: Dict[str, str],
+                 chunk_size: int = None,
+                 batch_size: int = 1024,
+                 shuffle: bool = False,
+                 num_workers: int = 0,
+                 **kwargs) -> DataLoader:
+    """
+    Create a DataLoader from AnnData that returns batches of (X, obs).
+    """
+    data_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['data'])
+    if chunk_size is None:
+        if getattr(adata, 'isbacked', False):
+            chunk_size = 5000
+        else:
+            chunk_size = len(adata)
+
+    dataset = AnnDataDataset(adata, formula, chunk_size)
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        collate_fn=dict_collate_fn,
+        **data_kwargs
+    )
+
+def obs_loader(obs: pd.DataFrame, marginal_formula, **kwargs):
+        adata = AnnData(X=np.zeros((len(obs), 1)), obs=obs)
+        return adata_loader(
+            adata,
+            marginal_formula,
+            **kwargs
+        )
+
+################################################################################
+## Helper functions
+################################################################################
+
+def dict_collate_fn(batch):
+    """
+    Custom collate function for handling dictionary obs tensors.
+    """
+    X_batch = torch.stack([item[0] for item in batch])
+    obs_batch = [item[1] for item in batch]
+
+    obs_dict = {}
+    for key in obs_batch[0].keys():
+        obs_dict[key] = torch.stack([obs[key] for obs in obs_batch])
+    return X_batch, obs_dict
+
+def to_tensor(X):
+    # If the tensor is 2D with second dim == 1, squeeze only the first
+    # dim when appropriate (e.g. converting a single-row X to 1D samples)
+    t = torch.tensor(X, dtype=torch.float32)
+    if t.dim() == 2 and t.size(1) == 1:
+        if t.size(0) == 1:
+            return t.view(1)
+        return t
+    return t.squeeze()
+
+def categories(obs):
+    levels = {}
+    for k in obs.columns:
+        obs_type = str(obs[k].dtype)
+        if obs_type in ["category", "object"]:
+            levels[k] = obs[k].unique()
+    return levels
+
+
+def code_levels(obs, categories):
+    for k in obs.columns:
+        if str(obs[k].dtype) == "category":
+            obs[k] = obs[k].astype(pd.CategoricalDtype(categories[k]))
+    return obs
+
+###############################################################################
+## Misc. Helper functions
+###############################################################################
+
+def _to_numpy(*tensors):
+    """Convenience helper: detach, move to CPU, and convert tensors to numpy arrays."""
+    return tuple(t.detach().cpu().numpy() for t in tensors)

scdesigner/minimal/marginal.py

@@ -0,0 +1,140 @@
+from .kwargs import DEFAULT_ALLOWED_KWARGS, _filter_kwargs
+from .loader import adata_loader
+from anndata import AnnData
+from typing import Union, Dict, Optional, Tuple
+import pandas as pd
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+
+
+class Marginal:
+    def __init__(self, formula: Union[Dict, str]):
+        self.formula = formula
+        self.feature_dims = None
+        self.loader = None
+        self.n_outcomes = None
+        self.predict = None
+        self.predictor_names = None
+        self.parameters = None
+
+    def setup_data(self, adata: AnnData, batch_size: int = 1024, **kwargs):
+        """Set up the dataloader for the AnnData object."""
+        # keep a reference to the AnnData for later use (e.g., var_names)
+        self.adata = adata
+        self.loader = adata_loader(adata, self.formula, batch_size=batch_size, **kwargs)
+        X_batch, obs_batch = next(iter(self.loader))
+        self.n_outcomes = X_batch.shape[1]
+        self.feature_dims = {k: v.shape[1] for k, v in obs_batch.items()}
+        self.predictor_names = self.loader.dataset.predictor_names
+
+    def fit(self, **kwargs):
+        """Fit the marginal predictor"""
+        if self.predict is None:
+            self.setup_optimizer(**kwargs)
+        trainer_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['trainer'])
+        trainer = pl.Trainer(**trainer_kwargs)
+        trainer.fit(self.predict, train_dataloaders=self.loader)
+        self.parameters = self.format_parameters()
+
+    def setup_optimizer(self, **kwargs):
+        raise NotImplementedError
+
+    def likelihood(self, batch: Tuple[torch.Tensor, Dict[str, torch.Tensor]]):
+        """Compute the (negative) log-likelihood or loss for a batch.
+        """
+        raise NotImplementedError
+
+    def invert(self, u: torch.Tensor, x: Dict[str, torch.Tensor]):
+        """Invert pseudoobservations."""
+        raise NotImplementedError
+
+    def uniformize(self, y: torch.Tensor, x: Dict[str, torch.Tensor]):
+        """Uniformize using learned CDF.
+        """
+        raise NotImplementedError
+
+    def format_parameters(self):
+        """Convert fitted coefficient tensors into pandas DataFrames.
+
+        Returns:
+            dict: mapping from parameter name -> pandas.DataFrame with rows
+                corresponding to predictor column names (from
+                `self.predictor_names[param]`) and columns corresponding to
+                `self.adata.var_names` (gene names). The values are moved to
+                CPU and converted to numpy floats.
+        """
+        var_names = list(self.adata.var_names)
+
+        dfs = {}
+        for param, tensor in self.predict.coefs.items():
+            coef_np = tensor.detach().cpu().numpy()
+            row_names = list(self.predictor_names[param])
+            dfs[param] = pd.DataFrame(coef_np, index=row_names, columns=var_names)
+        return dfs
+
+    def num_params(self):
+        """Return the number of parameters."""
+        if self.predict is None:
+            return 0
+        return sum(p.numel() for p in self.predict.parameters() if p.requires_grad)
+
+
+class GLMPredictor(pl.LightningModule):
+    """GLM-style predictor with arbitrary named parameters.
+
+    Args:
+        n_outcomes: number of model outputs (e.g. genes)
+        feature_dims: mapping from param name -> number of covariate features
+        link_fns: optional mapping from param name -> callable(link) applied to linear predictor
+
+    The module will create one coefficient matrix per named parameter with shape
+    (n_features_for_param, n_outcomes) and expose them as Parameters under
+    `self.coefs[param_name]`.
+    """
+    def __init__(
+        self,
+        n_outcomes: int,
+        feature_dims: Dict[str, int],
+        link_fns: Dict[str, callable] = None,
+        loss_fn: Optional[callable] = None,
+        optimizer_class: Optional[callable] = torch.optim.Adam,
+        optimizer_kwargs: Optional[Dict] = None,
+    ):
+        super().__init__()
+        self.n_outcomes = int(n_outcomes)
+        self.feature_dims = dict(feature_dims)
+        self.param_names = list(self.feature_dims.keys())
+
+        # create default link functions and parameter matrices
+        self.link_fns = link_fns or {k: torch.exp for k in self.param_names}
+        self.coefs = nn.ParameterDict()
+        for key, dim in self.feature_dims.items():
+            self.coefs[key] = nn.Parameter(torch.zeros(dim, self.n_outcomes))
+
+        # optimization parameters
+        self.reset_parameters()
+        self.loss_fn = loss_fn
+        self.optimizer_class = optimizer_class
+        self.optimizer_kwargs = optimizer_kwargs
+
+    def reset_parameters(self):
+        for p in self.coefs.values():
+            nn.init.normal_(p, mean=0.0, std=1e-2)
+
+    def forward(self, obs_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        out = {}
+        for name in self.param_names:
+            x_beta = obs_dict[name] @ self.coefs[name]
+            link = self.link_fns.get(name, torch.exp)
+            out[name] = link(x_beta)
+        return out
+
+    def training_step(self, batch):
+        loss = self.loss_fn(batch)
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self, **kwargs):
+        optimizer_kwargs = _filter_kwargs(kwargs, DEFAULT_ALLOWED_KWARGS['optimizer'])
+        return self.optimizer_class(self.parameters(), **optimizer_kwargs)