Release 0.1.5 (#25)

* Refactor benchmark

* Add CELU for MA

* Update version

* Revert callable for latent mean / var activation

* Update CHANGELOG.md

* [pre-commit.ci] pre-commit autoupdate (#24)

updates:
- [github.com/astral-sh/ruff-pre-commit: v0.11.5 → v0.11.8](astral-sh/ruff-pre-commit@v0.11.5...v0.11.8)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update _benchmark.py (align with ruff)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: moinfar

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
     hooks:
       - id: prettier
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.5
+    rev: v0.11.8
     hooks:
       - id: ruff
         types_or: [python, pyi, jupyter]

CHANGELOG.md

@@ -4,6 +4,15 @@
 
 -
 
+## [0.1.5] - 2025-05-09
+
+- Refactor benchmarking code for better reusability
+- Revert callable for mean and var activation
+
+## [0.1.4] - 2025-04-17
+
+- Limit anndata version for compatibility with old scvi-tools
+
 ## [0.1.3] - 2025-02-12
 
 - Introduce mean activation to make non-negative latents possible (docs will come later)

pyproject.toml

@@ -4,7 +4,7 @@ requires = ["hatchling"]
 
 [project]
 name = "drvi-py"
-version = "0.1.3"
+version = "0.1.5"
 description = "Disentangled Generative Representation of Single Cell Omics"
 readme = "README.md"
 requires-python = ">=3.10,<3.13"

src/drvi/scvi_tools_based/module/_drvi.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from typing import Literal
 
 import numpy as np
@@ -148,8 +148,8 @@ def __init__(
         ] = "pnb_softmax",
         prior: Literal["normal", "gmm_x", "vamp_x"] = "normal",
         prior_init_dataloader: DataLoader | None = None,
-        var_activation: Literal["exp", "pow2"] = "exp",
-        mean_activation: str = "identity",
+        var_activation: Callable | Literal["exp", "pow2", "2sig"] | Callable = "exp",
+        mean_activation: Callable | str = "identity",
         encoder_layer_factory: LayerFactory = None,
         decoder_layer_factory: LayerFactory = None,
         extra_encoder_kwargs: dict | None = None,

src/drvi/scvi_tools_based/nn/_base_components.py

@@ -1,6 +1,6 @@
 import collections
 import math
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from typing import Literal
 
 import torch
@@ -421,8 +421,8 @@ def __init__(
         dropout_rate: float = 0.1,
         distribution: str = "normal",
         var_eps: float = 1e-4,
-        var_activation: Literal["exp", "pow2"] = "exp",
-        mean_activation: str = "identity",
+        var_activation: Callable | Literal["exp", "pow2"] = "exp",
+        mean_activation: Callable | str = "identity",
         layer_factory: LayerFactory = None,
         covariate_modeling_strategy: Literal[
             "one_hot",
@@ -499,8 +499,11 @@ def __init__(
             self.var_activation = torch.exp
         elif var_activation == "pow2":
             self.var_activation = lambda x: torch.pow(x, 2)
+        elif var_activation == "2sig":
+            self.var_activation = lambda x: 2 * torch.sigmoid(x)
         else:
-            raise NotImplementedError()
+            assert callable(var_activation)
+            self.var_activation = var_activation
 
         if mean_activation == "identity":
             self.mean_activation = nn.Identity()
@@ -516,8 +519,14 @@ def __init__(
                 mean_activation = "elu_1.0"
             alpha = float(mean_activation.split("elu_")[1])
             self.mean_activation = nn.ELU(alpha=alpha)
+        elif mean_activation.startswith("celu"):
+            if mean_activation == "celu":
+                mean_activation = "celu_1.0"
+            alpha = float(mean_activation.split("celu_")[1])
+            self.mean_activation = nn.CELU(alpha=alpha)
         else:
-            raise NotImplementedError()
+            assert callable(mean_activation)
+            self.mean_activation = mean_activation
 
     def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor, cont_full_tensor: torch.Tensor = None):
         r"""The forward computation for a single sample.

src/drvi/utils/metrics/__init__.py

@@ -1,4 +1,5 @@
 from ._aggregation import latent_matching_score, most_similar_averaging_score, most_similar_gap_score
+from ._benchmark import DiscreteDisentanglementBenchmark
 from ._pairwise import (
     global_dim_mutual_info_score,
     local_mutual_info_score,
@@ -14,4 +15,5 @@
     "most_similar_averaging_score",
     "latent_matching_score",
     "most_similar_gap_score",
+    "DiscreteDisentanglementBenchmark",
 ]

src/drvi/utils/metrics/_benchmark.py

@@ -0,0 +1,170 @@
+import pickle
+
+import numpy as np
+import pandas as pd
+
+from drvi.utils.metrics._aggregation import latent_matching_score, most_similar_averaging_score, most_similar_gap_score
+from drvi.utils.metrics._pairwise import local_mutual_info_score, nn_alignment_score, spearman_correlataion_score
+
+AVAILABLE_METRICS = {
+    "ASC": spearman_correlataion_score,
+    "SPN": nn_alignment_score,
+    "SMI": local_mutual_info_score,
+}
+
+
+AVAILABLE_AGGREGATION_METHODS = {
+    "LMS": latent_matching_score,
+    "MSAS": most_similar_averaging_score,
+    "MSGS": most_similar_gap_score,
+}
+
+
+class DiscreteDisentanglementBenchmark:
+    version = "v1"
+
+    def __init__(
+        self,
+        embed,
+        discrete_target=None,
+        one_hot_target=None,
+        dim_titles=None,
+        metrics=("SMI", "SPN", "ASC"),
+        aggregation_methods=("LMS", "MSAS", "MSGS"),
+    ):
+        if discrete_target is None and one_hot_target is None:
+            raise ValueError("Either discrete_target or one_hot_target must be provided.")
+        if discrete_target is not None and one_hot_target is not None:
+            raise ValueError("Only one of discrete_target or one_hot_target should be provided.")
+
+        if discrete_target is not None:
+            if isinstance(discrete_target, pd.Series):
+                discrete_target = discrete_target.astype("category")
+            elif isinstance(discrete_target, np.ndarray):
+                discrete_target = pd.Series(discrete_target, dtype="category")
+            else:
+                raise ValueError("discrete_target must be a pandas Series or numpy array")
+            one_hot_target = pd.DataFrame(
+                np.eye(len(discrete_target.cat.categories))[discrete_target.cat.codes],
+                columns=discrete_target.cat.categories,
+            )
+
+        if isinstance(one_hot_target, pd.DataFrame):
+            pass
+        elif isinstance(one_hot_target, np.ndarray):
+            one_hot_target = pd.DataFrame(
+                one_hot_target, columns=[f"process_{i}" for i in range(one_hot_target.shape[1])]
+            )
+        else:
+            raise ValueError("one_hot_target must be a pandas DataFrame or numpy array")
+
+        if dim_titles is None:
+            dim_titles = [f"dim_{d}" for d in range(embed.shape[1])]
+
+        self.embed = embed.copy()
+        self.one_hot_target = one_hot_target.copy()
+        self.dim_titles = dim_titles
+        self.metrics = metrics
+        self.aggregation_methods = aggregation_methods
+
+        self.results = {}
+        self.aggregated_results = {}
+
+    @staticmethod
+    def _compute_metrics(embed, one_hot_target, dim_titles=None, metrics=()):
+        if dim_titles is None:
+            dim_titles = [f"dim_{d}" for d in range(embed.shape[1])]
+
+        results = {}
+        for metric_name in metrics:
+            result_df = pd.DataFrame(
+                AVAILABLE_METRICS[metric_name](embed, gt_one_hot=one_hot_target.values),
+                index=dim_titles,
+                columns=one_hot_target.columns,
+            )
+            results[metric_name] = result_df
+
+        return results
+
+    @staticmethod
+    def _aggregate_metrics(results, aggregation_methods=()):
+        aggregated_results = {}
+        for aggregation_method in aggregation_methods:
+            for metric_name in results:
+                aggregated_results[f"{aggregation_method}-{metric_name}"] = AVAILABLE_AGGREGATION_METHODS[
+                    aggregation_method
+                ](results[metric_name].values)
+        return aggregated_results
+
+    def is_complete(self):
+        for metric in self.metrics:
+            if metric not in self.results:
+                return False
+        for aggregation_method in self.aggregation_methods:
+            for metric in self.metrics:
+                if f"{aggregation_method}-{metric}" not in self.aggregated_results:
+                    return False
+        return True
+
+    def evaluate(self):
+        if not self.is_complete():
+            remaining_metrics = [metric for metric in self.metrics if metric not in self.results]
+            self.results = {
+                **self.results,
+                **self._compute_metrics(self.embed, self.one_hot_target, self.dim_titles, remaining_metrics),
+            }
+            # Aggregation is cheap. Do it always.
+            self.aggregated_results = {
+                **self.aggregated_results,
+                **self._aggregate_metrics(self.results, self.aggregation_methods),
+            }
+
+    def get_results(self):
+        return {
+            f"{aggregation_method}-{metric}": self.aggregated_results[f"{aggregation_method}-{metric}"]
+            for aggregation_method in self.aggregation_methods
+            for metric in self.metrics
+        }
+
+    def get_results_details(self):
+        return {f"{metric}": self.results[metric] for metric in self.metrics}
+
+    def save(self, path):
+        data = {
+            "version": self.version,
+            "results": self.results,
+            "aggregated_results": self.aggregated_results,
+            "metrics": self.metrics,
+            "aggregation_methods": self.aggregation_methods,
+            "dim_titles": self.dim_titles,
+        }
+
+        with open(path, "wb") as f:
+            pickle.dump(data, f)
+
+    @classmethod
+    def load(cls, path, embed, discrete_target=None, one_hot_target=None, metrics=None, aggregation_methods=None):
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+
+        assert cls.version == data["version"]
+        if metrics is None:
+            metrics = data["metrics"]
+        if aggregation_methods is None:
+            aggregation_methods = data["aggregation_methods"]
+        instance = cls(embed, discrete_target, one_hot_target, data["dim_titles"], metrics, aggregation_methods)
+        instance.results = data["results"]
+        instance.aggregated_results = data["aggregated_results"]
+        return instance
+
+    @classmethod
+    def load_results(cls, path):
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+        return data["aggregated_results"]
+
+    @classmethod
+    def load_results_details(cls, path):
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+        return data["results"]

src/drvi/utils/metrics/_pairwise.py

@@ -4,53 +4,71 @@
 from sklearn.feature_selection import mutual_info_classif
 
 
-def _nn_alignment_score_per_dim(var_continues, ct_cat_series):
+def check_discrete_metric_input(gt_cat_series=None, gt_one_hot=None):
+    if gt_cat_series is not None and gt_one_hot is not None:
+        raise ValueError("Only one of gt_cat_series or gt_one_hot should be provided.")
+    if gt_cat_series is None and gt_one_hot is None:
+        raise ValueError("Either gt_cat_series or gt_one_hot must be provided.")
+
+
+def get_one_hot_encoding(gt_cat_series):
+    return np.eye(len(gt_cat_series.cat.categories))[gt_cat_series.cat.codes]
+
+
+def _nn_alignment_score_per_dim(var_continues, gt_01):
     order = var_continues.argsort()
-    ct_cat_series = ct_cat_series[order]
-    ct_01 = np.eye(len(ct_cat_series.cat.categories))[ct_cat_series.cat.codes]
+    gt_01 = gt_01[order]
     alignment = np.clip(
         (
-            np.sum(ct_01[:-1, :] * ct_01[1:, :], axis=0) / (np.sum(ct_01, axis=0) - 1)
+            np.sum(gt_01[:-1, :] * gt_01[1:, :], axis=0) / (np.sum(gt_01, axis=0) - 1)
         )  # fraction of cells of this type that are next to a cell of the same type
-        - (np.sum(ct_01, axis=0) / ct_01.shape[0]),  # cancel random neighbors when CT is frequent
+        - (np.sum(gt_01, axis=0) / gt_01.shape[0]),  # cancel random neighbors when GT (ground-truth) is frequent
         0,
         None,
-    ) / (1 - (np.sum(ct_01, axis=0) / ct_01.shape[0]))
+    ) / (1 - (np.sum(gt_01, axis=0) / gt_01.shape[0]))
     return alignment
 
 
-def _local_mutual_info_score_per_binary_ct(all_vars_continues, ct_binary):
-    mi_score = mutual_info_classif(all_vars_continues, ct_binary, n_jobs=-1)
-    ct_prob = np.sum(ct_binary == 1) / ct_binary.shape[0]
-    ct_entropy = stats.entropy([ct_prob, 1 - ct_prob])
-    return mi_score / ct_entropy
-
+def nn_alignment_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None):
+    check_discrete_metric_input(gt_cat_series, gt_one_hot)
+    gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot
 
-def nn_alignment_score(all_vars_continues, ct_cat_series):
     n_vars = all_vars_continues.shape[1]
-    result = np.zeros([n_vars, len(ct_cat_series.cat.categories)])
+    result = np.zeros([n_vars, gt_01.shape[1]])
     for i in range(n_vars):
-        result[i, :] = _nn_alignment_score_per_dim(all_vars_continues[:, i], ct_cat_series)
+        result[i, :] = _nn_alignment_score_per_dim(all_vars_continues[:, i], gt_01)
     return result
 
 
-def local_mutual_info_score(all_vars_continues, ct_cat_series):
+def _local_mutual_info_score_per_binary_gt(all_vars_continues, gt_binary):
+    mi_score = mutual_info_classif(all_vars_continues, gt_binary, n_jobs=-1)
+    gt_prob = np.sum(gt_binary == 1) / gt_binary.shape[0]
+    gt_entropy = stats.entropy([gt_prob, 1 - gt_prob])
+    return mi_score / gt_entropy
+
+
+def local_mutual_info_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None):
+    check_discrete_metric_input(gt_cat_series, gt_one_hot)
+    gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot
+
     n_vars = all_vars_continues.shape[1]
-    result = np.zeros([n_vars, len(ct_cat_series.cat.categories)])
-    ct_01 = np.eye(len(ct_cat_series.cat.categories))[ct_cat_series.cat.codes].T
-    for j in range(ct_01.shape[0]):
-        result[:, j] = _local_mutual_info_score_per_binary_ct(all_vars_continues, ct_01[j])
+    result = np.zeros([n_vars, gt_01.shape[1]])
+    for j in range(gt_01.shape[1]):
+        result[:, j] = _local_mutual_info_score_per_binary_gt(all_vars_continues, gt_01[:, j])
     return result
 
 
-def global_dim_mutual_info_score(all_vars_continues, ct_cat_series):
-    mi_score = mutual_info_classif(all_vars_continues, ct_cat_series)
-    ct_entropy = stats.entropy(pd.Series(ct_cat_series).value_counts(normalize=True, sort=False))
-    return mi_score / ct_entropy
-
+def spearman_correlataion_score(all_vars_continues, gt_cat_series=None, gt_one_hot=None):
+    check_discrete_metric_input(gt_cat_series, gt_one_hot)
+    gt_01 = get_one_hot_encoding(gt_cat_series) if gt_cat_series is not None else gt_one_hot
 
-def spearman_correlataion_score(all_vars_continues, ct_cat_series):
     n_vars = all_vars_continues.shape[1]
-    ct_01 = np.eye(len(ct_cat_series.cat.categories))[ct_cat_series.cat.codes]
-    result = np.abs(stats.spearmanr(all_vars_continues, ct_01).statistic[:n_vars, n_vars:])
+    result = np.abs(stats.spearmanr(all_vars_continues, gt_01).statistic[:n_vars, n_vars:])
     return result
+
+
+def global_dim_mutual_info_score(all_vars_continues, gt_cat_series):
+    # This metric is not used in any analysis, but is provided for completeness.
+    mi_score = mutual_info_classif(all_vars_continues, gt_cat_series)
+    gt_entropy = stats.entropy(pd.Series(gt_cat_series).value_counts(normalize=True, sort=False))
+    return mi_score / gt_entropy