Support non-consecutive n_clusters in ClusterAutoK

pyproject.toml

@@ -47,6 +47,7 @@ dev = [
 ]
 test = [
   "coverage>=7.10",
+  "filelock",
   "pytest",
   "pytest-cov",     # For VS Code’s coverage functionality
 ]

src/cellcharter/pl/_autok.py

@@ -29,12 +29,12 @@ def autok_stability(autok: ClusterAutoK, save: str | Path | None = None, return_
         Nothing, just plots the figure and optionally saves the plot.
     """
     robustness_df = pd.melt(
-        pd.DataFrame.from_dict({k: autok.stability[i] for i, k in enumerate(autok.n_clusters[1:-1])}, orient="columns"),
+        pd.DataFrame.from_dict({k: autok.stability[i] for i, k in enumerate(autok.n_clusters)}, orient="columns"),
         var_name="N. clusters",
         value_name="Stability",
     )
     ax = sns.lineplot(data=robustness_df, x="N. clusters", y="Stability")
-    ax.set_xticks(autok.n_clusters[1:-1])
+    ax.set_xticks(autok.n_clusters)
     if save:
         plt.savefig(save)
 

src/cellcharter/tl/_autok.py

@@ -33,7 +33,10 @@ class ClusterAutoK:
     Parameters
     ----------
     n_clusters
-        Range for number of clusters (bounds included).
+        Numbers of clusters to evaluate. Either a tuple ``(min, max)`` defining a range (bounds included), or a list
+        of values, which can be non-consecutive (e.g. ``[5, 10, 15, 20]``). The stability of each `K` is computed
+        against `K-1` and `K+1`, so these neighbors are also fitted when they are not already included.
+        Values lower than 2 cannot be evaluated and are ignored.
     max_runs
         Maximum number of repetitions for each value of number of clusters.
     convergence_tol
@@ -66,7 +69,7 @@ class ClusterAutoK:
     #: The cluster assignments for each repetition and number of clusters.
     labels: dict
 
-    #: The stability values of all combinations of runs between K and K-1, and between K and K+1
+    #: The stability values of all combinations of runs between K and K-1, and between K and K+1, for each K in `n_clusters`
     stability: np.ndarray
 
     def __init__(
@@ -78,18 +81,26 @@ def __init__(
         model_params: dict = None,
         similarity_function: callable = None,
     ):
-        self.n_clusters = (
-            list(range(*(max(1, n_clusters[0] - 1), n_clusters[1] + 2)))
-            if isinstance(n_clusters, tuple)
-            else n_clusters
-        )
+        n_clusters = range(n_clusters[0], n_clusters[1] + 1) if isinstance(n_clusters, tuple) else n_clusters
+        n_clusters = sorted({int(k) for k in n_clusters})
+        if any(k < 2 for k in n_clusters):
+            warnings.warn("Stability cannot be computed for less than 2 clusters, ignoring them.", stacklevel=2)
+            n_clusters = [k for k in n_clusters if k >= 2]
+        if len(n_clusters) == 0:
+            raise ValueError("`n_clusters` must contain at least one value greater than or equal to 2.")
+        self.n_clusters = n_clusters
         self.max_runs = max_runs
         self.convergence_tol = convergence_tol
         self.model_class = model_class if model_class else cc.tl.GaussianMixture
         self.model_params = model_params if model_params else {}
         self.similarity_function = similarity_function if similarity_function else fowlkes_mallows_score
         self.stability = []
 
+    @property
+    def _fit_n_clusters(self) -> list[int]:
+        """Numbers of clusters to fit: each K in `n_clusters` together with its neighbors K-1 and K+1."""
+        return sorted({k + d for k in self.n_clusters for d in (-1, 0, 1)})
+
     def fit(self, adata: ad.AnnData, use_rep: str = "X_cellcharter"):
         """
         Cluster data multiple times for each number of clusters (K) in the selected range and compute the average stability for each them.
@@ -122,12 +133,15 @@ def fit(self, adata: ad.AnnData, use_rep: str = "X_cellcharter"):
         trainer_params = self.model_params.get("trainer_params") or {}
         self.model_params["trainer_params"] = {"enable_progress_bar": False, **trainer_params}
 
+        fit_n_clusters = self._fit_n_clusters
+        # Stability between K and K+1, for each K such that both K and K+1 are fitted
+        pair_stability = defaultdict(list)
         previous_stability = None
         for i in range(self.max_runs):
             print(f"Iteration {i + 1}/{self.max_runs}")
             new_labels = {}
 
-            for k in tqdm(self.n_clusters, disable=(len(self.n_clusters) == 1)):
+            for k in tqdm(fit_n_clusters, disable=(len(fit_n_clusters) == 1)):
                 with silence_lightning():
                     clustering = self.model_class(n_clusters=k, random_state=run_seeds[i], **self.model_params)
                     clustering.fit(X)
@@ -137,18 +151,16 @@ def fit(self, adata: ad.AnnData, use_rep: str = "X_cellcharter"):
                     self.best_models[k] = clustering
 
             if i > 0:
+                pair_ks = [k for k in fit_n_clusters if k + 1 in new_labels]
                 with concurrent.futures.ThreadPoolExecutor() as executor:
-                    pairs = [
-                        (new_labels[k], self.labels[k + 1][i])
-                        for i in range(len(list(self.labels.values())[0]))
-                        for k in list(self.labels.keys())[:-1]
-                    ]
-                    self.stability.extend(list(executor.map(lambda x: self.similarity_function(*x), pairs)))
+                    for k in pair_ks:
+                        pairs = [(new_labels[k], previous_labels) for previous_labels in self.labels[k + 1]]
+                        pair_stability[k].extend(executor.map(lambda x: self.similarity_function(*x), pairs))
 
                 if previous_stability is not None:
                     stability_change = mean_absolute_percentage_error(
                         np.mean(self._mirror_stability(previous_stability), axis=1),
-                        np.mean(self._mirror_stability(self.stability), axis=1),
+                        np.mean(self._mirror_stability(pair_stability), axis=1),
                     )
                     if stability_change < self.convergence_tol:
                         for k, new_l in new_labels.items():
@@ -158,40 +170,36 @@ def fit(self, adata: ad.AnnData, use_rep: str = "X_cellcharter"):
                         )
                         break
 
-                previous_stability = deepcopy(self.stability)
+                previous_stability = deepcopy(pair_stability)
 
             for k, new_l in new_labels.items():
                 self.labels[k].append(new_l)
 
         if self.max_runs > 1:
-            self.stability = self._mirror_stability(self.stability)
+            self.stability = self._mirror_stability(pair_stability)
         else:
             self.stability = None
 
-    def _mirror_stability(self, stability):
-        stability = [
-            stability[i : i + len(self.n_clusters) - 1] for i in range(0, len(stability), len(self.n_clusters) - 1)
-        ]
-        stability = list(map(list, zip(*stability, strict=False)))
-        return np.array([stability[i] + stability[i - 1] for i in range(1, len(stability))])
+    def _mirror_stability(self, pair_stability: dict) -> np.ndarray:
+        """For each K in `n_clusters`, combine its stability with K+1 and with K-1."""
+        return np.array([pair_stability[k] + pair_stability[k - 1] for k in self.n_clusters])
 
     @property
     def best_k(self) -> int:
         """The number of clusters with the highest stability."""
         if self.max_runs <= 1:
             raise ValueError("Cannot compute stability with max_runs <= 1")
-        stability_mean = np.array([np.mean(self.stability[k]) for k in range(len(self.n_clusters[1:-1]))])
-        best_idx = np.argmax(stability_mean)
-        return self.n_clusters[best_idx + 1]
+        stability_mean = np.mean(self.stability, axis=1)
+        return self.n_clusters[np.argmax(stability_mean)]
 
     @property
     def peaks(self) -> list[int]:
         """Find the peaks in the stability curve."""
         if self.max_runs <= 1:
             raise ValueError("Cannot compute stability with max_runs <= 1")
-        stability_mean = np.array([np.mean(self.stability[k]) for k in range(len(self.n_clusters[1:-1]))])
+        stability_mean = np.mean(self.stability, axis=1)
         peaks, _ = find_peaks(stability_mean)
-        return np.array(self.n_clusters[1:-1])[peaks]
+        return np.array(self.n_clusters)[peaks]
 
     def predict(self, adata: ad.AnnData, use_rep: str = None, k: int = None) -> pd.Categorical:
         """
@@ -207,7 +215,7 @@ def predict(self, adata: ad.AnnData, use_rep: str = None, k: int = None) -> pd.C
             Number of clusters to predict using the fitted model. If ``None``, the number of clusters with the highest stability will be selected. If ``max_runs > 1``, the model with the largest marginal likelihood will be used among the ones fitted on ``k``.
         """
         k = self.best_k if k is None else k
-        assert k is None or k in self.n_clusters
+        assert k in self.best_models, f"No model was fitted with {k} clusters"
 
         X = (
             adata.obsm[use_rep]
@@ -356,7 +364,10 @@ def _load_parameters(cls, path: Path):
             with pickle_path.open("rb") as f:
                 params = pickle.load(f)
 
-        return cls(**params)
+        with warnings.catch_warnings():
+            # Models saved with older versions also stored K=1 as neighbor of K=2
+            warnings.filterwarnings("ignore", message="Stability cannot be computed")
+            return cls(**params)
 
     def _load_attributes(self, path: Path) -> None:
         """
@@ -385,6 +396,13 @@ def _load_attributes(self, path: Path) -> None:
             with pickle_path.open("rb") as f:
                 self.set_params(pickle.load(f))
 
+        # Legacy parameters stored every fitted K, including the border neighbors.
+        # The constructor may already have discarded K=1. Detect this layout even
+        # for a single run, which has no stability values.
+        fitted_counts = sorted(self.labels)
+        if self.n_clusters == [k for k in fitted_counts if k >= 2]:
+            self.n_clusters = fitted_counts[1:-1]
+
         self.best_models = {}
         for model_dir in os.listdir(path / "best_models"):
             model = self.model_class()

tests/conftest.py

@@ -7,6 +7,7 @@
 import numpy as np
 import pytest
 import scanpy as sc
+from filelock import FileLock
 from squidpy._constants._pkg_constants import Key
 
 _adata = sc.read("tests/_data/test_data.h5ad")
@@ -57,19 +58,30 @@ def adata() -> ad.AnnData:
 
 @pytest.fixture(scope="session")
 def codex_adata() -> ad.AnnData:
+    """Load the two CODEX samples shared by the graph and shape tests."""
+    return _load_codex(_CODEX_PATH)
+
+
+def _load_codex(path: Path) -> ad.AnnData:
+    """Download and read the shared fixture without races between pytest workers."""
     max_retries = 3
     retry_delay = 5  # seconds
 
-    for attempt in range(max_retries):
-        try:
-            if not _is_hdf5_file(_CODEX_PATH):
-                _CODEX_PATH.unlink(missing_ok=True)
-                _download_codex(_CODEX_PATH)
-            adata = ad.read_h5ad(_CODEX_PATH)
-            adata.obs_names_make_unique()
-            return adata[adata.obs["sample"].isin(["BALBc-1", "MRL-5"])].copy()
-        except (HTTPError, URLError, OSError) as e:
-            _CODEX_PATH.unlink(missing_ok=True)  # Force re-download on next attempt
-            if attempt == max_retries - 1:  # Last attempt
-                pytest.skip(f"Failed to download test data after {max_retries} attempts: {str(e)}")
-            time.sleep(retry_delay)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with FileLock(str(path) + ".lock"):
+        for attempt in range(max_retries):
+            try:
+                if not _is_hdf5_file(path):
+                    path.unlink(missing_ok=True)
+                    _download_codex(path)
+                adata = ad.read_h5ad(path)
+            except (HTTPError, URLError, OSError, KeyError, ValueError) as e:
+                path.unlink(missing_ok=True)  # Force re-download on next attempt
+                if attempt == max_retries - 1:  # Last attempt
+                    pytest.skip(f"Failed to download test data after {max_retries} attempts: {str(e)}")
+                time.sleep(retry_delay)
+            else:
+                break
+
+    adata.obs_names_make_unique()
+    return adata[adata.obs["sample"].isin(["BALBc-1", "MRL-5"])].copy()

tests/test_fixtures.py

@@ -0,0 +1,70 @@
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
+
+import anndata as ad
+import numpy as np
+import pytest
+
+from . import conftest
+
+
+def _write_codex(path):
+    """Write a small fixture containing the selected and excluded samples."""
+    adata = ad.AnnData(np.arange(6).reshape(3, 2))
+    adata.obs["sample"] = ["BALBc-1", "MRL-5", "other"]
+    adata.write_h5ad(path)
+
+
+def test_concurrent_codex_loads(tmp_path, monkeypatch):
+    """Concurrent workers download once and each read the complete dataset."""
+    downloads = []
+    barrier = Barrier(4)
+    path = tmp_path / "codex.h5ad"
+
+    def download(path):
+        downloads.append(path)
+        _write_codex(path)
+
+    def load():
+        barrier.wait(timeout=10)
+        return conftest._load_codex(path)
+
+    monkeypatch.setattr(conftest, "_download_codex", download)
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        futures = [executor.submit(load) for _ in range(4)]
+        results = [future.result(timeout=30) for future in futures]
+
+    assert downloads == [path]
+    for adata in results:
+        assert adata.obs["sample"].tolist() == ["BALBc-1", "MRL-5"]
+        np.testing.assert_array_equal(adata.X, np.arange(4).reshape(2, 2))
+
+
+@pytest.mark.parametrize("error", [OSError, KeyError, ValueError])
+def test_corrupt_codex_cache_is_retried(tmp_path, monkeypatch, error):
+    """Unreadable cached HDF5 files are downloaded again before use."""
+    path = tmp_path / "codex.h5ad"
+    _write_codex(path)
+    read_h5ad = ad.read_h5ad
+    reads = []
+    downloads = []
+
+    def read(path):
+        reads.append(path)
+        if len(reads) == 1:
+            raise error("corrupt HDF5 data")
+        return read_h5ad(path)
+
+    def download(path):
+        assert not path.exists()
+        downloads.append(path)
+        _write_codex(path)
+
+    monkeypatch.setattr(conftest.ad, "read_h5ad", read)
+    monkeypatch.setattr(conftest, "_download_codex", download)
+    monkeypatch.setattr(conftest.time, "sleep", lambda _: None)
+    adata = conftest._load_codex(path)
+
+    assert downloads == [path]
+    assert len(reads) == 2
+    assert adata.obs["sample"].tolist() == ["BALBc-1", "MRL-5"]