Merge pull request #1052 from adilraza99/GH435-cohorts-robustness

Improve robustness of cohorts() metadata access

Author: jonbrenas

malariagen_data/anoph/sample_metadata.py

@@ -82,6 +82,7 @@ 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(
@@ -1487,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,
@@ -1514,20 +1519,45 @@ 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
 
-        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:
+            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
 
-        # Read the manifest into a pandas dataframe.
-        with self.open_file(path) as f:
-            df_cohorts = pd.read_csv(f, sep=",", na_values="")
+            self._cache_cohorts[cache_key] = df_cohorts
 
-        # Ensure all column names are lower case.
-        df_cohorts.columns = [c.lower() for c in df_cohorts.columns]  # type: ignore
+        if query is not None:
+            df_cohorts = df_cohorts.query(query)
+            df_cohorts = df_cohorts.reset_index(drop=True)
 
-        return df_cohorts
+        return df_cohorts.copy()
 
     @_check_types
     @doc(

tests/anoph/conftest.py

@@ -1408,23 +1408,32 @@ def write_metadata(
             df_coh_ds.to_csv(dst_path, index=False)
 
             # Create cohorts data by sampling from some real files.
-            src_path = (
-                self.fixture_dir
-                / "vo_agam_release_master_us_central1"
-                / "v3_cohorts"
-                / "cohorts_20230516"
-                / "cohorts_admin1_month.csv"
-            )
-            dst_path = (
-                self.bucket_path
-                / "v3_cohorts"
-                / "cohorts_20230516"
-                / "cohorts_admin1_month.csv"
-            )
-            dst_path.parent.mkdir(parents=True, exist_ok=True)
-            with open(src_path, mode="r") as src, open(dst_path, mode="w") as dst:
-                for line in src.readlines()[:5]:
-                    print(line, file=dst)
+            cohort_files = [
+                "cohorts_admin1_month.csv",
+                "cohorts_admin1_year.csv",
+                "cohorts_admin2_month.csv",
+            ]
+            for cohort_file in cohort_files:
+                src_path = (
+                    self.fixture_dir
+                    / "vo_agam_release_master_us_central1"
+                    / "v3_cohorts"
+                    / "cohorts_20230516"
+                    / cohort_file
+                )
+                if src_path.exists():
+                    dst_path = (
+                        self.bucket_path
+                        / "v3_cohorts"
+                        / "cohorts_20230516"
+                        / cohort_file
+                    )
+                    dst_path.parent.mkdir(parents=True, exist_ok=True)
+                    with open(src_path, mode="r") as src, open(
+                        dst_path, mode="w"
+                    ) as dst:
+                        for line in src.readlines()[:5]:
+                            print(line, file=dst)
 
             # Copy cohort GeoJSON fixtures.
             geojson_files = [

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

@@ -0,0 +1,5 @@
+cohort_id,cohort_size,country,country_alpha2,country_alpha3,taxon,year,admin1_name,admin1_iso,admin1_geoboundaries_shape_id,admin1_representative_longitude,admin1_representative_latitude
+AO-LUA_colu_2009,81,Angola,AO,AGO,coluzzii,2009,Luanda,AO-LUA,26408823B49174064004395,13.679075010193182,-9.592222213499952
+BF-01_arab_2008,1,Burkina Faso,BF,BFA,arabiensis,2008,Boucle du Mouhoun,BF-01,92566538B98190668782446,-3.592255305233366,12.479899304500035
+BF-01_colu_2008,4,Burkina Faso,BF,BFA,coluzzii,2008,Boucle du Mouhoun,BF-01,92566538B98190668782446,-3.592255305233366,12.479899304500035
+BF-02_colu_2011,18,Burkina Faso,BF,BFA,coluzzii,2011,Cascades,BF-02,92566538B44525923588019,-4.482809923408134,10.30846

tests/anoph/fixture/vo_agam_release_master_us_central1/v3_cohorts/cohorts_20230516/cohorts_admin2_month.csv

@@ -0,0 +1,3 @@
+cohort_id,cohort_size,country,country_alpha2,country_alpha3,taxon,year,quarter,month,admin1_name,admin1_iso,admin1_geoboundaries_shape_id,admin1_representative_longitude,admin1_representative_latitude,admin2_name,admin2_iso,admin2_geoboundaries_shape_id,admin2_representative_longitude,admin2_representative_latitude
+AO-LUA-LUA_colu_2009_04,81,Angola,AO,AGO,coluzzii,2009,2,4,Luanda,AO-LUA,26408823B49174064004395,13.679075010193182,-9.592222213499952,Luanda,AO-LUA-LUA,26408823B49174064004396,13.679075010193182,-9.592222213499952
+BF-01-BAN_arab_2008_11,1,Burkina Faso,BF,BFA,arabiensis,2008,4,11,Boucle du Mouhoun,BF-01,92566538B98190668782446,-3.592255305233366,12.479899304500035,Banwa,BF-01-BAN,92566538B98190668782447,-3.592255305233366,12.479899304500035

tests/anoph/test_sample_metadata.py

@@ -1446,6 +1446,47 @@ def cohort_data_expected_columns():
     }
 
 
+def cohort_data_admin1_year_expected_columns():
+    return {
+        "cohort_id": "O",
+        "cohort_size": "i",
+        "country": "O",
+        "country_alpha2": "O",
+        "country_alpha3": "O",
+        "taxon": "O",
+        "year": "i",
+        "admin1_name": "O",
+        "admin1_iso": "O",
+        "admin1_geoboundaries_shape_id": "O",
+        "admin1_representative_longitude": "f",
+        "admin1_representative_latitude": "f",
+    }
+
+
+def cohort_data_admin2_month_expected_columns():
+    return {
+        "cohort_id": "O",
+        "cohort_size": "i",
+        "country": "O",
+        "country_alpha2": "O",
+        "country_alpha3": "O",
+        "taxon": "O",
+        "year": "i",
+        "quarter": "i",
+        "month": "i",
+        "admin1_name": "O",
+        "admin1_iso": "O",
+        "admin1_geoboundaries_shape_id": "O",
+        "admin1_representative_longitude": "f",
+        "admin1_representative_latitude": "f",
+        "admin2_name": "O",
+        "admin2_iso": "O",
+        "admin2_geoboundaries_shape_id": "O",
+        "admin2_representative_longitude": "f",
+        "admin2_representative_latitude": "f",
+    }
+
+
 def validate_cohort_data(df, expected_columns):
     # Check column names.
     # Note: insertion order in dictionary keys is guaranteed since Python 3.7
@@ -1467,6 +1508,40 @@ def test_cohort_data(fixture, api):
     validate_cohort_data(df_cohorts, cohort_data_expected_columns())
 
 
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_data_admin1_year(fixture, api):
+    df_cohorts = api.cohorts("admin1_year")
+    validate_cohort_data(df_cohorts, cohort_data_admin1_year_expected_columns())
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_data_admin2_month(fixture, api):
+    df_cohorts = api.cohorts("admin2_month")
+    validate_cohort_data(df_cohorts, cohort_data_admin2_month_expected_columns())
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_data_invalid_cohort_set(fixture, api):
+    with pytest.raises(ValueError, match="is not a valid cohort set"):
+        api.cohorts("invalid_name")
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_data_with_query(fixture, api):
+    df_all = api.cohorts("admin1_month")
+    df_filtered = api.cohorts("admin1_month", query="country == 'Burkina Faso'")
+    assert len(df_filtered) > 0
+    assert (df_filtered["country"] == "Burkina Faso").all()
+    assert len(df_filtered) < len(df_all)
+
+
+@parametrize_with_cases("fixture,api", cases=case_ag3_sim)
+def test_cohort_data_cached(fixture, api):
+    df1 = api.cohorts("admin1_month")
+    df2 = api.cohorts("admin1_month")
+    assert_frame_equal(df1, df2)
+
+
 # ------------------------------------------------------------------
 # Tests for cohort_geometries()
 # ------------------------------------------------------------------