WIP: make taxon_by optional in ...frequencies_advanced funcs. Fix bug in groupby.

Author: leehart

malariagen_data/anoph/cnv_frq.py

@@ -445,7 +445,7 @@ def gene_cnv_frequencies_advanced(
         ci_method: Optional[frq_params.ci_method] = frq_params.ci_method_default,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
-        taxon_by: frq_params.taxon_by = frq_params.taxon_by_default,
+        taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
         regions: List[Region] = parse_multi_region(self, region)
         del region
@@ -530,7 +530,7 @@ def _gene_cnv_frequencies_advanced(
         )
 
         debug("group samples to make cohorts")
-        group_samples_by_cohort = df_samples.groupby(["taxon", "area", "period"])
+        group_samples_by_cohort = df_samples.groupby([taxon_by, "area", "period"])
 
         debug("build cohorts dataframe")
         df_cohorts = build_cohorts_from_sample_grouping(

malariagen_data/anoph/frq_base.py

@@ -14,9 +14,7 @@
 from .base import AnophelesBase
 
 
-def prep_samples_for_cohort_grouping(
-    *, df_samples, area_by, period_by, taxon_by="taxon"
-):
+def prep_samples_for_cohort_grouping(*, df_samples, area_by, period_by, taxon_by):
     # Take a copy, as we will modify the dataframe.
     df_samples = df_samples.copy()
 

malariagen_data/anoph/hap_frq.py

@@ -153,7 +153,7 @@ def haplotypes_frequencies_advanced(
         ci_method: Optional[frq_params.ci_method] = frq_params.ci_method_default,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
-        taxon_by: frq_params.taxon_by = frq_params.taxon_by_default,
+        taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
         # Load sample metadata.
         df_samples = self.sample_metadata(
@@ -171,7 +171,7 @@ def haplotypes_frequencies_advanced(
         )
 
         # Group samples to make cohorts.
-        group_samples_by_cohort = df_samples.groupby(["taxon_by", "area", "period"])
+        group_samples_by_cohort = df_samples.groupby([taxon_by, "area", "period"])
 
         # Build cohorts dataframe.
         df_cohorts = build_cohorts_from_sample_grouping(

malariagen_data/anoph/snp_frq.py

@@ -451,7 +451,7 @@ def snp_allele_frequencies_advanced(
         ci_method: Optional[frq_params.ci_method] = frq_params.ci_method_default,
         chunks: base_params.chunks = base_params.native_chunks,
         inline_array: base_params.inline_array = base_params.inline_array_default,
-        taxon_by: frq_params.taxon_by = frq_params.taxon_by_default,
+        taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
         # Load sample metadata.
         df_samples = self.sample_metadata(