Merge branch 'fix/deprecation-warning' of https://github.com/aman-dev-maker/malariagen-data-python into fix/deprecation-warning

# Please enter a commit message to explain why this merge is necessary,
# especially if it merges an updated upstream into a topic branch.
#
# Lines starting with '#' will be ignored, and an empty message aborts
# the commit.

Author: aman-dev-maker

WINDOWS_SETUP.md

@@ -1,11 +1,11 @@
-# Windows Setup Guide
+Windows Setup Guide
 
 To get setup for development on Windows, see 
 [this video if you prefer VS Code](https://youtu.be/zddl3n1DCFM), 
 or [this older video if you prefer PyCharm](https://youtu.be/QniQi-Hoo9A), 
 and the instructions below.
 
-## 1. Fork and clone this repo
+ 1. Fork and clone this repo
 
 Go to https://github.com/malariagen/malariagen-data-python 
 and click the Fork button (top right corner).
@@ -16,7 +16,7 @@ git clone https://github.com/[username]/malariagen-data-python.git
 cd malariagen-data-python
 ```
 
-## 2. Install Python
+ 2. Install Python
 
 Download and install Python 3.10 from:
 https://www.python.org/downloads/windows/
@@ -29,7 +29,7 @@ Verify it worked:
 python --version
 ```
 
-## 3. Install pipx and poetry
+ 3. Install pipx and poetry
 ```bash
 python -m pip install --user pipx
 python -m pipx ensurepath
@@ -39,19 +39,19 @@ pipx install poetry
 > ⚠️ After running ensurepath, close and reopen 
 > PowerShell before continuing.
 
-## 4. Create and activate development environment
+ 4. Create and activate development environment
 ```bash
 poetry install
 poetry shell
 ```
 
-## 5. Install pre-commit hooks
+ 5. Install pre-commit hooks
 ```bash
 pipx install pre-commit
 pre-commit install
 ```
 
-## 6. Add upstream remote and get latest code
+ 6. Add upstream remote and get latest code
 ```bash
 git remote add upstream https://github.com/malariagen/malariagen-data-python
 git pull upstream master
@@ -60,29 +60,29 @@ git pull upstream master
 > ℹ️ Note: On Windows the default branch is called 
 > master not main.
 
-## 7. Verify everything works
+ 7. Verify everything works
 ```bash
 python -c "import malariagen_data; print('Setup successful!')"
 ```
 
-## Common Issues on Windows
+ Common Issues on Windows
 
-**poetry not found after install**  
+poetry not found after install
 Close and reopen PowerShell, then try again.
 
-**git not recognized**  
+git not recognized  
 Install Git from https://git-scm.com/download/win 
 and restart PowerShell.
 
-**python not recognized**  
+python not recognized
 Reinstall Python and make sure to check 
 "Add Python to PATH" during installation.
 
-**fatal: not a git repository**  
+fatal: not a git repository
 Make sure you are inside the malariagen-data-python 
 folder before running any git commands.
 Use: cd malariagen-data-python
 
-**error: pathspec main did not match**  
+error: pathspec main did not match
 On Windows use master instead of main.
-Run: git checkout master
+Run: git checkout master

malariagen_data/anoph/sample_metadata.py

@@ -1,4 +1,5 @@
 import io
+import json
 from itertools import cycle
 from typing import (
     Any,
@@ -81,6 +82,8 @@ def __init__(
 
         # Initialize cache attributes.
         self._cache_sample_metadata: Dict = dict()
+        self._cache_cohorts: Dict = dict()
+        self._cache_cohort_geometries: Dict = dict()
 
     def _metadata_paths(
         self,
@@ -1485,7 +1488,11 @@ def _setup_cohort_queries(
                 A cohort set name. Accepted values are:
                 "admin1_month", "admin1_quarter", "admin1_year",
                 "admin2_month", "admin2_quarter", "admin2_year".
-            """
+            """,
+            query="""
+                An optional pandas query string to filter the resulting
+                dataframe, e.g., "country == 'Burkina Faso'".
+            """,
         ),
         returns="""A dataframe of cohort data, one row per cohort. There are up to 18 columns:
         `cohort_id` is the identifier of the cohort,
@@ -1512,20 +1519,98 @@ def _setup_cohort_queries(
     def cohorts(
         self,
         cohort_set: base_params.cohorts,
+        query: Optional[str] = None,
     ) -> pd.DataFrame:
-        major_version_path = self._major_version_path
+        valid_cohort_sets = {
+            "admin1_month",
+            "admin1_quarter",
+            "admin1_year",
+            "admin2_month",
+            "admin2_quarter",
+            "admin2_year",
+        }
+        if cohort_set not in valid_cohort_sets:
+            raise ValueError(
+                f"{cohort_set!r} is not a valid cohort set. "
+                f"Accepted values are: {sorted(valid_cohort_sets)}."
+            )
+
+        cohorts_analysis = self._cohorts_analysis
+
+        # Cache to avoid repeated reads.
+        cache_key = (cohorts_analysis, cohort_set)
+        try:
+            df_cohorts = self._cache_cohorts[cache_key]
+        except KeyError:
+            major_version_path = self._major_version_path
+            path = f"{major_version_path[:2]}_cohorts/cohorts_{cohorts_analysis}/cohorts_{cohort_set}.csv"
+
+            with self.open_file(path) as f:
+                df_cohorts = pd.read_csv(f, sep=",", na_values="")
+
+            # Ensure all column names are lower case.
+            df_cohorts.columns = [c.lower() for c in df_cohorts.columns]  # type: ignore
+
+            self._cache_cohorts[cache_key] = df_cohorts
+
+        if query is not None:
+            df_cohorts = df_cohorts.query(query)
+            df_cohorts = df_cohorts.reset_index(drop=True)
+
+        return df_cohorts.copy()
+
+    @_check_types
+    @doc(
+        summary="""
+            Read GeoJSON geometry data for a specific cohort set,
+            providing boundary geometries for each cohort.
+        """,
+        parameters=dict(
+            cohort_set="""
+                A cohort set name. Accepted values are:
+                "admin1_month", "admin1_quarter", "admin1_year",
+                "admin2_month", "admin2_quarter", "admin2_year".
+            """,
+        ),
+        returns="""
+            A dict containing the parsed GeoJSON FeatureCollection,
+            with boundary geometries for each cohort in the set.
+        """,
+    )
+    def cohort_geometries(
+        self,
+        cohort_set: base_params.cohorts,
+    ) -> dict:
+        valid_cohort_sets = {
+            "admin1_month",
+            "admin1_quarter",
+            "admin1_year",
+            "admin2_month",
+            "admin2_quarter",
+            "admin2_year",
+        }
+        if cohort_set not in valid_cohort_sets:
+            raise ValueError(
+                f"{cohort_set!r} is not a valid cohort set. "
+                f"Accepted values are: {sorted(valid_cohort_sets)}."
+            )
+
         cohorts_analysis = self._cohorts_analysis
 
-        path = f"{major_version_path[:2]}_cohorts/cohorts_{cohorts_analysis}/cohorts_{cohort_set}.csv"
+        # Cache to avoid repeated reads.
+        cache_key = (cohorts_analysis, cohort_set)
+        try:
+            geojson_data = self._cache_cohort_geometries[cache_key]
+        except KeyError:
+            major_version_path = self._major_version_path
+            path = f"{major_version_path[:2]}_cohorts/cohorts_{cohorts_analysis}/cohorts_{cohort_set}.geojson"
 
-        # Read the manifest into a pandas dataframe.
-        with self.open_file(path) as f:
-            df_cohorts = pd.read_csv(f, sep=",", na_values="")
+            with self.open_file(path) as f:
+                geojson_data = json.load(f)
 
-        # Ensure all column names are lower case.
-        df_cohorts.columns = [c.lower() for c in df_cohorts.columns]  # type: ignore
+            self._cache_cohort_geometries[cache_key] = geojson_data
 
-        return df_cohorts
+        return geojson_data
 
     @_check_types
     @doc(

malariagen_data/util.py

@@ -855,9 +855,7 @@ def _value_error(
     value,
     expectation,
 ):
-    message = (
-        f"Bad value for parameter {name}; expected {expectation}, " f"found {value!r}"
-    )
+    message = f"Bad value for parameter {name}; expected {expectation}, found {value!r}"
     raise ValueError(message)
 
 
@@ -935,6 +933,7 @@ def info(self, msg):
         self.flush()
 
     def set_level(self, level):
+        self._logger.setLevel(level)
         if self._handler is not None:
             self._handler.setLevel(level)
 

notebooks/cohort_geometries.ipynb

@@ -0,0 +1,115 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Cohort Geometries\n",
+    "\n",
+    "Demonstrates the `cohort_geometries()` method for accessing GeoJSON boundary data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import malariagen_data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set up the Ag3 data resource"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ag3 = malariagen_data.Ag3(\n",
+    "    \"simplecache::gs://vo_agam_release_master_us_central1\",\n",
+    "    simplecache=dict(cache_storage=\"../gcs_cache\"),\n",
+    ")\n",
+    "ag3"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Access cohort geometries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "geojson = ag3.cohort_geometries(cohort_set=\"admin1_year\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inspect the returned data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "type(geojson)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "geojson.keys()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(geojson[\"features\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for f in geojson[\"features\"][:3]:\n",
+    "    print(f[\"properties\"])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

notebooks/plot_haplotypes_frequencies.ipynb

@@ -102,7 +102,7 @@
     "    sample_sets=[\"1232-VO-KE-OCHOMO-VMF00044\"],\n",
     "    min_cohort_size=10,\n",
     ")\n",
-    "ag3.plot_frequencies_time_series(hap_xr)"
+    "af1.plot_frequencies_time_series(hap_xr)"
    ]
   },
   {