feat: add custom goterm set for the 3 supported species

Author: mffrank

grassp/tools/cluster_merging.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 from itertools import combinations
-from pathlib import Path
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Literal, Optional
 
 import numpy as np
 import scanpy as sc
@@ -12,58 +11,21 @@
 from scipy.spatial.distance import squareform
 from scipy.stats import fisher_exact
 
-from .enrichment import calculate_cluster_enrichment
+from .enrichment import _load_gmt, calculate_cluster_enrichment
+
+__all__ = [  # re-export private helper for callers/tests that imported it here
+    "_load_gmt",
+    "calculate_cluster_enrichment",
+    "dendrogram_cherry_pairs",
+    "merge_clusters_go",
+    "merge_small_clusters",
+    "paga_dendrogram",
+]
 
 if TYPE_CHECKING:
     from anndata import AnnData
 
 
-# ── Gene-set loading ──────────────────────────────────────────────────────────
-
-
-def _load_gmt(path: str | dict[str, list[str]] | None) -> dict[str, list[str]]:
-    """Parse a GMT file or gseapy library name into a gene-set dict.
-
-    Parameters
-    ----------
-    path
-        A ``dict`` (returned as-is), a path to a ``.gmt`` file, a gseapy
-        library name (fetched via ``gseapy.get_library``), or ``None``
-        (returns the default UniProt subcellular compartment gene sets
-        bundled with grassp).
-
-    Returns
-    -------
-    dict[str, list[str]]
-        Mapping of term name to list of gene symbols.
-    """
-    if path is None:
-        path = str(
-            Path(__file__).parent.parent
-            / 'datasets'
-            / 'external'
-            / 'custom_goterms_genes_reviewed.gmt'
-        )
-    if isinstance(path, dict):
-        return path
-    import os
-
-    if not os.path.exists(path):
-        import gseapy as gp
-
-        return gp.get_library(name=path)
-    gene_sets: dict[str, list[str]] = {}
-    with open(path) as f:
-        for line in f:
-            parts = line.strip().split('\t')
-            if len(parts) < 3:
-                continue
-            term = parts[0]
-            genes = [g for g in parts[2:] if g]
-            gene_sets[term] = genes
-    return gene_sets
-
-
 # ── Pair-testing helpers ──────────────────────────────────────────────────────
 
 
@@ -907,6 +869,7 @@ def merge_clusters_go(
     connectivity_lower: float = 0.5,
     cluster_col: str = 'leiden',
     gene_sets_path: Optional[str | dict] = None,
+    species: Literal['hsap', 'mmus', 'scer'] = 'hsap',
     gene_name_key: str = 'Gene_name_canonical',
     compartment_col: str = 'Cell_compartment',
     key_added: str = 'leiden_merged',
@@ -951,8 +914,15 @@ def merge_clusters_go(
         ``adata.obs`` column with initial cluster labels.
     gene_sets_path
         Path to a GMT file, a gseapy library name, a pre-loaded
-        ``dict[str, list[str]]``, or ``None`` (uses the default UniProt
-        subcellular compartment gene sets bundled with grassp).
+        ``dict[str, list[str]]``, or ``None`` (uses the consolidated UniProt
+        subcellular compartment gene sets for the chosen ``species``).
+    species
+        Species code used to pick the default gene-set file when
+        ``gene_sets_path`` is ``None``. One of ``"hsap"`` (human,
+        ``consolidated_goterms_human.gmt``), ``"mmus"`` (mouse,
+        ``consolidated_goterms_mouse.gmt``), or ``"scer"`` (yeast,
+        ``consolidated_goterms_yeast.gmt``). Default ``"hsap"``. Ignored when
+        an explicit ``gene_sets_path`` is provided.
     gene_name_key
         ``adata.obs`` column with gene/protein names used for enrichment.
     compartment_col
@@ -971,7 +941,7 @@ def merge_clusters_go(
         If True, plot the initial PAGA dendrogram after convergence with leaf
         lines colored by compartment term and merge nodes colored by p-value.
     """
-    gene_sets = _load_gmt(gene_sets_path)
+    gene_sets = _load_gmt(gene_sets_path, species=species)
 
     # Initialise working column from the original clustering
     adata.obs[key_added] = adata.obs[cluster_col].astype(str).astype('category')

grassp/tools/enrichment.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from pathlib import Path
 from typing import TYPE_CHECKING, Literal, Optional, Union
 
 if TYPE_CHECKING:
@@ -13,6 +14,74 @@
 rank_proteins_groups = scanpy.tl.rank_genes_groups
 
 
+# Map from species code → filename of the bundled consolidated GMT.
+# Shared between `calculate_cluster_enrichment` and `merge_clusters_go`.
+_SPECIES_TO_GMT_FILENAME: dict[str, str] = {
+    "hsap": "consolidated_goterms_human.gmt",
+    "mmus": "consolidated_goterms_mouse.gmt",
+    "scer": "consolidated_goterms_yeast.gmt",
+}
+
+
+def _load_gmt(
+    path: str | dict[str, list[str]] | None = None,
+    species: Literal["hsap", "mmus", "scer"] = "hsap",
+) -> dict[str, list[str]]:
+    """Resolve a gene-set source into a ``{term: [gene, ...]}`` dict.
+
+    Parameters
+    ----------
+    path
+        One of:
+        - ``dict`` — returned as-is.
+        - existing file path — parsed as GMT.
+        - non-existing string — treated as a ``gseapy`` library name and
+          fetched via :func:`gseapy.get_library`.
+        - ``None`` — uses the consolidated UniProt subcellular-location gene
+          sets bundled with grassp, picked according to ``species``.
+    species
+        Used only when ``path is None``. One of ``"hsap"``, ``"mmus"``,
+        ``"scer"``; selects the matching ``consolidated_goterms_*.gmt`` file
+        in ``grassp/datasets/external/``.
+
+    Returns
+    -------
+    dict[str, list[str]]
+        Mapping of term name → list of gene symbols.
+    """
+    if isinstance(path, dict):
+        return path
+    if path is None:
+        if species not in _SPECIES_TO_GMT_FILENAME:
+            raise ValueError(
+                f"species must be one of {sorted(_SPECIES_TO_GMT_FILENAME)}, "
+                f"got {species!r}"
+            )
+        path = str(
+            Path(__file__).parent.parent
+            / "datasets"
+            / "external"
+            / _SPECIES_TO_GMT_FILENAME[species]
+        )
+
+    import os
+
+    if not os.path.exists(path):
+        import gseapy as gp
+
+        return gp.get_library(name=path)
+    gene_sets: dict[str, list[str]] = {}
+    with open(path) as f:
+        for line in f:
+            parts = line.strip().split("\t")
+            if len(parts) < 3:
+                continue
+            term = parts[0]
+            genes = [g for g in parts[2:] if g]
+            gene_sets[term] = genes
+    return gene_sets
+
+
 def calculate_cluster_enrichment(
     data: AnnData,
     cluster_key: str = "leiden",
@@ -101,27 +170,10 @@ def calculate_cluster_enrichment(
     )
     # print(f"Sorting {enrichment_ranking_metric} in {'ascending' if sort_ascending else 'descending'} order")
 
-    if gene_sets is None:
-        # Pick the consolidated UniProt subcellular gene-set file for the
-        # requested species. Mapping matches the files produced by
-        # marker_curation/fetch_custom_goterms.py.
-        from pathlib import Path
-
-        species_to_filename = {
-            "hsap": "consolidated_goterms_human.gmt",
-            "mmus": "consolidated_goterms_mouse.gmt",
-            "scer": "consolidated_goterms_yeast.gmt",
-        }
-        if species not in species_to_filename:
-            raise ValueError(
-                f"species must be one of {sorted(species_to_filename)}, got {species!r}"
-            )
-        gene_sets = str(
-            Path(__file__).parent.parent
-            / "datasets"
-            / "external"
-            / species_to_filename[species]
-        )
+    # Resolve the gene-set source to a {term: [gene, ...]} dict. Handles dict
+    # passthrough, file paths, gseapy library names, and the species-specific
+    # default when `gene_sets is None`.
+    gene_sets = _load_gmt(gene_sets, species=species)
 
     for n, group in groups:
         gene_list = group[gene_name_key].tolist()