Merge remote-tracking branch 'upstream/master' into refactor/xpehh-analysis

Author: 31puneet

malariagen_data/anoph/fst.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Tuple, Optional
 
 import numpy as np
@@ -43,6 +44,8 @@ def _fst_gwss(
         inline_array,
         chunks,
         clip_min,
+        min_snps_threshold,
+        window_adjustment_factor,
     ):
         # Compute allele counts.
         ac1 = self.snp_allele_counts(
@@ -81,6 +84,24 @@ def _fst_gwss(
                 chunks=chunks,
             ).compute()
 
+        n_snps = len(pos)
+        if n_snps < min_snps_threshold:
+            raise ValueError(
+                f"Too few SNP sites ({n_snps}) available for Fst GWSS. "
+                f"At least {min_snps_threshold} sites are required. "
+                "Try a larger genomic region or different site selection criteria."
+            )
+        if window_size >= n_snps:
+            adjusted_window_size = max(1, n_snps // window_adjustment_factor)
+            warnings.warn(
+                f"window_size ({window_size}) is >= the number of SNP sites "
+                f"available ({n_snps}); automatically adjusting window_size to "
+                f"{adjusted_window_size} (= {n_snps} // {window_adjustment_factor}).",
+                UserWarning,
+                stacklevel=2,
+            )
+            window_size = adjusted_window_size
+
         with self._spinner(desc="Compute Fst"):
             with np.errstate(divide="ignore", invalid="ignore"):
                 fst = allel.moving_hudson_fst(ac1, ac2, size=window_size)
@@ -96,8 +117,23 @@ def _fst_gwss(
     @doc(
         summary="""
             Run a Fst genome-wide scan to investigate genetic differentiation
-            between two cohorts.
+            between two cohorts. If window_size is >= the number of available
+            SNP sites, a UserWarning is issued and window_size is automatically
+            adjusted to number_of_snps // window_adjustment_factor. A ValueError
+            is raised if the number of available SNP sites is below
+            min_snps_threshold.
         """,
+        parameters=dict(
+            min_snps_threshold="""
+                Minimum number of SNP sites required. If fewer sites are
+                available a ValueError is raised.
+            """,
+            window_adjustment_factor="""
+                If window_size is >= the number of available SNP sites,
+                window_size is automatically set to
+                number_of_snps // window_adjustment_factor.
+            """,
+        ),
         returns=dict(
             x="An array containing the window centre point genomic positions",
             fst="An array with Fst statistic values for each window.",
@@ -123,6 +159,8 @@ def fst_gwss(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
         clip_min: fst_params.clip_min = 0.0,
+        min_snps_threshold: fst_params.min_snps_threshold = 1000,
+        window_adjustment_factor: fst_params.window_adjustment_factor = 10,
     ) -> Tuple[np.ndarray, np.ndarray]:
         # Change this name if you ever change the behaviour of this function, to
         # invalidate any previously cached data.
@@ -147,7 +185,13 @@ def fst_gwss(
             results = self.results_cache_get(name=name, params=params)
 
         except CacheMiss:
-            results = self._fst_gwss(**params, inline_array=inline_array, chunks=chunks)
+            results = self._fst_gwss(
+                **params,
+                inline_array=inline_array,
+                chunks=chunks,
+                min_snps_threshold=min_snps_threshold,
+                window_adjustment_factor=window_adjustment_factor,
+            )
             self.results_cache_set(name=name, params=params, results=results)
 
         x = results["x"]

malariagen_data/anoph/fst_params.py

@@ -34,6 +34,22 @@
     """,
 ]
 
+min_snps_threshold: TypeAlias = Annotated[
+    int,
+    """
+    Minimum number of SNP sites required for the Fst GWSS computation. If
+    fewer sites are available, a ValueError is raised.
+    """,
+]
+
+window_adjustment_factor: TypeAlias = Annotated[
+    int,
+    """
+    If window_size is >= the number of available SNP sites, the window_size
+    is automatically adjusted to number_of_snps // window_adjustment_factor.
+    """,
+]
+
 annotation: TypeAlias = Annotated[
     Optional[Literal["standard error", "Z score", "lower triangle"]],
     """

malariagen_data/anoph/snp_data.py

@@ -1,4 +1,5 @@
 import warnings
+from collections import OrderedDict
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -75,7 +76,9 @@ def __init__(
             base_params.site_mask, zarr.hierarchy.Group
         ] = dict()
         self._cache_site_annotations: Optional[zarr.hierarchy.Group] = None
-        self._cache_locate_site_class: Dict[Tuple[Any, ...], np.ndarray] = dict()
+        self._cache_locate_site_class: OrderedDict[
+            Tuple[Any, ...], np.ndarray
+        ] = OrderedDict()
 
         # Create the SNP-calls cache as a per-instance lru_cache wrapping the
         # bound method.  Storing it on the instance (rather than using a
@@ -878,6 +881,9 @@ def _locate_site_class(
 
         try:
             loc_ann = self._cache_locate_site_class[cache_key]
+            # Promote to most-recently-used so the LRU eviction below
+            # always removes the *least*-recently-used entry.
+            self._cache_locate_site_class.move_to_end(cache_key)
 
         except KeyError as exc:
             # Access site annotations data.
@@ -1023,9 +1029,10 @@ def _locate_site_class(
 
             self._cache_locate_site_class[cache_key] = loc_ann
 
-            # Evict the oldest entry when the cache exceeds its size limit.
-            # Plain dicts preserve insertion order (Python 3.7+), so the first
-            # key is always the oldest.
+            # Evict the least-recently-used entry when the cache exceeds its
+            # size limit.  Because the cache is an OrderedDict and both hits
+            # (move_to_end above) and inserts append to the right, the first
+            # key is always the *least*-recently-used entry.
             while len(self._cache_locate_site_class) > _LOCATE_SITE_CLASS_CACHE_MAXSIZE:
                 oldest = next(iter(self._cache_locate_site_class))
                 del self._cache_locate_site_class[oldest]

tests/anoph/test_fst.py

@@ -139,6 +139,52 @@ def test_fst_gwss(fixture, api: AnophelesFstAnalysis):
     assert isinstance(fig, bokeh.models.GridPlot)
 
 
+@parametrize_with_cases("fixture,api", cases=".")
+def test_fst_gwss_window_size_too_large(fixture, api: AnophelesFstAnalysis):
+    # When window_size exceeds available SNPs, a UserWarning must be issued and
+    # the function must still return a valid result using the adjusted window_size.
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    all_countries = api.sample_metadata()["country"].dropna().unique().tolist()
+    countries = np.random.choice(all_countries, size=2, replace=False).tolist()
+    cohort1_query = f"country == {countries[0]!r}"
+    cohort2_query = f"country == {countries[1]!r}"
+    with pytest.warns(UserWarning, match="window_size"):
+        x, fst = api.fst_gwss(
+            contig=str(np.random.choice(api.contigs)),
+            sample_sets=all_sample_sets,
+            cohort1_query=cohort1_query,
+            cohort2_query=cohort2_query,
+            site_mask=str(np.random.choice(api.site_mask_ids)),
+            window_size=10_000_000,  # far larger than any fixture SNP count
+            min_cohort_size=1,
+        )
+    assert isinstance(x, np.ndarray)
+    assert isinstance(fst, np.ndarray)
+    assert len(x) > 0
+    assert x.shape == fst.shape
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_fst_gwss_too_few_snps(fixture, api: AnophelesFstAnalysis):
+    # When min_snps_threshold exceeds available SNPs, a ValueError must be raised.
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    all_countries = api.sample_metadata()["country"].dropna().unique().tolist()
+    countries = np.random.choice(all_countries, size=2, replace=False).tolist()
+    cohort1_query = f"country == {countries[0]!r}"
+    cohort2_query = f"country == {countries[1]!r}"
+    with pytest.raises(ValueError, match="Too few SNP sites"):
+        api.fst_gwss(
+            contig=str(np.random.choice(api.contigs)),
+            sample_sets=all_sample_sets,
+            cohort1_query=cohort1_query,
+            cohort2_query=cohort2_query,
+            site_mask=str(np.random.choice(api.site_mask_ids)),
+            window_size=100,
+            min_cohort_size=1,
+            min_snps_threshold=10_000_000,  # far larger than any fixture SNP count (~28k-70k)
+        )
+
+
 @parametrize_with_cases("fixture,api", cases=".")
 def test_average_fst(fixture, api: AnophelesFstAnalysis):
     # Set up test parameters.

tests/anoph/test_snp_data.py

@@ -973,6 +973,64 @@ def test_locate_site_class_cache_is_bounded(ag3_sim_api: AnophelesSnpData):
     assert len(ag3_sim_api._cache_locate_site_class) <= _LOCATE_SITE_CLASS_CACHE_MAXSIZE
 
 
+def test_locate_site_class_cache_lru_eviction(ag3_sim_api: AnophelesSnpData):
+    """Verify true LRU semantics: recently *accessed* entries survive eviction,
+    while least-recently-used entries are evicted first."""
+    from collections import OrderedDict
+
+    from malariagen_data.anoph.snp_data import _LOCATE_SITE_CLASS_CACHE_MAXSIZE
+
+    cache = ag3_sim_api._cache_locate_site_class
+
+    # Start from a clean cache.
+    cache.clear()
+    assert isinstance(cache, OrderedDict)
+
+    maxsize = _LOCATE_SITE_CLASS_CACHE_MAXSIZE  # 64
+
+    # --- Phase 1: fill the cache to exactly maxsize ---
+    dummy = np.array([True, False])
+    for i in range(maxsize):
+        key = (f"contig_{i}", f"mask_{i}", f"class_{i}")
+        cache[key] = dummy
+    assert len(cache) == maxsize
+
+    # Remember the first key inserted (the oldest / least-recently-used).
+    first_key = ("contig_0", "mask_0", "class_0")
+    assert first_key in cache
+
+    # --- Phase 2: simulate an access (LRU promotion) on the first key ---
+    # move_to_end makes it the most-recently-used entry.
+    cache.move_to_end(first_key)
+
+    # Insert one more entry, exceeding maxsize.
+    overflow_key = ("overflow", "mask", "class")
+    cache[overflow_key] = dummy
+
+    # Evict to maintain the bound (same logic as _locate_site_class).
+    while len(cache) > maxsize:
+        oldest = next(iter(cache))
+        del cache[oldest]
+
+    # The first key should STILL be present because it was promoted.
+    assert (
+        first_key in cache
+    ), "LRU promotion via move_to_end must keep recently accessed entries alive"
+
+    # The second key ("contig_1", ...) — which was never re-accessed —
+    # should have been evicted as the new least-recently-used entry.
+    second_key = ("contig_1", "mask_1", "class_1")
+    assert (
+        second_key not in cache
+    ), "The least-recently-used entry should be evicted when cache exceeds maxsize"
+
+    # The overflow key should be present (it was just inserted).
+    assert overflow_key in cache
+
+    # Cache size must remain bounded.
+    assert len(cache) == maxsize
+
+
 def test_snp_calls_cache_is_per_instance(ag3_sim_api: AnophelesSnpData):
     """_cached_snp_calls must be a per-instance lru_cache, not a class-level one.