Add taxon, area params to plot_frequencies_time_series()

Author: leehart

malariagen_data/anoph/plotly_params.py

@@ -215,3 +215,13 @@
     ],
     "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,6 +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,
         **kwargs,
     ) -> plotly_params.figure:
         # Handle title.
@@ -947,6 +949,14 @@ 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]
+
         # Extract variant labels.
         variant_labels = ds["variant_label"].values
 

notebooks/plot_frequencies_space_time.ipynb

@@ -1,14 +1,22 @@
 {
  "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47f669f3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import malariagen_data"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "f820bc66-2fb2-4ca2-9b54-824e50d61a0a",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import malariagen_data\n",
-    "\n",
     "ag3 = malariagen_data.Ag3(\n",
     "    \"simplecache::gs://vo_agam_release_master_us_central1\",\n",
     "    simplecache=dict(cache_storage=\"../gcs_cache\"),\n",
@@ -23,8 +31,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import malariagen_data\n",
-    "\n",
     "af1 = malariagen_data.Af1(\n",
     "    \"simplecache::gs://vo_afun_release_master_us_central1\",\n",
     "    simplecache=dict(cache_storage=\"../gcs_cache\"),\n",
@@ -69,6 +75,16 @@
     "ag3.plot_frequencies_time_series(ds, height=500, width=1000)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "790c99e8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ag3.plot_frequencies_time_series(ds, taxon=\"gambiae\", height=500, width=1000)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -252,6 +268,16 @@
     "ag3.plot_frequencies_time_series(ds, height=900, width=900)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e16ab3fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ag3.plot_frequencies_time_series(ds, area=\"BF-09\", height=400, width=900)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -336,19 +362,11 @@
    "source": [
     "af1.plot_frequencies_interactive_map(ds)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a512b459",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "mgen_data_py3.11",
    "language": "python",
    "name": "python3"
   },
@@ -362,7 +380,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.5"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {