feat: Allow running of a single slug

Author: lewisjared

scripts/fetch_test_data.py

@@ -133,6 +133,7 @@ def process_sample_data_request(
 DATASETS_TO_FETCH = [
     # # Example metric data
     CMIP6Request(
+        "example",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -144,6 +145,7 @@ def process_sample_data_request(
     ),
     # ESMValTool Climate at global warmings levels data
     CMIP6Request(
+        slug="esmvaltool-gwl",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -155,6 +157,7 @@ def process_sample_data_request(
     ),
     # ESMValTool Cloud radiative effects
     CMIP6Request(
+        slug="esmvaltool-cloud-radiative-effects",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -166,6 +169,7 @@ def process_sample_data_request(
     ),
     # ESMValTool ECS data
     CMIP6Request(
+        slug="esmvaltool-ecs",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -177,6 +181,7 @@ def process_sample_data_request(
     ),
     # ESMValTool TCR data
     CMIP6Request(
+        slug="esmvaltool-tcr",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -188,6 +193,7 @@ def process_sample_data_request(
     ),
     # ESMValTool TCRE data
     CMIP6Request(
+        slug="esmvaltool-tcre",
         facets=dict(
             source_id="MPI-ESM1-2-LR",
             frequency=["fx", "mon"],
@@ -198,6 +204,7 @@ def process_sample_data_request(
         time_span=("1850", "1915"),
     ),
     CMIP6Request(
+        slug="esmvaltool-extra",
         facets=dict(
             source_id="MPI-ESM1-2-LR",
             frequency=["fx", "mon"],
@@ -209,6 +216,7 @@ def process_sample_data_request(
     ),
     # ESMValTool ZEC data
     CMIP6Request(
+        slug="esmvaltool-zec",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -220,6 +228,7 @@ def process_sample_data_request(
     ),
     # ESMValTool Sea Ice Area Seasonal Cycle data
     CMIP6Request(
+        slug="esmvaltool-sea-ice",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -231,6 +240,7 @@ def process_sample_data_request(
     ),
     # ILAMB data
     CMIP6Request(
+        slug="ilamb",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -242,6 +252,7 @@ def process_sample_data_request(
     ),
     # ILAMB data, nbp requires a longer time span
     CMIP6Request(
+        slug="ilamb-nbp",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["mon"],
@@ -253,6 +264,7 @@ def process_sample_data_request(
     ),
     # IOMB data
     CMIP6Request(
+        slug="iomb",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -264,6 +276,7 @@ def process_sample_data_request(
     ),
     # PMP modes of variability data
     CMIP6Request(
+        slug="pmp-modes-of-variability",
         facets=dict(
             source_id="ACCESS-ESM1-5",
             frequency=["fx", "mon"],
@@ -274,6 +287,19 @@ def process_sample_data_request(
         remove_ensembles=False,
         time_span=("2000", "2025"),
     ),
+    # PMP ENSO data
+    CMIP6Request(
+        slug="pmp-enso",
+        facets=dict(
+            source_id="ACCESS-ESM1-5",
+            frequency=["fx", "mon"],
+            variable_id=["areacella", "ts", "tauu", "taux", "hfls", "hfss", "rlds", "rlus", "rsds", "rsus"],
+            experiment_id=["historical"],
+            variant_label=["r1i1p1f1"],
+        ),
+        remove_ensembles=False,
+        time_span=("2000", "2025"),
+    ),
     # # Obs4MIPs AIRS data
     # Obs4MIPsRequest(
     #     facets=dict(
@@ -295,11 +321,16 @@ def process_sample_data_request(
 def create_sample_data(
     decimate: bool = True,
     output: Path = OUTPUT_PATH,
+    slug: str | None = None,
 ) -> None:
     """Fetch and create sample datasets"""
     processed_datasets = pd.DataFrame(columns=["source_type", "key", "files", "time_start", "time_end"])
 
     for dataset_requested in DATASETS_TO_FETCH:
+        if slug is not None and dataset_requested.slug != slug:
+            logger.info(f"Skipping dataset {dataset_requested.slug} as it does not match the slug {slug}")
+            continue
+
         # Process the request
         new_datasets = process_sample_data_request(
             processed_datasets,

src/ref_sample_data/data_request/base.py

@@ -15,6 +15,7 @@ class DataRequest(Protocol):
     differently to generate the sample data.
     """
 
+    slug: str
     source_type: str
     time_span: tuple[str, str] | None = None
 

src/ref_sample_data/data_request/cmip6.py

@@ -67,7 +67,9 @@ class CMIP6Request(IntakeESGFDataRequest):
         "grid_label",
     )
 
-    def __init__(self, facets: dict[str, Any], remove_ensembles: bool, time_span: tuple[str, str] | None):
+    def __init__(
+        self, slug: str, facets: dict[str, Any], remove_ensembles: bool, time_span: tuple[str, str] | None
+    ):
         self.avail_facets = [
             "mip_era",
             "activity_drs",
@@ -81,7 +83,7 @@ def __init__(self, facets: dict[str, Any], remove_ensembles: bool, time_span: tu
             "version",
             "data_node",
         ]
-
+        self.slug = slug
         self.facets = facets
         self.remove_ensembles = remove_ensembles
         self.time_span = time_span

src/ref_sample_data/data_request/obs4mips.py

@@ -32,7 +32,9 @@ class Obs4MIPsRequest(IntakeESGFDataRequest):
         "grid_label",
     )
 
-    def __init__(self, facets: dict[str, Any], remove_ensembles: bool, time_span: tuple[str, str] | None):
+    def __init__(
+        self, slug: str, facets: dict[str, Any], remove_ensembles: bool, time_span: tuple[str, str] | None
+    ):
         self.avail_facets = [
             "activity_id",
             "institution_id",
@@ -44,6 +46,7 @@ def __init__(self, facets: dict[str, Any], remove_ensembles: bool, time_span: tu
             "data_node",
         ]
 
+        self.slug = slug
         self.facets = facets
         self.remove_ensembles = remove_ensembles
         self.time_span = time_span

src/ref_sample_data/data_request/obs4ref.py

@@ -20,6 +20,7 @@ class Obs4REFRequest(DataRequest):
     ```
     """
 
+    slug = "obs4ref"
     source_type = "obs4REF"
 
     def fetch_datasets(self) -> pd.DataFrame: