Progress on haplotype frequencies

Author: jonbrenas

malariagen_data/anoph/hap_frq.py

@@ -120,7 +120,7 @@ def haplotypes_frequencies(
         df_haps_sorted["label"] = ["H" + str(i) for i in range(len(df_haps_sorted))]
 
         # Reset index after filtering.
-        df_haps_sorted.set_index(keys="label", drop=True)
+        df_haps_sorted.set_index(keys="label", drop=True, inplace=True)
 
         return df_haps_sorted
 

tests/anoph/test_frq.py

@@ -15,9 +15,13 @@ def test_plot_frequencies_heatmap(api, frq_df):
 
     # Test index parameter - if None, should use dataframe index.
     fig = api.plot_frequencies_heatmap(frq_df, show=False, index=None, max_len=None)
-    # Not unique.
-    with pytest.raises(ValueError):
-        api.plot_frequencies_heatmap(frq_df, show=False, index="contig", max_len=None)
+
+    if "contig" in list(frq_df.columns):
+        # Not unique.
+        with pytest.raises(ValueError):
+            api.plot_frequencies_heatmap(
+                frq_df, show=False, index="contig", max_len=None
+            )
 
 
 @pytest.mark.skip

tests/anoph/test_hap_frq.py

@@ -8,6 +8,12 @@
 
 from malariagen_data import ag3 as _ag3
 from malariagen_data.anoph.hap_frq import AnophelesHapFrequencyAnalysis
+from .test_frq import (
+    test_plot_frequencies_heatmap,
+    test_plot_frequencies_time_series,
+    test_plot_frequencies_time_series_with_taxa,
+    test_plot_frequencies_time_series_with_areas,
+)
 
 
 @pytest.fixture
@@ -70,9 +76,8 @@ def check_hap_frequencies(*, api, df, sample_sets, cohorts, min_cohort_size):
     cohort_counts = df_samples[cohort_column].value_counts()
     cohort_labels = cohort_counts[cohort_counts >= min_cohort_size].index.to_list()
 
-    universal_fields = ["label"]
     frq_fields = ["frq_" + s for s in cohort_labels] + ["max_af"]
-    expected_fields = universal_fields + frq_fields
+    expected_fields = frq_fields
     assert sorted(df.columns.tolist()) == sorted(expected_fields)
 
 
@@ -82,9 +87,9 @@ def check_hap_frequencies_advanced(
     ds,
 ):
     assert isinstance(ds, xr.Dataset)
-    # test_plot_frequencies_time_series(api, ds)
-    # test_plot_frequencies_time_series_with_taxa(api, ds)
-    # test_plot_frequencies_time_series_with_areas(api, ds)
+    test_plot_frequencies_time_series(api, ds)
+    test_plot_frequencies_time_series_with_taxa(api, ds)
+    test_plot_frequencies_time_series_with_areas(api, ds)
     # test_plot_frequencies_interactive_map(api, ds)
     assert set(ds.dims) == {"cohorts", "variants"}
 
@@ -158,7 +163,7 @@ def test_hap_frequencies_with_str_cohorts(
     # Run the function under test.
     df_hap = api.haplotypes_frequencies(**params)
 
-    # test_plot_frequencies_heatmap(api, df_hap)
+    test_plot_frequencies_heatmap(api, df_hap)
 
     # Standard checks.
     check_hap_frequencies(