Merge pull request malariagen#1218 from khushthecoder/fix/GH1217-jitter-random-state

fix: make _jitter() accept optional random_state parameter (malariagen#1217)

Author: jonbrenas

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_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_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)