Add cohort_geometries() for accessing cohort GeoJSON metadata

Author: adilraza99

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,7 @@ def __init__(
 
         # Initialize cache attributes.
         self._cache_sample_metadata: Dict = dict()
+        self._cache_cohort_geometries: Dict = dict()
 
     def _metadata_paths(
         self,
@@ -1527,6 +1529,59 @@ def cohorts(
 
         return df_cohorts
 
+    @_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
+
+        # 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"
+
+            with self.open_file(path) as f:
+                geojson_data = json.load(f)
+
+            self._cache_cohort_geometries[cache_key] = geojson_data
+
+        return geojson_data
+
     @_check_types
     @doc(
         summary="""

tests/anoph/conftest.py

@@ -1426,6 +1426,29 @@ def write_metadata(
                 for line in src.readlines()[:5]:
                     print(line, file=dst)
 
+            # Copy cohort GeoJSON fixtures.
+            geojson_files = [
+                "cohorts_admin1_month.geojson",
+                "cohorts_admin1_year.geojson",
+            ]
+            for geojson_file in geojson_files:
+                src_path = (
+                    self.fixture_dir
+                    / "vo_agam_release_master_us_central1"
+                    / "v3_cohorts"
+                    / "cohorts_20230516"
+                    / geojson_file
+                )
+                if src_path.exists():
+                    dst_path = (
+                        self.bucket_path
+                        / "v3_cohorts"
+                        / "cohorts_20230516"
+                        / geojson_file
+                    )
+                    dst_path.parent.mkdir(parents=True, exist_ok=True)
+                    shutil.copy2(src_path, dst_path)
+
         # Create data catalog by sampling from some real metadata files.
         src_path = (
             self.fixture_dir

tests/anoph/fixture/vo_agam_release_master_us_central1/v3_cohorts/cohorts_20230516/cohorts_admin1_month.geojson

@@ -0,0 +1,29 @@
+{
+  "type": "FeatureCollection",
+  "features": [
+    {
+      "type": "Feature",
+      "properties": {
+        "cohort_id": "BF-01_arab_2008_04",
+        "admin1_name": "Boucle du Mouhoun",
+        "admin1_iso": "BF-01"
+      },
+      "geometry": {
+        "type": "Polygon",
+        "coordinates": [[[-4.5, 12.0], [-3.5, 12.0], [-3.5, 13.0], [-4.5, 13.0], [-4.5, 12.0]]]
+      }
+    },
+    {
+      "type": "Feature",
+      "properties": {
+        "cohort_id": "BF-02_colu_2011_07",
+        "admin1_name": "Cascades",
+        "admin1_iso": "BF-02"
+      },
+      "geometry": {
+        "type": "Polygon",
+        "coordinates": [[[-5.0, 10.0], [-4.0, 10.0], [-4.0, 11.0], [-5.0, 11.0], [-5.0, 10.0]]]
+      }
+    }
+  ]
+}

tests/anoph/fixture/vo_agam_release_master_us_central1/v3_cohorts/cohorts_20230516/cohorts_admin1_year.geojson

@@ -0,0 +1,29 @@
+{
+  "type": "FeatureCollection",
+  "features": [
+    {
+      "type": "Feature",
+      "properties": {
+        "cohort_id": "AO-LUA_colu_2009",
+        "admin1_name": "Luanda",
+        "admin1_iso": "AO-LUA"
+      },
+      "geometry": {
+        "type": "Polygon",
+        "coordinates": [[[13.0, -10.0], [14.0, -10.0], [14.0, -9.0], [13.0, -9.0], [13.0, -10.0]]]
+      }
+    },
+    {
+      "type": "Feature",
+      "properties": {
+        "cohort_id": "BF-01_arab_2008",
+        "admin1_name": "Boucle du Mouhoun",
+        "admin1_iso": "BF-01"
+      },
+      "geometry": {
+        "type": "Polygon",
+        "coordinates": [[[-4.5, 12.0], [-3.5, 12.0], [-3.5, 13.0], [-4.5, 13.0], [-4.5, 12.0]]]
+      }
+    }
+  ]
+}

tests/anoph/test_sample_metadata.py

@@ -1465,3 +1465,48 @@ def test_cohort_data(fixture, api):
     df_cohorts = api.cohorts(cohort_name)
     # Check output.
     validate_cohort_data(df_cohorts, cohort_data_expected_columns())
+
+
+# ------------------------------------------------------------------
+# Tests for cohort_geometries()
+# ------------------------------------------------------------------
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_geometries(fixture, api):
+    """Test that GeoJSON geometry can be loaded for a valid cohort set."""
+    geojson = api.cohort_geometries("admin1_month")
+    assert isinstance(geojson, dict)
+    assert geojson["type"] == "FeatureCollection"
+    assert "features" in geojson
+    assert len(geojson["features"]) > 0
+    for feature in geojson["features"]:
+        assert feature["type"] == "Feature"
+        assert "geometry" in feature
+        assert "properties" in feature
+        assert "coordinates" in feature["geometry"]
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_geometries_admin1_year(fixture, api):
+    """Test that GeoJSON geometry can be loaded for admin1_year."""
+    geojson = api.cohort_geometries("admin1_year")
+    assert isinstance(geojson, dict)
+    assert geojson["type"] == "FeatureCollection"
+    assert len(geojson["features"]) > 0
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_geometries_invalid_cohort_set(fixture, api):
+    """Test that an invalid cohort_set raises ValueError."""
+    with suppress_type_checks():
+        with pytest.raises(ValueError, match="not a valid cohort set"):
+            api.cohort_geometries("invalid_set")
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_geometries_cached(fixture, api):
+    """Test that the second call returns the same cached object."""
+    g1 = api.cohort_geometries("admin1_month")
+    g2 = api.cohort_geometries("admin1_month")
+    assert g1 is g2