feat: Add obs4ref datasets

Author: lewisjared

scripts/fetch_test_data.py

@@ -2,45 +2,16 @@
 from pathlib import Path
 from typing import Annotated
 
-import pandas as pd
 import pooch
 import typer
 import xarray as xr
 
-from ref_sample_data import CMIP6Request, DataRequest, Obs4MIPsRequest
+from ref_sample_data import CMIP6Request, DataRequest, Obs4MIPsRequest, Obs4REFRequest
 
 OUTPUT_PATH = Path("data")
 app = typer.Typer()
 
 
-def deduplicate_datasets(datasets: pd.DataFrame) -> pd.DataFrame:
-    """
-    Deduplicate a dataset collection.
-
-    Uses the metadata from the first dataset in each group,
-    but expands the time range to the min/max timespan of the group.
-
-    Parameters
-    ----------
-    datasets
-        The dataset collection
-
-    Returns
-    -------
-    pd.DataFrame
-        The deduplicated dataset collection spanning the times requested
-    """
-
-    def _deduplicate_group(group: pd.DataFrame) -> pd.DataFrame:
-        first = group.iloc[0].copy()
-        first.time_start = group.time_start.min()
-        first.time_end = group.time_end.max()
-
-        return first
-
-    return datasets.groupby("key").apply(_deduplicate_group, include_groups=False).reset_index()
-
-
 def process_sample_data_request(
     request: DataRequest, decimate: bool, output_directory: Path, quiet: bool
 ) -> None:
@@ -61,7 +32,6 @@ def process_sample_data_request(
         Whether to suppress progress messages
     """
     datasets = request.fetch_datasets()
-    datasets = deduplicate_datasets(datasets)
 
     for _, dataset in datasets.iterrows():
         for ds_filename in dataset["files"]:
@@ -183,6 +153,8 @@ def process_sample_data_request(
         remove_ensembles=False,
         time_span=("2002", "2016"),
     ),
+    # All unpublished obs4mips datasets
+    Obs4REFRequest(),
 ]
 
 

src/ref_sample_data/__init__.py

@@ -10,5 +10,6 @@
 from .data_request.base import DataRequest
 from .data_request.cmip6 import CMIP6Request
 from .data_request.obs4mips import Obs4MIPsRequest
+from .data_request.obs4ref import Obs4REFRequest
 
-__all__ = ["CMIP6Request", "DataRequest", "Obs4MIPsRequest"]
+__all__ = ["CMIP6Request", "DataRequest", "Obs4MIPsRequest", "Obs4REFRequest"]

src/ref_sample_data/data_request/base.py

@@ -34,6 +34,34 @@ def generate_filename(
         ...
 
 
+def _deduplicate_datasets(datasets: pd.DataFrame) -> pd.DataFrame:
+    """
+    Deduplicate a dataset collection.
+
+    Uses the metadata from the first dataset in each group,
+    but expands the time range to the min/max timespan of the group.
+
+    Parameters
+    ----------
+    datasets
+        The dataset collection
+
+    Returns
+    -------
+    pd.DataFrame
+        The deduplicated dataset collection spanning the times requested
+    """
+
+    def _deduplicate_group(group: pd.DataFrame) -> pd.DataFrame:
+        first = group.iloc[0].copy()
+        first.time_start = group.time_start.min()
+        first.time_end = group.time_end.max()
+
+        return first
+
+    return datasets.groupby("key").apply(_deduplicate_group, include_groups=False).reset_index()
+
+
 class IntakeESGFDataRequest(DataRequest):
     """
     A data request that fetches datasets from ESGF using intake-esgf.
@@ -56,4 +84,4 @@ def fetch_datasets(self) -> pd.DataFrame:
         if self.time_span:
             merged_df["time_start"] = self.time_span[0]
             merged_df["time_end"] = self.time_span[1]
-        return merged_df
+        return _deduplicate_datasets(merged_df)

src/ref_sample_data/data_request/cmip6.py

@@ -94,8 +94,6 @@ def decimate_dataset(self, dataset: xr.Dataset) -> xr.Dataset | None:
         ----------
         dataset
             The dataset to downscale
-        time_span
-            The time span to extract from a dataset
 
         Returns
         -------

src/ref_sample_data/data_request/obs4mips.py

@@ -73,8 +73,6 @@ def decimate_dataset(self, dataset: xr.Dataset) -> xr.Dataset | None:
         ----------
         dataset
             The dataset to downscale
-        time_span
-            The time span to extract from a dataset
 
         Returns
         -------

src/ref_sample_data/data_request/obs4ref.py

@@ -0,0 +1,88 @@
+import pathlib
+from pathlib import Path
+
+import pandas as pd
+import xarray as xr
+from cmip_ref.dataset_registry import build_reference_data_registry
+
+from ref_sample_data.data_request.base import DataRequest
+from ref_sample_data.resample import decimate_curvilinear, decimate_rectilinear
+
+
+class Obs4REFRequest(DataRequest):
+    """
+    Fetch the unpublished Obs4MIPs datasets from the PMP registry
+
+    This includes all files that would be downloaded if you ran:
+    ```
+    ref datasets fetch-obs4ref-data --output-data ...
+    ```
+    """
+
+    def fetch_datasets(self) -> pd.DataFrame:
+        """
+        Fetch the datasets from the source
+
+        Returns a dataframe of the metadata and paths to the fetched datasets.
+        """
+        registry = build_reference_data_registry()
+
+        datasets = []
+        for key in registry.registry.keys():
+            dataset_path = registry.fetch(key)
+            datasets.append(
+                {
+                    "key": key,
+                    "files": [dataset_path],
+                }
+            )
+        return pd.DataFrame(datasets)
+
+    def decimate_dataset(self, dataset: xr.Dataset) -> xr.Dataset | None:
+        """
+        Downscale the dataset to a smaller size.
+
+        Parameters
+        ----------
+        dataset
+            The dataset to downscale
+
+        Returns
+        -------
+        xr.Dataset
+            The downscaled dataset
+        """
+        has_latlon = "lat" in dataset.dims and "lon" in dataset.dims
+        has_ij = "i" in dataset.dims and "j" in dataset.dims
+
+        if has_latlon:
+            assert len(dataset.lat.dims) == 1 and len(dataset.lon.dims) == 1
+
+            result = decimate_rectilinear(dataset)
+        elif has_ij:
+            # 2d curvilinear grid (generally ocean variables)
+            result = decimate_curvilinear(dataset)
+        else:
+            raise ValueError("Cannot decimate this grid: too many dimensions")
+
+        return result
+
+    def generate_filename(self, metadata: pd.Series, ds: xr.Dataset, ds_filename: pathlib.Path) -> Path:
+        """
+        Create the output filename for the dataset.
+
+        Parameters
+        ----------
+        metadata
+            Metadata from the file
+        ds
+            Loaded dataset
+
+        ds_filename:
+            Filename of the dataset (Unused)
+
+        Returns
+        -------
+            The output filename
+        """
+        return Path("obs4REF") / metadata.key