Extend test_allele_frequencies_advanced_with_period_by() to include period_by as random_year

Author: leehart

malariagen_data/anoph/cnv_frq.py

@@ -447,10 +447,12 @@ def gene_cnv_frequencies_advanced(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
-        # Check that the taxon_by default hasn't been subverted.
-        # This avoids type-checking errors, e.g. with `getattr`.
+        # Check that the taxon_by default hasn't been subverted, e.g. via `taxon_by=None`.
+        # This also satisfies type-checking, e.g. with `getattr`.
         if taxon_by is None:
-            raise ValueError("`taxon_by` cannot be set to `None`.")
+            raise ValueError(
+                f"`taxon_by` cannot be set to `None`. The default would be {frq_params.taxon_by_default!r}."
+            )
 
         regions: List[Region] = parse_multi_region(self, region)
         del region

malariagen_data/anoph/frq_base.py

@@ -57,7 +57,7 @@ def prep_samples_for_cohort_grouping(*, df_samples, area_by, period_by, taxon_by
             )
         ):
             raise TypeError(
-                "Invalid values in {period_by!r} column. Must be either pandas.Period or null."
+                f"Invalid values in {period_by!r} column. Must be either pandas.Period or null."
             )
 
         # Copy the specified period_by column to a new "period" column.

malariagen_data/anoph/hap_frq.py

@@ -155,10 +155,12 @@ def haplotypes_frequencies_advanced(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
-        # Check that the taxon_by default hasn't been subverted.
-        # This avoids type-checking errors, e.g. with `getattr`.
+        # Check that the taxon_by default hasn't been subverted, e.g. via `taxon_by=None`.
+        # This also satisfies type-checking, e.g. with `getattr`.
         if taxon_by is None:
-            raise ValueError("`taxon_by` cannot be set to `None`.")
+            raise ValueError(
+                f"`taxon_by` cannot be set to `None`. The default would be {frq_params.taxon_by_default!r}."
+            )
 
         # Load sample metadata.
         df_samples = self.sample_metadata(

malariagen_data/anoph/snp_frq.py

@@ -453,10 +453,12 @@ def snp_allele_frequencies_advanced(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
-        # Check that the taxon_by default hasn't been subverted.
-        # This avoids type-checking errors, e.g. with `getattr`.
+        # Check that the taxon_by default hasn't been subverted, e.g. via `taxon_by=None`.
+        # This also satisfies type-checking, e.g. with `getattr`.
         if taxon_by is None:
-            raise ValueError("`taxon_by` cannot be set to `None`.")
+            raise ValueError(
+                f"`taxon_by` cannot be set to `None`. The default would be {frq_params.taxon_by_default!r}."
+            )
 
         # Load sample metadata.
         df_samples = self.sample_metadata(
@@ -687,10 +689,12 @@ def aa_allele_frequencies_advanced(
         inline_array: base_params.inline_array = base_params.inline_array_default,
         taxon_by: Optional[frq_params.taxon_by] = frq_params.taxon_by_default,
     ) -> xr.Dataset:
-        # Check that the taxon_by default hasn't been subverted.
-        # This avoids type-checking errors, e.g. with `getattr`.
+        # Check that the taxon_by default hasn't been subverted, e.g. via `taxon_by=None`.
+        # This also satisfies type-checking, e.g. with `getattr`.
         if taxon_by is None:
-            raise ValueError("`taxon_by` cannot be set to `None`.")
+            raise ValueError(
+                f"`taxon_by` cannot be set to `None`. The default would be {frq_params.taxon_by_default!r}."
+            )
 
         # Begin by computing SNP allele frequencies.
         ds_snp_frq = self.snp_allele_frequencies_advanced(

tests/anoph/test_snp_frq.py

@@ -114,6 +114,30 @@ def random_transcript(*, api):
     return transcript
 
 
+def add_random_year(*, api):
+    # Add a 'random_year' column to the sample_metadata, if it doesn't exist.
+
+    # Get the existing sample metadata.
+    sample_metadata_df = api.sample_metadata()
+
+    # Only create the new column if it doesn't already exist.
+    # Otherwise we'll get multiple columns with different suffixes, e.g. 'random_year_x' and 'random_year_y'.
+    if "random_year" not in sample_metadata_df.columns:
+        random_years_as_list = np.random.choice(
+            range(1900, 2100), len(sample_metadata_df)
+        )
+        random_years_as_period_index = pd.PeriodIndex(random_years_as_list, freq="Y")
+        extra_metadata_df = pd.DataFrame(
+            {
+                "sample_id": sample_metadata_df["sample_id"],
+                "random_year": random_years_as_period_index,
+            }
+        )
+        api.add_extra_metadata(extra_metadata_df)
+
+    return api
+
+
 @parametrize_with_cases("fixture,api", cases=".")
 def test_snp_effects(fixture, api: AnophelesSnpFrequencyAnalysis):
     # Pick a random transcript.
@@ -888,7 +912,7 @@ def check_snp_allele_frequencies_advanced(
     if area_by is None:
         area_by = random.choice(["country", "admin1_iso", "admin2_name"])
     if period_by is None:
-        period_by = random.choice(["year", "quarter", "month"])
+        period_by = random.choice(["year", "quarter", "month", "random_year"])
     if sample_sets is None:
         all_sample_sets = api.sample_sets()["sample_set"].to_list()
         sample_sets = random.choice(all_sample_sets)
@@ -897,6 +921,10 @@ def check_snp_allele_frequencies_advanced(
     if site_mask is None:
         site_mask = random.choice(api.site_mask_ids + (None,))
 
+    if period_by == "random_year":
+        # Add a random_year column to the sample metadata, if there isn't already.
+        api = add_random_year(api=api)
+
     # Run function under test.
     ds = api.snp_allele_frequencies_advanced(
         transcript=transcript,
@@ -1002,6 +1030,8 @@ def check_snp_allele_frequencies_advanced(
         expected_freqstr = "M"
     elif period_by == "quarter":
         expected_freqstr = "Q-DEC"
+    elif period_by == "random_year":
+        expected_freqstr = "Y-DEC"
     else:
         assert False, "not implemented"
     for p in period_values:
@@ -1082,13 +1112,17 @@ def check_aa_allele_frequencies_advanced(
     if area_by is None:
         area_by = random.choice(["country", "admin1_iso", "admin2_name"])
     if period_by is None:
-        period_by = random.choice(["year", "quarter", "month"])
+        period_by = random.choice(["year", "quarter", "month", "random_year"])
     if sample_sets is None:
         all_sample_sets = api.sample_sets()["sample_set"].to_list()
         sample_sets = random.choice(all_sample_sets)
     if min_cohort_size is None:
         min_cohort_size = random.randint(0, 2)
 
+    if period_by == "random_year":
+        # Add a random_year column to the sample metadata, if there isn't already.
+        api = add_random_year(api=api)
+
     # Run function under test.
     ds = api.aa_allele_frequencies_advanced(
         transcript=transcript,
@@ -1185,6 +1219,8 @@ def check_aa_allele_frequencies_advanced(
         expected_freqstr = "M"
     elif period_by == "quarter":
         expected_freqstr = "Q-DEC"
+    elif period_by == "random_year":
+        expected_freqstr = "Y-DEC"
     else:
         assert False, "not implemented"
     for p in period_values:
@@ -1266,7 +1302,7 @@ def test_allele_frequencies_advanced_with_area_by(
     )
 
 
-@pytest.mark.parametrize("period_by", ["year", "quarter", "month"])
+@pytest.mark.parametrize("period_by", ["year", "quarter", "month", "random_year"])
 @parametrize_with_cases("fixture,api", cases=".")
 def test_allele_frequencies_advanced_with_period_by(
     fixture,