Expand taxa, areas params for plot_frequencies_time_series()

Author: leehart

malariagen_data/anoph/frq_params.py

@@ -1,6 +1,6 @@
 """Parameter definitions for functions computing and plotting allele frequencies."""
 
-from typing import Literal
+from typing import Literal, List, Optional, Tuple, Union
 
 import xarray as xr
 from typing_extensions import Annotated, TypeAlias
@@ -70,3 +70,13 @@
     bool,
     "Include columns with allele counts and number of non-missing allele calls (nobs).",
 ]
+
+taxa: TypeAlias = Annotated[
+    Optional[Union[str, List[str], Tuple[str, ...]]],
+    "The taxon or taxa to restrict the dataset to.",
+]
+
+areas: TypeAlias = Annotated[
+    Optional[Union[str, List[str], Tuple[str, ...]]],
+    "The area or areas to restrict the dataset to.",
+]

malariagen_data/anoph/plotly_params.py

@@ -215,13 +215,3 @@
     ],
     "Scope of the map.",
 ]
-
-taxon: TypeAlias = Annotated[
-    Optional[str],
-    "The taxon to restrict the dataset to.",
-]
-
-area: TypeAlias = Annotated[
-    Optional[str],
-    "The area to restrict the dataset to.",
-]

malariagen_data/anoph/snp_frq.py

@@ -936,8 +936,8 @@ def plot_frequencies_time_series(
         legend_sizing: plotly_params.legend_sizing = "constant",
         show: plotly_params.show = True,
         renderer: plotly_params.renderer = None,
-        taxon: plotly_params.taxon = None,
-        area: plotly_params.area = None,
+        taxa: frq_params.taxa = None,
+        areas: frq_params.areas = None,
         **kwargs,
     ) -> plotly_params.figure:
         # Handle title.
@@ -949,13 +949,17 @@ def plot_frequencies_time_series(
         df_cohorts = ds[cohort_vars].to_dataframe()
         df_cohorts.columns = [c.split("cohort_")[1] for c in df_cohorts.columns]  # type: ignore
 
-        # If specified, restrict the dataframe by taxon.
-        if taxon:
-            df_cohorts = df_cohorts[df_cohorts["taxon"] == taxon]
-
-        # If specified, restrict the dataframe by area.
-        if area:
-            df_cohorts = df_cohorts[df_cohorts["area"] == area]
+        # If specified, restrict the dataframe by taxa.
+        if isinstance(taxa, str):
+            df_cohorts = df_cohorts[df_cohorts["taxon"] == taxa]
+        elif isinstance(taxa, (list, tuple)):
+            df_cohorts = df_cohorts[df_cohorts["taxon"].isin(taxa)]
+
+        # If specified, restrict the dataframe by areas.
+        if isinstance(areas, str):
+            df_cohorts = df_cohorts[df_cohorts["area"] == areas]
+        elif isinstance(areas, (list, tuple)):
+            df_cohorts = df_cohorts[df_cohorts["area"].isin(areas)]
 
         # Extract variant labels.
         variant_labels = ds["variant_label"].values

notebooks/plot_frequencies_space_time.ipynb

@@ -82,7 +82,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ag3.plot_frequencies_time_series(ds, taxon=\"gambiae\", height=500, width=1000)"
+    "ag3.plot_frequencies_time_series(ds, taxa=\"gambiae\", height=500, width=1000)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1bfc7298",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ag3.plot_frequencies_time_series(ds, taxa=(\"gambiae\", \"arabiensis\"), height=500, width=1000)"
    ]
   },
   {
@@ -275,7 +285,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ag3.plot_frequencies_time_series(ds, area=\"BF-09\", height=400, width=900)"
+    "ag3.plot_frequencies_time_series(ds, areas=\"BF-09\", height=400, width=900)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26af27a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ag3.plot_frequencies_time_series(ds, areas=(\"BF-09\", \"TZ-25\"), height=400, width=900)"
    ]
   },
   {