diff --git a/pyproject.toml b/pyproject.toml index 42e9044..5e01d7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ dev = [ ] test = [ "coverage>=7.10", + "filelock", "pytest", "pytest-cov", # For VS Code’s coverage functionality ] diff --git a/src/cellcharter/pl/_autok.py b/src/cellcharter/pl/_autok.py index 0da7f80..68d0dc5 100644 --- a/src/cellcharter/pl/_autok.py +++ b/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) diff --git a/src/cellcharter/tl/_autok.py b/src/cellcharter/tl/_autok.py index bfcb9f7..af47a1c 100644 --- a/src/cellcharter/tl/_autok.py +++ b/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,11 +81,14 @@ 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 @@ -90,6 +96,11 @@ def __init__( 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() diff --git a/tests/conftest.py b/tests/conftest.py index 7688d2f..24da045 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_fixtures.py b/tests/test_fixtures.py new file mode 100644 index 0000000..7458304 --- /dev/null +++ b/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"] diff --git a/tests/tools/test_autok.py b/tests/tools/test_autok.py index 0191035..473e5c0 100644 --- a/tests/tools/test_autok.py +++ b/tests/tools/test_autok.py @@ -1,3 +1,6 @@ +import pickle +from shutil import copytree + import anndata as ad import numpy as np import pytest @@ -11,6 +14,7 @@ class TestClusterAutoK: @pytest.mark.parametrize("dataset_name", ["mibitof"]) def test_spatial_proteomics(self, dataset_name: str): + """Fit and predict spatial clusters on the proteomics reference dataset.""" download_dataset = getattr(sq.datasets, dataset_name) adata = download_dataset() if sps.issparse(adata.X): @@ -34,11 +38,13 @@ def test_spatial_proteomics(self, dataset_name: str): @pytest.mark.parametrize("random_state", [0, -1]) def test_random_state_reproducible(self, random_state: int): + """Reusing a seed reproduces the selected K and stability values.""" X, _ = make_blobs(n_samples=500, n_features=10, centers=4, cluster_std=3.0, random_state=0) adata = ad.AnnData(X=np.zeros((X.shape[0], 2), dtype=np.float32)) adata.obsm["X_rep"] = X.astype(np.float32) def fit(): + """Fit a fresh model with the selected seed.""" model_params = { "random_state": random_state, "trainer_params": {"accelerator": "cpu", "enable_progress_bar": False}, @@ -52,3 +58,100 @@ def fit(): autok1, autok2 = fit(), fit() assert autok1.best_k == autok2.best_k np.testing.assert_allclose(autok1.stability, autok2.stability) + + @staticmethod + def _blobs_adata(): + """Create a small reproducible dataset for clustering regressions.""" + X, _ = make_blobs(n_samples=300, n_features=10, centers=4, cluster_std=3.0, random_state=0) + adata = ad.AnnData(X=np.zeros((X.shape[0], 2), dtype=np.float32)) + adata.obsm["X_rep"] = X.astype(np.float32) + return adata + + @staticmethod + def _model_params(): + """Use deterministic CPU models without progress bars.""" + return {"random_state": 0, "trainer_params": {"accelerator": "cpu", "enable_progress_bar": False}} + + def test_non_consecutive_n_clusters(self): + """Fit neighbors of every requested K and predict with the selected model.""" + adata = self._blobs_adata() + autok = cc.tl.ClusterAutoK(n_clusters=[6, 2, 4], model_params=self._model_params(), max_runs=3) + autok.fit(adata, use_rep="X_rep") + + assert autok.n_clusters == [2, 4, 6] + # Each K is fitted together with its neighbors K-1 and K+1 + assert sorted(autok.labels.keys()) == [1, 2, 3, 4, 5, 6, 7] + assert autok.stability.shape[0] == 3 + assert autok.best_k in autok.n_clusters + assert len(autok.predict(adata, use_rep="X_rep").categories) == autok.best_k + + def test_tuple_equivalent_to_list(self): + """A tuple range and its explicit list produce identical stability.""" + adata = self._blobs_adata() + autok_tuple = cc.tl.ClusterAutoK(n_clusters=(2, 4), model_params=self._model_params(), max_runs=2) + autok_list = cc.tl.ClusterAutoK(n_clusters=[2, 3, 4], model_params=self._model_params(), max_runs=2) + autok_tuple.fit(adata, use_rep="X_rep") + autok_list.fit(adata, use_rep="X_rep") + + assert autok_tuple.n_clusters == autok_list.n_clusters == [2, 3, 4] + np.testing.assert_allclose(autok_tuple.stability, autok_list.stability) + + def test_n_clusters_lower_than_two_ignored(self): + """Ignore unsupported counts and reject an entirely invalid range.""" + with pytest.warns(UserWarning): + autok = cc.tl.ClusterAutoK(n_clusters=(1, 4)) + assert autok.n_clusters == [2, 3, 4] + with pytest.raises(ValueError): + cc.tl.ClusterAutoK(n_clusters=[1]) + + def test_load_legacy_model(self): + """Convert the stored border counts in the legacy multi-run model.""" + # Saved with n_clusters=(2, 5), when `n_clusters` also stored the neighbors 1 and 6 + autok = cc.tl.ClusterAutoK.load("tests/_models/cellcharter_autok_mibitof") + assert autok.n_clusters == [2, 3, 4, 5] + assert autok.best_k in autok.n_clusters + + @pytest.mark.parametrize("lower", [2, 3]) + def test_load_legacy_single_run_model(self, tmp_path, lower): + """Legacy single-run models discard border counts before they are refitted.""" + path = tmp_path / "legacy" + copytree("tests/_models/cellcharter_autok_mibitof", path) + with (path / "params.pickle").open("rb") as f: + params = pickle.load(f) + params["n_clusters"] = list(range(lower - 1, 7)) + params["max_runs"] = 1 + with (path / "params.pickle").open("wb") as f: + pickle.dump(params, f) + with (path / "attributes.pickle").open("rb") as f: + attributes = pickle.load(f) + attributes["stability"] = None + attributes["labels"] = {k: labels[:1] for k, labels in attributes["labels"].items() if k >= lower - 1} + with (path / "attributes.pickle").open("wb") as f: + pickle.dump(attributes, f) + + autok = cc.tl.ClusterAutoK.load(path) + assert autok.n_clusters == list(range(lower, 6)) + assert autok.stability is None + autok.fit(self._blobs_adata(), use_rep="X_rep") + assert sorted(autok.labels) == list(range(lower - 1, 7)) + + @pytest.mark.parametrize("max_runs", [1, 2]) + @pytest.mark.parametrize("n_clusters", [[2, 3, 4], [2, 5, 9]]) + def test_save_load_preserves_requested_counts(self, tmp_path, max_runs, n_clusters): + """Loading current models preserves all requested counts and predictions.""" + adata = self._blobs_adata() + autok = cc.tl.ClusterAutoK(n_clusters=n_clusters, max_runs=max_runs, model_params=self._model_params()) + autok.fit(adata, use_rep="X_rep") + autok.save(tmp_path) + + loaded = cc.tl.ClusterAutoK.load(tmp_path) + assert loaded.n_clusters == n_clusters + assert loaded._fit_n_clusters == autok._fit_n_clusters + if max_runs == 1: + assert loaded.stability is None + else: + np.testing.assert_array_equal(loaded.stability, autok.stability) + for k in n_clusters: + np.testing.assert_array_equal( + loaded.predict(adata, use_rep="X_rep", k=k), autok.predict(adata, use_rep="X_rep", k=k) + )