Merge pull request #1147 from khushthecoder/GH1064-fix-transient-aa-allele-frequencies-test

Fix #1064: Return empty DataFrame instead of raising ValueError in aa…

Author: jonbrenas

malariagen_data/anoph/snp_frq.py

@@ -345,11 +345,55 @@ def aa_allele_frequencies(
         # We just want aa change.
         df_ns_snps = df_snps.query(AA_CHANGE_QUERY).copy()
 
-        # Early check for no matching SNPs.
-        if len(df_ns_snps) == 0:  # pragma: no cover
-            raise ValueError(
-                "No amino acid change SNPs found for the given transcript and site mask."
+        # Handle case where no amino acid change SNPs are found.
+        # N.B., this can legitimately happen for some transcript/site_mask/query
+        # combinations. Return a well-formed empty DataFrame rather than raising,
+        # to avoid transient test failures and to allow downstream code to handle
+        # the empty result gracefully. See also:
+        # https://github.com/malariagen/malariagen-data-python/issues/1064
+        if len(df_ns_snps) == 0:
+            warnings.warn(
+                "No amino acid change SNPs found for the given transcript "
+                "and site mask. Returning an empty DataFrame.",
+                stacklevel=2,
+            )
+            # Build an empty DataFrame with the expected schema.
+            freq_cols = [col for col in df_snps.columns if col.startswith("frq_")]
+            count_cols = [col for col in df_snps.columns if col.startswith("count_")]
+            nobs_cols = [col for col in df_snps.columns if col.startswith("nobs_")]
+            keep_cols = [
+                "contig",
+                "transcript",
+                "aa_pos",
+                "ref_allele",
+                "ref_aa",
+                "alt_aa",
+                "effect",
+                "impact",
+                "grantham_score",
+                "sneath_score",
+            ]
+            all_cols = (
+                ["aa_change"]
+                + freq_cols
+                + ["max_af"]
+                + keep_cols
+                + ["alt_allele", "label", "position"]
             )
+            if include_counts:
+                all_cols = all_cols + count_cols + nobs_cols
+            df_empty = pd.DataFrame(columns=all_cols)
+            df_empty.set_index(["aa_change", "contig", "position"], inplace=True)
+
+            # Add metadata.
+            gene_name = self._transcript_to_parent_name(transcript)
+            title = transcript
+            if gene_name:
+                title += f" ({gene_name})"
+            title += " SNP frequencies"
+            df_empty.attrs["title"] = title
+
+            return df_empty
 
         # N.B., we need to worry about the possibility of the
         # same aa change due to SNPs at different positions. We cannot
@@ -375,7 +419,7 @@ def np_sum(g):
             for c in nobs_cols:
                 agg[c] = "first"
 
-        keep_cols = (
+        keep_cols = [
             "contig",
             "transcript",
             "aa_pos",
@@ -386,7 +430,7 @@ def np_sum(g):
             "impact",
             "grantham_score",
             "sneath_score",
-        )
+        ]
         for c in keep_cols:
             agg[c] = "first"
         agg["alt_allele"] = lambda v: "{" + ",".join(v) + "}" if len(v) > 1 else v

tests/anoph/test_snp_frq.py

@@ -403,14 +403,21 @@ def test_allele_frequencies_with_str_cohorts(
     # Run the function under test.
     df_aa = api.aa_allele_frequencies(**params)
 
-    check_plot_frequencies_heatmap(api, df_aa)
-
-    # Standard checks.
-    check_aa_allele_frequencies(
-        df=df_aa,
-        cohort_labels=cohort_labels,
-        transcript=transcript,
-    )
+    # Handle the case where no amino acid change SNPs are found.
+    # In this case, aa_allele_frequencies returns an empty DataFrame
+    # instead of raising (see issue #1064).
+    if len(df_aa) > 0:
+        check_plot_frequencies_heatmap(api, df_aa)
+
+        # Standard checks.
+        check_aa_allele_frequencies(
+            df=df_aa,
+            cohort_labels=cohort_labels,
+            transcript=transcript,
+        )
+    else:
+        assert isinstance(df_aa, pd.DataFrame)
+        assert df_aa.index.names == ["aa_change", "contig", "position"]
 
 
 @pytest.mark.parametrize("min_cohort_size", [0, 10, 100])
@@ -531,14 +538,22 @@ def test_allele_frequencies_with_str_cohorts_and_sample_query(
     # Run the function under test.
     df_aa = api.aa_allele_frequencies(**params)
 
-    check_plot_frequencies_heatmap(api, df_aa)
+    # Handle the case where no amino acid change SNPs are found.
+    # In this case, aa_allele_frequencies returns an empty DataFrame
+    # instead of raising (see issue #1064).
+    if len(df_aa) > 0:
+        check_plot_frequencies_heatmap(api, df_aa)
 
-    # Standard checks.
-    check_aa_allele_frequencies(
-        df=df_aa,
-        cohort_labels=cohort_labels,
-        transcript=transcript,
-    )
+        # Standard checks.
+        check_aa_allele_frequencies(
+            df=df_aa,
+            cohort_labels=cohort_labels,
+            transcript=transcript,
+        )
+    else:
+        # Verify the empty DataFrame has the expected structure.
+        assert isinstance(df_aa, pd.DataFrame)
+        assert df_aa.index.names == ["aa_change", "contig", "position"]
 
 
 @parametrize_with_cases(
@@ -604,14 +619,21 @@ def test_allele_frequencies_with_str_cohorts_and_sample_query_options(
     # Run the function under test.
     df_aa = api.aa_allele_frequencies(**params)
 
-    check_plot_frequencies_heatmap(api, df_aa)
+    # Handle the case where no amino acid change SNPs are found.
+    # In this case, aa_allele_frequencies returns an empty DataFrame
+    # instead of raising (see issue #1064).
+    if len(df_aa) > 0:
+        check_plot_frequencies_heatmap(api, df_aa)
 
-    # Standard checks.
-    check_aa_allele_frequencies(
-        df=df_aa,
-        cohort_labels=cohort_labels,
-        transcript=transcript,
-    )
+        # Standard checks.
+        check_aa_allele_frequencies(
+            df=df_aa,
+            cohort_labels=cohort_labels,
+            transcript=transcript,
+        )
+    else:
+        assert isinstance(df_aa, pd.DataFrame)
+        assert df_aa.index.names == ["aa_change", "contig", "position"]
 
 
 @parametrize_with_cases("fixture,api", cases=".")