Merge branch 'master' into optimize/cohort-heterozygosity-vectorized

Author: kunal-10-cloud

malariagen_data/anoph/heterozygosity.py

@@ -501,6 +501,27 @@ def cohort_count_het(
     def _roh_hmm_cache_name(self):
         return "roh_hmm_v1"
 
+    def _get_roh_hmm_cache_name(self):
+        """Safely resolve the ROH HMM cache name.
+
+        Supports class attribute, property, or legacy method override.
+        Falls back to the default "roh_hmm_v1" if resolution fails.
+
+        See also: https://github.com/malariagen/malariagen-data-python/issues/1151
+        """
+        try:
+            name = self._roh_hmm_cache_name
+            # Handle legacy case where _roh_hmm_cache_name might be a
+            # callable method rather than a property or class attribute.
+            if callable(name):
+                name = name()
+            if isinstance(name, str) and len(name) > 0:
+                return name
+        except NotImplementedError:
+            pass
+        # Fallback to default.
+        return "roh_hmm_v1"
+
     @_check_types
     @doc(
         summary="Infer runs of homozygosity for a single sample over a genome region.",
@@ -522,7 +543,7 @@ def roh_hmm(
 
         resolved_region: Region = _parse_single_region(self, region)
 
-        name = self._roh_hmm_cache_name
+        name = self._get_roh_hmm_cache_name()
 
         params = dict(
             sample=sample,

malariagen_data/anoph/pca.py

@@ -399,9 +399,9 @@ def plot_pca_coords(
 
         # Apply jitter if desired - helps spread out points when tightly clustered.
         if jitter_frac:
-            np.random.seed(random_seed)
-            data[x] = _jitter(data[x], jitter_frac)
-            data[y] = _jitter(data[y], jitter_frac)
+            rng = np.random.default_rng(seed=random_seed)
+            data[x] = _jitter(data[x], jitter_frac, random_state=rng)
+            data[y] = _jitter(data[y], jitter_frac, random_state=rng)
 
         # Convenience variables.
         # Prevent lint error (mypy): Unsupported operand types for + ("Series[Any]" and "str")
@@ -503,10 +503,10 @@ def plot_pca_coords_3d(
 
         # Apply jitter if desired - helps spread out points when tightly clustered.
         if jitter_frac:
-            np.random.seed(random_seed)
-            data[x] = _jitter(data[x], jitter_frac)
-            data[y] = _jitter(data[y], jitter_frac)
-            data[z] = _jitter(data[z], jitter_frac)
+            rng = np.random.default_rng(seed=random_seed)
+            data[x] = _jitter(data[x], jitter_frac, random_state=rng)
+            data[y] = _jitter(data[y], jitter_frac, random_state=rng)
+            data[z] = _jitter(data[z], jitter_frac, random_state=rng)
 
         # Convenience variables.
         # Prevent lint error (mypy): Unsupported operand types for + ("Series[Any]" and "str")

malariagen_data/util.py

@@ -897,10 +897,33 @@ def _hash_params(params):
     return h, s
 
 
-def _jitter(a, fraction):
-    """Jitter data in `a` using the fraction `f`."""
+def _jitter(a, fraction, random_state=np.random):
+    """Jitter data by adding uniform noise scaled by the data range.
+
+    Parameters
+    ----------
+    a : array-like
+        Input data to jitter. Can be a numpy array or pandas Series.
+    fraction : float
+        Controls the amplitude of the jitter relative to the data range.
+    random_state : numpy.random.Generator or module, optional
+        Random number generator to use. Accepts a ``numpy.random.Generator``
+        (from ``np.random.default_rng()``) or the ``numpy.random`` module.
+        Defaults to ``np.random`` (global RNG) for backward compatibility.
+
+    Returns
+    -------
+    array-like
+        Jittered copy of the input data with the same shape and type.
+
+    Notes
+    -----
+    Prefer passing a local ``np.random.default_rng(seed=...)`` to avoid
+    mutating global RNG state and to ensure reproducibility.
+
+    """
     r = a.max() - a.min()
-    return a + fraction * np.random.uniform(-r, r, a.shape)
+    return a + fraction * random_state.uniform(-r, r, a.shape)
 
 
 class CacheMiss(Exception):

tests/anoph/test_fst.py

@@ -346,8 +346,16 @@ def test_pairwise_average_fst_with_sample_query(fixture, api: AnophelesFstAnalys
         n_jack=random.randint(10, 200),
     )
 
-    # Run checks.
-    check_pairwise_average_fst(api=api, fst_params=fst_params)
+    # Run checks - skip if random parameter selection results in insufficient cohorts.
+    try:
+        check_pairwise_average_fst(api=api, fst_params=fst_params)
+    except ValueError as e:
+        if "No cohorts remain" in str(e):
+            pytest.skip(
+                f"Skipping: random parameter selection produced insufficient "
+                f"cohorts for taxon={taxon!r}: {e}"
+            )
+        raise
 
 
 @parametrize_with_cases("fixture,api", cases=".")

tests/anoph/test_hap_frq.py

@@ -187,7 +187,12 @@ def test_hap_frequencies_with_str_cohorts(
         return
 
     # Run the function under test.
-    df_hap = api.haplotypes_frequencies(**params)
+    try:
+        df_hap = api.haplotypes_frequencies(**params)
+    except ValueError as e:
+        if "No SNPs available for the given region" in str(e):
+            pytest.skip("Random region contained no SNPs")
+        raise
 
     check_plot_frequencies_heatmap(api, df_hap)
 

tests/anoph/test_heterozygosity.py

@@ -186,6 +186,19 @@ def test_roh_hmm(fixture, api: AnophelesHetAnalysis):
         assert col in df_roh.columns
 
 
+@parametrize_with_cases("fixture,api", cases=".")
+def test_roh_hmm_cache_name_resolution(fixture, api: AnophelesHetAnalysis):
+    """Regression test for GH#1151: _roh_hmm_cache_name must resolve to a string.
+
+    Verifies that the cache name resolver handles class attributes, properties,
+    and legacy method overrides without raising NotImplementedError.
+    """
+    # The resolver should always return a non-empty string.
+    name = api._get_roh_hmm_cache_name()
+    assert isinstance(name, str), f"Expected str, got {type(name)}"
+    assert len(name) > 0, "Cache name must be non-empty"
+
+
 @parametrize_with_cases("fixture,api", cases=".")
 def test_plot_roh(fixture, api: AnophelesHetAnalysis):
     # Set up test parameters.

tests/anoph/test_pca.py

@@ -313,3 +313,52 @@ def test_pca_fit_exclude_samples(fixture, api: AnophelesPca):
         len(pca_df.query(f"sample_id in {exclude_samples} and not pca_fit"))
         == n_samples_excluded
     )
+
+
+# --- _jitter() determinism unit tests ---
+
+
+def test_jitter_determinism():
+    """_jitter with the same seed must produce identical results."""
+    from malariagen_data.util import _jitter
+
+    a = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    fraction = 0.1
+
+    rng1 = np.random.default_rng(seed=42)
+    result1 = _jitter(a, fraction, random_state=rng1)
+
+    rng2 = np.random.default_rng(seed=42)
+    result2 = _jitter(a, fraction, random_state=rng2)
+
+    np.testing.assert_array_equal(result1, result2)
+
+
+def test_jitter_different_seeds():
+    """_jitter with different seeds must produce different results."""
+    from malariagen_data.util import _jitter
+
+    a = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
+    fraction = 0.1
+
+    rng1 = np.random.default_rng(seed=42)
+    result1 = _jitter(a, fraction, random_state=rng1)
+
+    rng2 = np.random.default_rng(seed=99)
+    result2 = _jitter(a, fraction, random_state=rng2)
+
+    assert not np.array_equal(result1, result2)
+
+
+def test_jitter_no_global_rng_side_effect():
+    """_jitter with explicit random_state must not alter global RNG state."""
+    from malariagen_data.util import _jitter
+
+    np.random.seed(0)
+    state_before = np.random.get_state()[1].copy()
+
+    rng = np.random.default_rng(seed=42)
+    _jitter(np.array([1.0, 2.0, 3.0]), 0.1, random_state=rng)
+
+    state_after = np.random.get_state()[1].copy()
+    np.testing.assert_array_equal(state_before, state_after)