Merge branch 'master' into fix/pr817-ci-cleanup

Author: adilraza99

malariagen_data/anoph/fst.py

@@ -404,8 +404,14 @@ def average_fst(
         )
 
         # Calculate block length for jackknife.
-        n_sites = ac1.shape[0]  # number of sites
-        block_length = n_sites // n_jack  # number of sites in each block
+        n_sites = ac1.shape[0]
+        block_length = n_sites // n_jack
+
+        if block_length < 1:
+            raise ValueError(
+                f"Not enough sites ({n_sites}) for {n_jack} jackknife blocks. "
+                "Choose a larger region or reduce n_jack."
+            )
 
         # Calculate average Fst.
         fst, se, _, _ = allel.blockwise_hudson_fst(ac1, ac2, blen=block_length)
@@ -530,7 +536,7 @@ def plot_pairwise_average_fst(
 
         # Set up plot title.
         title = "<i>F</i><sub>ST</sub>"
-        if annotation is not None:
+        if annotation is not None and annotation != "lower triangle":
             title += " ⧅ " + annotation
 
         # Fill the figure dataframe from the Fst dataframe.
@@ -543,12 +549,15 @@ def plot_pairwise_average_fst(
                     fig_df.loc[cohort1, cohort2] = np.nan
                 else:
                     fig_df.loc[cohort1, cohort2] = fst / se
+            elif annotation == "lower triangle":
+                # Leave the upper triangle as NaN (empty).
+                pass
             else:
                 fig_df.loc[cohort1, cohort2] = fst
 
         # Don't colour the plot if the upper triangle is SE or Z score,
         # as the colouring doesn't really make sense.
-        if annotation is not None and zmax is None:
+        if annotation is not None and annotation != "lower triangle" and zmax is None:
             zmax = 1e9
 
         # Dynamically size the figure based on number of cohorts.

malariagen_data/anoph/fst_params.py

@@ -35,10 +35,12 @@
 ]
 
 annotation: TypeAlias = Annotated[
-    Optional[Literal["standard error", "Z score"]],
+    Optional[Literal["standard error", "Z score", "lower triangle"]],
     """
-    How to annotate the upper-right corner of the plot. Default behaviour (None) is using Fst, other options
-    are using the standard error (if annotation is 'standard error') or the Z score of the two
-    cohorts being the same (if annotation is 'Z score').
+    How to annotate the upper-right corner of the plot. Default behaviour (None)
+    is using Fst, other options are using the standard error (if annotation is
+    'standard error'), the Z score of the two cohorts being the same (if
+    annotation is 'Z score'), or leaving the upper triangle empty (if annotation
+    is 'lower triangle').
     """,
 ]

tests/anoph/test_fst.py

@@ -190,6 +190,27 @@ def test_average_fst_with_min_cohort_size(fixture, api: AnophelesFstAnalysis):
         api.average_fst(**fst_params)
 
 
+@parametrize_with_cases("fixture,api", cases=".")
+def test_average_fst_region_too_small(fixture, api: AnophelesFstAnalysis):
+    """ValueError should be raised when block_length == 0 (n_jack > n_sites)."""
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    all_countries = api.sample_metadata()["country"].dropna().unique().tolist()
+    countries = random.sample(all_countries, 2)
+    cohort1_query = f"country == {countries[0]!r}"
+    cohort2_query = f"country == {countries[1]!r}"
+    fst_params = dict(
+        region=random.choice(api.contigs),
+        sample_sets=all_sample_sets,
+        cohort1_query=cohort1_query,
+        cohort2_query=cohort2_query,
+        site_mask=random.choice(api.site_mask_ids),
+        min_cohort_size=1,
+        n_jack=1_000_000,  # deliberately exceeds available sites
+    )
+    with pytest.raises(ValueError, match="Not enough sites"):
+        api.average_fst(**fst_params)
+
+
 def check_pairwise_average_fst(api: AnophelesFstAnalysis, fst_params):
     # Run main function under test.
     fst_df = api.pairwise_average_fst(**fst_params)
@@ -236,6 +257,10 @@ def check_pairwise_average_fst(api: AnophelesFstAnalysis, fst_params):
         assert isinstance(fig, go.Figure)
         fig = api.plot_pairwise_average_fst(fst_df, annotation="Z score", show=False)
         assert isinstance(fig, go.Figure)
+        fig = api.plot_pairwise_average_fst(
+            fst_df, annotation="lower triangle", show=False
+        )
+        assert isinstance(fig, go.Figure)
 
 
 @pytest.mark.parametrize("cohorts", ["country", "admin1_year", "cohort_admin2_month"])