Merge branch 'master' into fix/TODO-strings

Author: blankirigaya

malariagen_data/anoph/cnv_frq.py

@@ -446,6 +446,7 @@ def gene_cnv_frequencies_advanced(
         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,
+        filter_unassigned: Optional[frq_params.filter_unassigned] = None,
     ) -> xr.Dataset:
         regions: List[Region] = _parse_multi_region(self, region)
         del region
@@ -468,6 +469,7 @@ def gene_cnv_frequencies_advanced(
                     chunks=chunks,
                     inline_array=inline_array,
                     taxon_by=taxon_by,
+                    filter_unassigned=filter_unassigned,
                 )
                 for r in regions
             ],
@@ -497,6 +499,7 @@ def _gene_cnv_frequencies_advanced(
         chunks,
         inline_array,
         taxon_by,
+        filter_unassigned,
     ):
         debug = self._log.debug
 
@@ -527,6 +530,7 @@ def _gene_cnv_frequencies_advanced(
             area_by=area_by,
             period_by=period_by,
             taxon_by=taxon_by,
+            filter_unassigned=filter_unassigned,
         )
 
         debug("group samples to make cohorts")

malariagen_data/anoph/frq_base.py

@@ -16,20 +16,30 @@
 from .base import AnophelesBase
 
 
-def _prep_samples_for_cohort_grouping(*, df_samples, area_by, period_by, taxon_by):
+def _prep_samples_for_cohort_grouping(
+    *, df_samples, area_by, period_by, taxon_by, filter_unassigned=None
+):
     # Take a copy, as we will modify the dataframe.
     df_samples = df_samples.copy()
 
-    # Fix "intermediate" or "unassigned" taxon values - we only want to build
-    # cohorts with clean taxon calls, so we set other values to None.
-    loc_intermediate_taxon = (
-        df_samples[taxon_by].str.startswith("intermediate").fillna(False)
-    )
-    df_samples.loc[loc_intermediate_taxon, taxon_by] = None
-    loc_unassigned_taxon = (
-        df_samples[taxon_by].str.startswith("unassigned").fillna(False)
-    )
-    df_samples.loc[loc_unassigned_taxon, taxon_by] = None
+    # Determine whether to filter "intermediate"/"unassigned" taxon values.
+    # See: https://github.com/malariagen/malariagen-data-python/issues/806
+    if filter_unassigned is None:
+        # Auto-apply filtering only when using the default "taxon" column.
+        # Users can explicitly override with True/False.
+        filter_unassigned = taxon_by == "taxon"
+
+    if filter_unassigned:
+        # Remove samples with "intermediate" or "unassigned" taxon values,
+        # as we only want cohorts with clean taxon calls.
+        loc_intermediate_taxon = (
+            df_samples[taxon_by].str.startswith("intermediate").fillna(False)
+        )
+        df_samples.loc[loc_intermediate_taxon, taxon_by] = None
+        loc_unassigned_taxon = (
+            df_samples[taxon_by].str.startswith("unassigned").fillna(False)
+        )
+        df_samples.loc[loc_unassigned_taxon, taxon_by] = None
 
     # Add period column.
 

malariagen_data/anoph/frq_params.py

@@ -87,3 +87,13 @@
 ]
 
 taxon_by_default: taxon_by = "taxon"
+
+filter_unassigned: TypeAlias = Annotated[
+    Optional[bool],
+    """
+    Whether to filter out samples with "intermediate" or "unassigned" taxon
+    values before building cohorts. If None (default), filtering is applied
+    only when using the default "taxon" column. Set True to always filter,
+    or False to never filter.
+    """,
+]

malariagen_data/anoph/genome_sequence.py

@@ -16,6 +16,8 @@
 
 
 class AnophelesGenomeSequenceData(AnophelesBase):
+    _cache_genome: Optional[zarr.hierarchy.Group]
+
     def __init__(
         self, virtual_contigs: Optional[Mapping[str, Sequence[str]]] = None, **kwargs
     ):
@@ -29,9 +31,6 @@ def __init__(
             virtual_contigs = dict()
         self._virtual_contigs = virtual_contigs
 
-        # Initialize cache attributes.
-        self._cache_genome = None
-
     @property
     def contigs(self) -> Tuple[str, ...]:
         """Contigs in the reference genome."""
@@ -62,13 +61,9 @@ def _genome_ref_id(self) -> str:
     def _genome_ref_name(self) -> str:
         return self.config["GENOME_REF_NAME"]
 
-    @_check_types
-    @doc(
-        summary="Open the reference genome zarr.",
-        returns="Zarr hierarchy containing the reference genome sequence.",
-    )
     def open_genome(self) -> zarr.hierarchy.Group:
-        if self._cache_genome is None:
+        """Open the reference genome zarr."""
+        if not hasattr(self, "_cache_genome") or self._cache_genome is None:
             path = f"{self._base_path}/{self._genome_zarr_path}"
             store = _init_zarr_store(fs=self._fs, path=path)
             self._cache_genome = zarr.open_consolidated(store=store)

malariagen_data/anoph/hap_frq.py

@@ -154,6 +154,7 @@ def haplotypes_frequencies_advanced(
         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,
+        filter_unassigned: Optional[frq_params.filter_unassigned] = None,
     ) -> xr.Dataset:
         # Load sample metadata.
         df_samples = self.sample_metadata(
@@ -168,6 +169,7 @@ def haplotypes_frequencies_advanced(
             area_by=area_by,
             period_by=period_by,
             taxon_by=taxon_by,
+            filter_unassigned=filter_unassigned,
         )
 
         # Group samples to make cohorts.

malariagen_data/anoph/snp_frq.py

@@ -1,3 +1,4 @@
+from functools import cached_property
 from typing import Optional, Dict, Union, Callable, List
 import warnings
 
@@ -40,9 +41,6 @@ def __init__(
         # to the superclass constructor.
         super().__init__(**kwargs)
 
-        # Set up cache variables.
-        self._cache_annotator = None
-
     def _snp_df_melt(self, *, ds_snp: xr.Dataset) -> pd.DataFrame:
         """Set up a dataframe with SNP site and filter data,
         melting each alternate allele into a separate row."""
@@ -81,13 +79,12 @@ def _snp_df_melt(self, *, ds_snp: xr.Dataset) -> pd.DataFrame:
 
         return df_snps
 
+    @cached_property
     def _snp_effect_annotator(self):
         """Set up variant effect annotator."""
-        if self._cache_annotator is None:
-            self._cache_annotator = veff.Annotator(
-                genome=self.open_genome(), genome_features=self.genome_features()
-            )
-        return self._cache_annotator
+        return veff.Annotator(
+            genome=self.open_genome(), genome_features=self.genome_features()
+        )
 
     @_check_types
     @doc(
@@ -112,7 +109,7 @@ def snp_effects(
         df_snps = self._snp_df_melt(ds_snp=ds_snp)
 
         # Setup variant effect annotator.
-        ann = self._snp_effect_annotator()
+        ann = self._snp_effect_annotator
 
         # Add effects to the dataframe.
         ann.get_effects(transcript=transcript, variants=df_snps)
@@ -249,7 +246,7 @@ def snp_allele_frequencies(
 
         if effects:
             # Add effect annotations.
-            ann = self._snp_effect_annotator()
+            ann = self._snp_effect_annotator
             ann.get_effects(
                 transcript=transcript, variants=df_snps, progress=self._progress
             )
@@ -453,6 +450,7 @@ def snp_allele_frequencies_advanced(
         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,
+        filter_unassigned: Optional[frq_params.filter_unassigned] = None,
     ) -> xr.Dataset:
         # Load sample metadata.
         df_samples = self.sample_metadata(
@@ -467,6 +465,7 @@ def snp_allele_frequencies_advanced(
             area_by=area_by,
             period_by=period_by,
             taxon_by=taxon_by,
+            filter_unassigned=filter_unassigned,
         )
 
         # Group samples to make cohorts.
@@ -586,7 +585,7 @@ def snp_allele_frequencies_advanced(
             frequency = np.compress(loc_variant, frequency, axis=0)
 
         # Set up variant effect annotator.
-        ann = self._snp_effect_annotator()
+        ann = self._snp_effect_annotator
 
         # Add effects to the dataframe.
         ann.get_effects(
@@ -684,6 +683,7 @@ def aa_allele_frequencies_advanced(
         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,
+        filter_unassigned: Optional[frq_params.filter_unassigned] = None,
     ) -> xr.Dataset:
         # Begin by computing SNP allele frequencies.
         ds_snp_frq = self.snp_allele_frequencies_advanced(
@@ -702,6 +702,7 @@ def aa_allele_frequencies_advanced(
             chunks=chunks,
             inline_array=inline_array,
             taxon_by=taxon_by,
+            filter_unassigned=filter_unassigned,
         )
 
         # N.B., we need to worry about the possibility of the
@@ -832,7 +833,7 @@ def snp_genotype_allele_counts(
         df_snps = pd.concat([df_snps, df_counts], axis=1)
 
         # Add effect annotations.
-        ann = self._snp_effect_annotator()
+        ann = self._snp_effect_annotator
         ann.get_effects(
             transcript=transcript, variants=df_snps, progress=self._progress
         )

tests/anoph/test_frq_base.py

@@ -0,0 +1,101 @@
+"""Tests for _prep_samples_for_cohort_grouping filter_unassigned behavior.
+
+See: https://github.com/malariagen/malariagen-data-python/issues/806
+"""
+
+import pandas as pd
+
+from malariagen_data.anoph.frq_base import _prep_samples_for_cohort_grouping
+
+
+def _make_test_df(taxon_col="taxon"):
+    """Create a test DataFrame with intermediate and unassigned taxon values."""
+    return pd.DataFrame(
+        {
+            taxon_col: [
+                "gambiae",
+                "intermediate_gambcolu_arabiensis",
+                "unassigned",
+                "coluzzii",
+            ],
+            "admin1_iso": ["KE-01", "KE-01", "KE-02", "KE-02"],
+            "year": [2020, 2020, 2020, 2020],
+            "month": [1, 1, 1, 1],
+        }
+    )
+
+
+class TestPrepSamplesFilterUnassigned:
+    """Tests for the filter_unassigned parameter in _prep_samples_for_cohort_grouping."""
+
+    def test_default_taxon_column_filters(self):
+        """When taxon_by='taxon' and filter_unassigned=None (default),
+        intermediate/unassigned values should be set to None (backward compat)."""
+        df = _make_test_df(taxon_col="taxon")
+        result = _prep_samples_for_cohort_grouping(
+            df_samples=df,
+            area_by="admin1_iso",
+            period_by="year",
+            taxon_by="taxon",
+        )
+        assert result["taxon"].iloc[0] == "gambiae"
+        assert result["taxon"].iloc[1] is None
+        assert result["taxon"].iloc[2] is None
+        assert result["taxon"].iloc[3] == "coluzzii"
+
+    def test_custom_column_preserves(self):
+        """When taxon_by is a custom column and filter_unassigned=None (default),
+        intermediate/unassigned values should be preserved."""
+        df = _make_test_df(taxon_col="custom_taxon")
+        result = _prep_samples_for_cohort_grouping(
+            df_samples=df,
+            area_by="admin1_iso",
+            period_by="year",
+            taxon_by="custom_taxon",
+        )
+        assert result["custom_taxon"].iloc[0] == "gambiae"
+        assert result["custom_taxon"].iloc[1] == "intermediate_gambcolu_arabiensis"
+        assert result["custom_taxon"].iloc[2] == "unassigned"
+        assert result["custom_taxon"].iloc[3] == "coluzzii"
+
+    def test_explicit_filter_true(self):
+        """When filter_unassigned=True, always filter regardless of column name."""
+        df = _make_test_df(taxon_col="custom_taxon")
+        result = _prep_samples_for_cohort_grouping(
+            df_samples=df,
+            area_by="admin1_iso",
+            period_by="year",
+            taxon_by="custom_taxon",
+            filter_unassigned=True,
+        )
+        assert result["custom_taxon"].iloc[0] == "gambiae"
+        assert result["custom_taxon"].iloc[1] is None
+        assert result["custom_taxon"].iloc[2] is None
+        assert result["custom_taxon"].iloc[3] == "coluzzii"
+
+    def test_explicit_filter_false(self):
+        """When filter_unassigned=False, never filter even for default 'taxon' column."""
+        df = _make_test_df(taxon_col="taxon")
+        result = _prep_samples_for_cohort_grouping(
+            df_samples=df,
+            area_by="admin1_iso",
+            period_by="year",
+            taxon_by="taxon",
+            filter_unassigned=False,
+        )
+        assert result["taxon"].iloc[0] == "gambiae"
+        assert result["taxon"].iloc[1] == "intermediate_gambcolu_arabiensis"
+        assert result["taxon"].iloc[2] == "unassigned"
+        assert result["taxon"].iloc[3] == "coluzzii"
+
+    def test_does_not_modify_original(self):
+        """Ensure the original DataFrame is not modified."""
+        df = _make_test_df(taxon_col="taxon")
+        original_values = df["taxon"].tolist()
+        _prep_samples_for_cohort_grouping(
+            df_samples=df,
+            area_by="admin1_iso",
+            period_by="year",
+            taxon_by="taxon",
+        )
+        assert df["taxon"].tolist() == original_values