Merge pull request #1014 from Sharon-codes/issue-798-cohort-diversity-cache

feat: add results caching to cohort_diversity_stats

Author: jonbrenas

malariagen_data/anopheles.py

@@ -371,6 +371,10 @@ def cohort_diversity_stats(
     ) -> pd.Series:
         debug = self._log.debug
 
+        # Change this name if you ever change the behaviour of this function, to
+        # invalidate any previously cached data.
+        name = "cohort_diversity_stats_v1"
+
         debug("process cohort parameter")
         cohort_query = None
         if isinstance(cohort, str):
@@ -391,28 +395,59 @@ def cohort_diversity_stats(
         else:
             raise TypeError(f"invalid cohort parameter: {cohort!r}")
 
-        debug("access allele counts")
-        ac = self.snp_allele_counts(
+        params = dict(
+            cohort_label=cohort_label,
+            cohort_query=cohort_query,
+            cohort_size=cohort_size,
             region=region,
+            min_cohort_size=min_cohort_size,
+            max_cohort_size=max_cohort_size,
             site_mask=site_mask,
             site_class=site_class,
-            sample_query=cohort_query,
             sample_sets=sample_sets,
-            cohort_size=cohort_size,
-            min_cohort_size=min_cohort_size,
-            max_cohort_size=max_cohort_size,
             random_seed=random_seed,
+            n_jack=n_jack,
+            confidence_level=confidence_level,
             chunks=chunks,
             inline_array=inline_array,
         )
 
-        debug("compute diversity stats")
-        stats = self._block_jackknife_cohort_diversity_stats(
-            cohort_label=cohort_label,
-            ac=ac,
-            n_jack=n_jack,
-            confidence_level=confidence_level,
-        )
+        # Try to retrieve results from the cache.
+        try:
+            results = self.results_cache_get(name=name, params=params)
+            stats = {
+                key: value.item()
+                if isinstance(value, np.ndarray) and value.shape == ()
+                else value
+                for key, value in results.items()
+            }
+
+        except CacheMiss:
+            debug("access allele counts")
+            ac = self.snp_allele_counts(
+                region=region,
+                site_mask=site_mask,
+                site_class=site_class,
+                sample_query=cohort_query,
+                sample_sets=sample_sets,
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                random_seed=random_seed,
+                chunks=chunks,
+                inline_array=inline_array,
+            )
+
+            debug("compute diversity stats")
+            stats = self._block_jackknife_cohort_diversity_stats(
+                cohort_label=cohort_label,
+                ac=ac,
+                n_jack=n_jack,
+                confidence_level=confidence_level,
+            )
+
+            cache_results = {key: np.asarray(value) for key, value in stats.items()}
+            self.results_cache_set(name=name, params=params, results=cache_results)
 
         debug("compute some extra cohort variables")
         df_samples = self.sample_metadata(

tests/anoph/test_cohort_diversity_stats.py

@@ -0,0 +1,58 @@
+import random
+
+import pandas as pd
+import pytest
+
+from malariagen_data import Ag3
+
+
+@pytest.fixture
+def ag3_sim_api(ag3_sim_fixture, tmp_path):
+    data_path = ag3_sim_fixture.bucket_path.as_posix()
+    return Ag3(
+        url=data_path,
+        public_url=data_path,
+        pre=True,
+        check_location=False,
+        bokeh_output_notebook=False,
+        results_cache=tmp_path.as_posix(),
+    )
+
+
+def test_cohort_diversity_stats_uses_cache(ag3_sim_api, monkeypatch):
+    api = ag3_sim_api
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    sample_set = random.choice(all_sample_sets)
+    df_samples = api.sample_metadata(sample_sets=[sample_set])
+    cohort_sample_ids = df_samples["sample_id"].head(10).to_list()
+    cohort_size = min(5, len(cohort_sample_ids))
+    if cohort_size < 2:
+        pytest.skip("not enough samples in simulated cohort")
+
+    params = dict(
+        cohort=("cache_test", f"sample_id in {cohort_sample_ids!r}"),
+        cohort_size=cohort_size,
+        region=random.choice(api.contigs),
+        sample_sets=[sample_set],
+        random_seed=42,
+        n_jack=10,
+        confidence_level=0.95,
+    )
+
+    stats_first = api.cohort_diversity_stats(**params)
+
+    def _unexpected_recompute(*args, **kwargs):  # noqa: ARG001, ARG002
+        raise AssertionError(
+            "cohort_diversity_stats recomputed instead of loading from cache"
+        )
+
+    monkeypatch.setattr(
+        api, "_block_jackknife_cohort_diversity_stats", _unexpected_recompute
+    )
+    stats_second = api.cohort_diversity_stats(**params)
+
+    pd.testing.assert_series_equal(
+        stats_first.sort_index(),
+        stats_second.sort_index(),
+        check_dtype=False,
+    )