# Copyright Contributors to the Cellarium project.
# SPDX-License-Identifier: BSD-3-Clause
import warnings
from typing import TypedDict
import lightning.pytorch as pl
import numpy as np
import pandas as pd
import torch
import torch.nn.functional
from cellarium.ml.models.model import CellariumModel, PredictMixin, ValidateMixin
from cellarium.ml.utilities.testing import (
assert_arrays_equal,
assert_columns_and_array_lengths_equal,
)
class NonleafInfo(TypedDict):
nonleaf_desc_cc: torch.Tensor
perm: torch.Tensor
inv_perm: torch.Tensor
def _expand_with_ancestors(
cl_name_subset: list[str],
cl_names: list[str],
descendant_tensor: torch.Tensor,
) -> list[str]:
"""Return ``cl_name_subset`` extended to include every ancestor of each node.
``descendant_tensor[k, j] == 1`` means j is a descendant of k, equivalently k
is an ancestor of j. So the ancestors of j are all k where column j is 1.
Args:
cl_name_subset: Category names requested by the user.
cl_names: Full ordered list of category names (rows/cols of ``descendant_tensor``).
descendant_tensor: Binary ``(C, C)`` tensor on any device.
Returns:
Sorted list containing every name in ``cl_name_subset`` plus all of their
ancestors that are present in ``cl_names``.
"""
index_map = {cat: i for i, cat in enumerate(cl_names)}
expanded: set[int] = set()
for name in cl_name_subset:
j = index_map[name]
# column j: entry k==1 means k is an ancestor of j (including j itself)
ancestor_indices = (descendant_tensor[:, j] > 0).nonzero(as_tuple=True)[0].tolist()
expanded.update(ancestor_indices)
return sorted(cl_names[i] for i in expanded)
def _build_nonleaf_info(desc_matrix_cc: torch.Tensor) -> NonleafInfo:
"""Precompute the non-leaf descriptor tensors needed by ``propagate_logits``.
Produces a permutation that places non-leaf categories first and leaves last,
so the hot path can use ``torch.cat`` and simple slices instead of boolean-mask
assignment. Boolean-mask assignment internally calls ``aten.nonzero`` which
``torch.compile`` / inductor cannot fuse; integer-index gather and ``torch.cat``
have no such restriction.
Runs once per unique ``desc_matrix_cc`` (construction time, not the hot path).
Args:
desc_matrix_cc: Binary ``(c, c)`` descendant tensor.
Returns:
A :class:`NonleafInfo` dict with keys ``nonleaf_desc_cc``, ``perm``, and
``inv_perm``. ``nonleaf_desc_cc`` is ``(c_nonleaf, c)`` with columns in
``perm`` order; ``perm`` places non-leaf indices first, leaf last;
``inv_perm`` restores the original column order.
"""
nonleaf_mask = desc_matrix_cc.sum(dim=1) > 1
nonleaf_indices = nonleaf_mask.nonzero(as_tuple=True)[0] # (c_nonleaf,)
leaf_indices = (~nonleaf_mask).nonzero(as_tuple=True)[0] # (c_leaf,)
perm = torch.cat([nonleaf_indices, leaf_indices]) # non-leaf first
inv_perm = torch.argsort(perm)
# Reorder desc columns to match the permuted input order used in propagate_logits
nonleaf_desc_cc = desc_matrix_cc[nonleaf_indices][:, perm] # (c_nonleaf, c)
return {"nonleaf_desc_cc": nonleaf_desc_cc, "perm": perm, "inv_perm": inv_perm}
def _propagate_probs_impl(probs_nc: torch.Tensor, descendant_tensor_cc: torch.Tensor) -> torch.Tensor:
"""
Propagate probabilities up the hierarchy defined by ``descendant_tensor_cc`` using matrix multiplication.
This effectively sums the probabilities of all descendant categories for each category.
The output is then clamped to a maximum of 1.0 to ensure valid probability values.
Args:
probs_nc: Tensor of shape (n, c) containing the probabilities for each category.
descendant_tensor_cc: Binary tensor of shape (c, c) defining descendant relationships.
Returns:
Tensor of shape (n, c) containing the propagated probabilities for each category
"""
propagated_probs_nc = torch.einsum(
"nc,kc->nk",
probs_nc,
descendant_tensor_cc,
)
return torch.clamp(propagated_probs_nc, max=1.0)
# Compiled wrapper (used by SOCAM). The plain ``_propagate_probs_impl`` is importable by callers
# (e.g. SCANVI) that want to avoid recompilation for variable-shape inputs.
propagate_probs = torch.compile(_propagate_probs_impl)
def _logsumexp_propagated(logits_nc: torch.Tensor, desc_matrix_cc: torch.Tensor) -> torch.Tensor:
c_nonleaf = desc_matrix_cc.shape[0]
expanded_logits = logits_nc.unsqueeze(dim=-1).expand(-1, -1, c_nonleaf)
temp = expanded_logits.masked_fill(desc_matrix_cc.T == 0, float("-inf"))
return temp.logsumexp(dim=1)
def _propagate_logits_impl(
logits_nc: torch.Tensor,
nonleaf_desc_cc: torch.Tensor,
perm: torch.Tensor,
inv_perm: torch.Tensor,
) -> torch.Tensor:
"""
Perform probability propagation in logit space.
Non-leaf output categories reduce over all their descendants via
``_logsumexp_propagated`` using a ``(c_nonleaf, c)`` submatrix, so the
intermediate tensor is ``(n, c, c_nonleaf)`` rather than ``(n, c, c)``.
Leaf output categories are the identity (logsumexp of a single element).
``perm`` / ``inv_perm`` sort columns so non-leaf outputs come first,
allowing assembly via ``torch.cat`` and a single integer-index gather —
avoiding ``aten.nonzero`` which breaks ``torch.compile``.
Args:
logits_nc: ``(n, c)`` raw logit tensor.
nonleaf_desc_cc: ``(c_nonleaf, c)`` descendant rows for non-leaf outputs,
with columns in ``perm`` order (from ``_build_nonleaf_info``).
perm: ``(c,)`` permutation — non-leaf indices first, leaf last
(from ``_build_nonleaf_info``).
inv_perm: ``(c,)`` inverse of ``perm`` (from ``_build_nonleaf_info``).
Returns:
``(n, c)`` propagated log-probability tensor in original column order.
"""
c_nonleaf = nonleaf_desc_cc.shape[0]
logits_reordered = logits_nc[:, perm] # (n, c): non-leaf first
nonleaf_part = _logsumexp_propagated(logits_reordered, nonleaf_desc_cc) # (n, c_nonleaf)
leaf_part = logits_reordered[:, c_nonleaf:] # (n, c_leaf)
out = torch.cat([nonleaf_part, leaf_part], dim=1)[:, inv_perm] # (n, c) original order
return out - torch.logsumexp(logits_nc, dim=1, keepdim=True)
# Compiled wrapper (used by SOCAM). The plain ``_propagate_logits_impl`` is importable by callers
# (e.g. SCANVI) that want an eager version to avoid recompilation for variable-shape inputs.
propagate_logits = torch.compile(_propagate_logits_impl)
def compute_class_weights(
active_cl_names: list[str],
class_counts: "pd.Series | None",
active_descendant_tensor_cc: torch.Tensor,
propagate_class_counts: bool = False,
normalize: str = "class_mean",
) -> torch.Tensor | None:
"""Compute per-class cross-entropy weights from training cell counts.
Classes absent from ``class_counts`` (typically pure-ancestor nodes with no direct cell
labels) and classes with a count of zero are treated as unlabeled and receive a neutral
weight of 1.0. Inverse-frequency weights are computed over the nonzero-count classes only.
Args:
active_cl_names: Ordered list of active category names (defines the output order).
class_counts: Optional pandas Series mapping class names to training cell counts.
When ``None`` this returns ``None`` (no weighting). Extra entries not in
``active_cl_names`` are ignored. Negative counts raise a ``ValueError``.
active_descendant_tensor_cc: ``(c, c)`` binary descendant submatrix over the active
categories (diagonal included). Used only when ``propagate_class_counts`` is True.
propagate_class_counts: If True, propagate raw counts up the ontology
(``active_descendant_tensor_cc @ counts``) so each node's effective count becomes
its own count plus the sum of all descendant counts.
normalize: Normalization convention for the inverse-frequency weights.
``"class_mean"`` (SOCAM, self-normalizing ``reduction="mean"`` CE) divides by the
unweighted class mean. ``"data_mean"`` (SCANVI, ``reduction="none"`` CE) leaves the
weights with a data-frequency-weighted mean of 1 so an outer scalar weight keeps
its meaning.
Returns:
A float tensor of length ``len(active_cl_names)``, or ``None`` if ``class_counts`` is None.
"""
if class_counts is None:
return None
provided_counts = {c: class_counts[c] for c in active_cl_names if c in class_counts.index}
if any(v < 0 for v in provided_counts.values()):
raise ValueError("All class_counts values must be >= 0.")
# Pin to CPU so this works inside a torch.device("meta") construction context.
counts = torch.tensor(
[float(provided_counts.get(c, 0.0)) for c in active_cl_names], dtype=torch.float, device="cpu"
)
if propagate_class_counts:
counts = active_descendant_tensor_cc.cpu() @ counts
nonzero = counts > 0
weights = torch.ones(len(active_cl_names), dtype=torch.float, device="cpu")
if nonzero.any():
total = counts[nonzero].sum()
n_nonzero = nonzero.sum().float()
raw = total / (n_nonzero * counts[nonzero]) # data-frequency-weighted mean of 1
if normalize == "class_mean":
weights[nonzero] = raw / raw.mean()
elif normalize == "data_mean":
weights[nonzero] = raw
else:
raise ValueError(f"normalize must be 'class_mean' or 'data_mean', got {normalize!r}.")
return weights
[docs]
class SOCAM(CellariumModel, PredictMixin, ValidateMixin):
"""
Logistic regression model for cell type ontology classification.
Args:
n_obs: Number of observations in the dataset (used to scale the cross-entropy loss).
var_names_g: The variable-name schema for the input data; used for validation.
output_categories: Total number of target categories expected at prediction/validation time.
Used when the trained model has fewer categories than the final output space.
descendant_tensor: Binary (0/1) tensor of shape (n_categories, n_categories) defining
the descendant relationships between categories. Row i contains ones
for all categories considered descendants of category i (plus self). Used for
probability-propagation.
cl_names: Full list of category identifiers matching the rows/columns of ``descendant_tensor``.
cl_name_subset: Optional list of category names (from ``cl_names``) to restrict
training and prediction to. The list is sorted internally so order does not matter.
When ``None``, all categories are used.
probability_propagation_flag: If True, applies hierarchical probability propagation
before predicting the output distribution.
W_prior_scale: Scale (b) parameter of the Laplace prior on the weight matrix `W_gc`.
W_init_scale: Standard deviation for initializing `W_gc`.
seed: Random seed used to initialize parameters.
log_metrics: If True, logs weight histograms (TensorBoard) during training.
If True, logs weight histograms (TensorBoard) during training.
class_counts: Optional pandas Series mapping class names to cell counts in the
training data. Classes absent from the Series (typically pure-ancestor nodes
with no direct cell labels) and classes with a count of zero are both treated
as unlabeled and receive a neutral weight of 1.0, which has no effect on the
loss since no cells carry those labels. Inverse-frequency weights are computed
and normalized to mean 1 over the nonzero-count classes only. Negative counts
raise a ``ValueError``. Extra Series entries not in the active category set are
ignored. When ``None``, all classes are weighted equally.
propagate_class_counts: If ``True`` and ``class_counts`` is provided, propagate
raw counts up the ontology before computing weights. Each node's effective
count becomes its own direct count plus the sum of all descendant counts. This
down-weights ancestor nodes (they accumulate large effective counts) relative
to rare leaf nodes. A pure-ancestor node with zero direct labels but nonzero
descendant counts will receive a real inverse-frequency weight rather than the
neutral 1.0 it would otherwise get. Has no effect when ``class_counts`` is
``None``.
"""
def __init__(
self,
n_obs: int,
var_names_g: np.ndarray,
descendant_tensor: torch.Tensor,
cl_names: list[str],
cl_name_subset: list[str] | None = None,
probability_propagation_flag: bool = True,
W_prior_scale: float = 1e-2,
W_init_scale: float = 1.0,
seed: int = 0,
log_metrics: bool = True,
include_ancestors_of_cl_name_subset: bool = True,
class_counts: pd.Series | None = None,
propagate_class_counts: bool = False,
) -> None:
super().__init__()
self.n_obs = n_obs
self.var_names_g = var_names_g
self.n_vars = len(var_names_g)
self.cl_names = cl_names
descendant_tensor = descendant_tensor.float()
if descendant_tensor.shape[0] != descendant_tensor.shape[1]:
raise ValueError("`descendant_tensor` should be a square matrix.")
if descendant_tensor.trace() != descendant_tensor.shape[0]:
raise ValueError(
"`descendant_tensor` should have ones on the diagonal (each category is a descendant of itself)."
)
if len(cl_names) != descendant_tensor.shape[0]:
raise ValueError("Length of `cl_names` should match the number of rows in `descendant_tensor`.")
self._descendant_tensor = descendant_tensor
self.register_buffer("descendant_tensor", descendant_tensor)
self.n_categories = descendant_tensor.shape[0]
self.include_ancestors_of_cl_name_subset = include_ancestors_of_cl_name_subset
if include_ancestors_of_cl_name_subset and cl_name_subset is not None:
cl_name_subset = _expand_with_ancestors(cl_name_subset, cl_names, descendant_tensor)
self.cl_name_subset = cl_name_subset
self.probability_propagation_flag = probability_propagation_flag
self.propagate_class_counts = propagate_class_counts
self.seed = seed
self.log_metrics = log_metrics
# Build active category set (subset if given, else full list)
active_cl_names: list[str] = cl_name_subset if cl_name_subset is not None else list(cl_names)
self.active_cl_names = active_cl_names
self.n_active_cats = len(active_cl_names)
self.label_lookup: dict[str, int] = {name: i for i, name in enumerate(active_cl_names)}
# Precompute the active-category submatrix of the descendant tensor.
# Use the raw (pre-registration) descendant_tensor so this works even when
# __init__ is called inside a meta-device context (register_buffer would
# create a meta tensor, but descendant_tensor here is still a real CPU tensor).
index_map_init = {cat: i for i, cat in enumerate(cl_names)}
active_indices = [index_map_init[cat] for cat in active_cl_names]
ix = torch.tensor(active_indices, dtype=torch.long)
active_descendant_tensor_cc = descendant_tensor[ix][:, ix]
nonleaf_info = _build_nonleaf_info(active_descendant_tensor_cc)
# Store CPU copies so reset_parameters() can repopulate after meta→real copy.
self._active_descendant_tensor_cc = active_descendant_tensor_cc
self._nonleaf_desc_cc = nonleaf_info["nonleaf_desc_cc"]
self._perm = nonleaf_info["perm"]
self._inv_perm = nonleaf_info["inv_perm"]
self.register_buffer("active_descendant_tensor_cc", active_descendant_tensor_cc)
self.register_buffer("nonleaf_desc_cc", nonleaf_info["nonleaf_desc_cc"])
self.register_buffer("perm", nonleaf_info["perm"])
self.register_buffer("inv_perm", nonleaf_info["inv_perm"])
# Trainable parameters — sized to active categories only
self._W_prior_scale = W_prior_scale
self.W_init_scale = W_init_scale
self.W_prior_scale: torch.Tensor
self.register_buffer("W_prior_scale", torch.empty(()))
self.W_gc = torch.nn.Parameter(torch.empty(self.n_vars, self.n_active_cats, dtype=torch.float))
self.b_c = torch.nn.Parameter(torch.empty(self.n_active_cats, dtype=torch.float))
# Class weights for cross-entropy loss
weights = compute_class_weights(
active_cl_names=active_cl_names,
class_counts=class_counts,
active_descendant_tensor_cc=active_descendant_tensor_cc,
propagate_class_counts=propagate_class_counts,
normalize="class_mean",
)
self._class_weights: torch.Tensor | None = weights
self.register_buffer("class_weights", weights.clone() if weights is not None else None)
self.reset_parameters()
def reset_parameters(self) -> None:
rng_device = self.W_gc.device.type if self.W_gc.device.type != "meta" else "cpu"
rng = torch.Generator(device=rng_device)
rng.manual_seed(self.seed)
self.W_prior_scale.fill_(self._W_prior_scale)
self.descendant_tensor.copy_(self._descendant_tensor)
self.active_descendant_tensor_cc.copy_(self._active_descendant_tensor_cc)
self.nonleaf_desc_cc.copy_(self._nonleaf_desc_cc)
self.perm.copy_(self._perm)
self.inv_perm.copy_(self._inv_perm)
if self._class_weights is not None:
self.class_weights.copy_(self._class_weights)
self.W_gc.data.normal_(0, self.W_init_scale, generator=rng)
self.b_c.data.zero_()
def _cl_names_to_indices(self, cl_names_n: np.ndarray) -> torch.Tensor:
"""
Convert a per-cell array of string category names to a 1-D integer tensor of
0-based indices into the active category list.
Args:
cl_names_n:
Array of length n containing category name strings for each cell.
Returns:
Long tensor of shape ``(n,)`` with integer category indices.
Raises:
ValueError: If any label in ``cl_names_n`` is not present in ``self.label_lookup``.
"""
try:
return torch.tensor([self.label_lookup[c] for c in cl_names_n], dtype=torch.long)
except KeyError as exc:
valid = sorted(self.label_lookup.keys())
raise ValueError(f"Label {exc} is not in the active category list. Valid labels are: {valid}") from exc
[docs]
def forward(
self,
x_ng: torch.Tensor,
var_names_g: np.ndarray,
cl_names_n: np.ndarray,
) -> dict[str, torch.Tensor | None]:
"""
Args:
x_ng:
The input data.
var_names_g:
The variable names for the input data.
cl_names_n:
Array of length n containing a category name string (from ``self.cl_names``) for
each cell. When ``self.cl_name_subset`` is set, every label must be a member of
that subset.
Returns:
A dictionary with the loss value.
"""
assert_columns_and_array_lengths_equal("x_ng", x_ng, "var_names_g", var_names_g)
assert_arrays_equal("var_names_g", var_names_g, "self.var_names_g", self.var_names_g)
y_n = self._cl_names_to_indices(cl_names_n).to(x_ng.device)
logits_nc = self._compute_regression(x_ng, self.W_gc, self.b_c)
if self.probability_propagation_flag:
logits_nc = propagate_logits(logits_nc, self.nonleaf_desc_cc, self.perm, self.inv_perm)
ce_loss = torch.nn.functional.cross_entropy(logits_nc, y_n, reduction="mean", weight=self.class_weights)
laplace_loss = (self.W_gc.abs().sum() / self.W_prior_scale) / self.n_obs
loss = ce_loss + laplace_loss
return {"loss": loss}
@torch.compile()
def _compute_regression(self, x_ng: torch.Tensor, W_gc: torch.Tensor, b_c: torch.Tensor) -> torch.Tensor:
return x_ng @ W_gc + b_c
[docs]
def validate(
self,
trainer: pl.Trainer,
pl_module: pl.LightningModule,
batch_idx: int,
x_ng: torch.Tensor,
var_names_g: np.ndarray,
cl_names_n: np.ndarray,
) -> None:
"""
Default validation method for models. This method logs the validation loss to TensorBoard.
Override this method to customize the validation behavior.
"""
output = self(
x_ng=x_ng,
var_names_g=var_names_g,
cl_names_n=cl_names_n,
)
loss = output.get("loss")
if loss is not None:
# Logging to TensorBoard by default
pl_module.log("val_loss", loss, sync_dist=True, on_epoch=True)
[docs]
def predict(
self,
x_ng: torch.Tensor,
var_names_g: np.ndarray,
) -> dict[str, np.ndarray | torch.Tensor]:
"""
Predict the target logits.
Args:
x_ng:
The input data.
var_names_g:
The variable names for the input data.
Returns:
A dictionary with the target logits. Output tensors have shape
``(n, n_active_cats)``.
"""
assert_columns_and_array_lengths_equal("x_ng", x_ng, "var_names_g", var_names_g)
assert_arrays_equal("var_names_g", var_names_g, "self.var_names_g", self.var_names_g)
logits_nc = self._compute_regression(x_ng, self.W_gc, self.b_c)
probs_nc = torch.nn.functional.softmax(logits_nc, dim=1)
if self.probability_propagation_flag:
probs_nc = propagate_probs(probs_nc, self.active_descendant_tensor_cc)
return {"y_logits_nc": logits_nc, "cell_type_probs_nc": probs_nc}
def on_train_epoch_end(self, trainer: pl.Trainer) -> None:
if trainer.global_rank != 0:
return
if not self.log_metrics:
return
for logger in trainer.loggers:
if isinstance(logger, pl.loggers.TensorBoardLogger):
try:
logger.experiment.add_histogram(
"W_gc",
self.W_gc,
global_step=trainer.global_step,
)
except ValueError as e:
warnings.warn(f"Failed to log histogram for W_gc step={trainer.global_step} due to {e}")