Add Pangeo CMIP6 ArchiveIndex accessor (closes #284)

.github/workflows/python-app.yml

@@ -40,7 +40,7 @@ jobs:
         pip install -e packages/bundled_models/fourcastnext
         pip install -e .[test]
 
-        pytest -m="not noci" --cov=packages/data --cov=packages/utils --cov=packages/pipeline --cov=packages/training --cov=packages/zoo --cov=packages/bundled_models/fourcastnext --ignore=packages/nci_site_archive
+        pytest -m="not noci" --cov=packages/data --cov=packages/utils --cov=packages/pipeline --cov=packages/training --cov=packages/zoo --cov=packages/bundled_models/fourcastnext --cov=packages/third_party/pangeo_site_archive --ignore=packages/nci_site_archive
     - name: Coveralls GitHub Action
       uses: coverallsapp/github-action@v2.3.0
       with:

docs/data.md

@@ -40,6 +40,20 @@ The package `site_archive_nci` (which is present in some of the tutorials) is th
 
 Additionally, you can explore the [geonetwork](https://geonetwork.nci.org.au/geonetwork/srv/eng/catalog.search#/home) to find additional data sets of interest which you might like to utilise by writing a custom PyEarthTools data accessor.
 
+### Accessing CMIP6 climate model data from the cloud
+
+The package `site_archive_pangeo` (in `packages/third_party/pangeo_site_archive`) reads CMIP6 model output from the [Pangeo CMIP6 cloud archive](https://pangeo-data.github.io/pangeo-cmip6-cloud/accessing_data.html) on Google Cloud. No account, download or HPC access is needed; the data is opened lazily, so only the times you request are transferred.
+
+```python
+import pyearthtools.data
+from site_archive_pangeo import PangeoCMIP6
+
+index = PangeoCMIP6(["tas", "pr"], source_id="CanESM5", experiment_id="historical")
+ds = index.retrieve(pyearthtools.data.Petdt("2000-01"))
+```
+
+Each variable is looked up in the Pangeo catalogue and the newest version is used. If a variable exists on several grids, pass `grid_label` to choose one. Scalar coordinates that differ between variables (for example `height`, which is 2 m for `tas` but 10 m for `uas`) are dropped when the variables are combined.
+
 ### Connecting PyEarthTools to a new dataset in any HPC facility by coding a new data accessor
 
 For on-disk data access, you will need to create a new accessor based on the [`pyearthtools.data.ArchiveIndex`](project:./api/data/data_api.md#pyearthtools.data.indexes.ArchiveIndex) class (or [`pyearthtools.data.Index`](project:./api/data/data_api.md#pyearthtools.data.indexes.Index) for some use cases). Additional instructions for this still need to be written. In the meantime, refer to the [NCI site archive source code](https://github.com/ACCESS-Community-Hub/PyEarthTools/tree/develop/packages/nci_site_archive) for examples.

packages/third_party/pangeo_site_archive/README.md

@@ -0,0 +1,3 @@
+# PyEarthTools archive: Pangeo CMIP6
+
+Accessor for the Pangeo analysis-ready, cloud-optimised (ARCO) CMIP6 data, hosted as Zarr stores on Google Cloud.

packages/third_party/pangeo_site_archive/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=61", "setuptools-scm>=8.0"]
+
+[project]
+name = "PyEarthTools-archive-Pangeo"
+description = "PyEarthTools archive for the Pangeo CMIP6 analysis-ready, cloud-optimised (ARCO) data"
+requires-python = ">=3.11, <3.14"
+keywords = ["pyearthtools", "cmip6", "pangeo", "zarr", "climate"]
+classifiers = [
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+dependencies = [
+  "pyearthtools-data>=0.6.3",
+  "pandas",
+  "xarray",
+  "zarr",
+  "gcsfs",
+]
+dynamic = ["version", "readme"]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0", "pytest-cov", "black==26.3.1", "ruff"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.dynamic]
+readme = {file = ["README.md"], content-type = "text/markdown"}
+
+[project.urls]
+homepage = "https://pyearthtools.readthedocs.io/en/latest/"
+documentation = "https://pyearthtools.readthedocs.io/en/latest/"
+repository = "https://github.com/ACCESS-Community-Hub/PyEarthTools"
+
+[tool.isort]
+profile = "black"
+
+[tool.black]
+line-length = 120

packages/third_party/pangeo_site_archive/src/site_archive_pangeo/__init__.py

@@ -0,0 +1,21 @@
+"""
+Pangeo CMIP6 archive for PyEarthTools.
+
+| Name                                            | Description                                         |
+| :---                                            | ----:                                               |
+| [PangeoCMIP6][site_archive_pangeo.PangeoCMIP6]  | CMIP6 analysis-ready, cloud-optimised (ARCO) data   |
+"""
+
+from pyearthtools.data.archive import register_archive
+
+import site_archive_pangeo
+from site_archive_pangeo.pangeo_cmip6 import PangeoCMIP6
+
+register_archive("pangeo")(site_archive_pangeo)
+
+try:
+    from ._version import __version__  # type: ignore
+except ImportError:
+    __version__ = "999"
+
+__all__ = ["PangeoCMIP6"]

packages/third_party/pangeo_site_archive/src/site_archive_pangeo/pangeo_cmip6.py

@@ -0,0 +1,188 @@
+"""
+Pangeo CMIP6 accessor.
+
+Reads the Pangeo analysis-ready, cloud-optimised (ARCO) CMIP6 data, which is published as one Zarr store
+per variable and described by a single catalog CSV.
+"""
+
+import os
+from typing import Any, ClassVar
+
+import pandas
+import pyearthtools.data
+import xarray
+from pyearthtools.data import Petdt
+from pyearthtools.data.archive import register_archive
+from pyearthtools.data.indexes import ArchiveIndex
+from pyearthtools.data.transforms import Transform, TransformCollection
+
+DEFAULT_CATALOG = "https://storage.googleapis.com/cmip6/pangeo-cmip6.csv"
+
+_CATALOG_COLUMNS = [
+    "source_id",
+    "experiment_id",
+    "member_id",
+    "table_id",
+    "variable_id",
+    "grid_label",
+    "zstore",
+    "version",
+]
+
+
+@register_archive(
+    "PangeoCMIP6",
+    sample_kwargs={"variables": ["tas"], "source_id": "CanESM5", "experiment_id": "historical"},
+)
+class PangeoCMIP6(ArchiveIndex):
+    """
+    Data accessor for the Pangeo CMIP6 analysis-ready, cloud-optimised (ARCO) data.
+
+    A catalog CSV lists the Zarr store of every variable of every model run. This accessor selects one
+    model run (``source_id``, ``experiment_id``, ``member_id``, ``table_id``) and opens the store of each
+    requested variable lazily, so only the parts of the data that are used are downloaded.
+
+    A store holds all times, so ``filesystem`` ignores the requested time and the time is selected after
+    loading.
+
+    Examples:
+        >>> index = PangeoCMIP6(["tas"], source_id="CanESM5", experiment_id="historical")
+        >>> index.retrieve("2000-01")
+    """
+
+    _OPEN_ARGS: ClassVar[dict[str, Any]] = {"consolidated": True, "storage_options": {"token": "anon"}}
+
+    def __init__(
+        self,
+        variables: list[str] | str,
+        *,
+        source_id: str,
+        experiment_id: str,
+        member_id: str = "r1i1p1f1",
+        table_id: str = "Amon",
+        grid_label: str | None = None,
+        catalog: str | os.PathLike = DEFAULT_CATALOG,
+        transforms: Transform | TransformCollection | None = None,
+    ):
+        """
+        Set up the Pangeo CMIP6 accessor.
+
+        The catalog is not read here, only when data is first requested.
+
+        Args:
+            variables: Variables to retrieve, e.g. ``["tas", "pr"]``.
+            source_id: Model name, e.g. ``"CanESM5"``.
+            experiment_id: Experiment, e.g. ``"historical"`` or ``"ssp585"``.
+            member_id: Ensemble member. Defaults to ``"r1i1p1f1"``.
+            table_id: Frequency and realm table, e.g. ``"Amon"`` for monthly atmosphere. Defaults to ``"Amon"``.
+            grid_label: Grid, e.g. ``"gn"`` (native) or ``"gr"`` (regridded). Needed when a variable exists
+                on more than one grid. Defaults to None.
+            catalog: Location of the Pangeo catalog CSV, a URL or a local path.
+            transforms: Extra transforms to apply after loading.
+
+        Raises:
+            ValueError: If no variables are given.
+        """
+        self._variables = [variables] if isinstance(variables, str) else list(variables)
+        if not self._variables:
+            raise ValueError("At least one variable is required.")
+
+        self._query = {
+            "source_id": source_id,
+            "experiment_id": experiment_id,
+            "member_id": member_id,
+            "table_id": table_id,
+        }
+        self._grid_label = grid_label
+        self._catalog_src = catalog
+        self._catalog = None  # read lazily, see _load_catalog
+
+        extra_transforms = TransformCollection() if transforms is None else TransformCollection(transforms)
+        super().__init__(
+            transforms=TransformCollection([pyearthtools.data.transforms.variables.Trim(self._variables)])
+            + extra_transforms,
+        )
+        self.record_initialisation()
+
+    def _load_catalog(self) -> pandas.DataFrame:
+        """Read the catalog once, keeping only the rows of the requested variables and model run."""
+        if self._catalog is None:
+            catalog = self.catalog(self._catalog_src)
+            catalog = catalog[catalog["variable_id"].isin(self._variables)]
+            for column, value in self._query.items():
+                catalog = catalog[catalog[column] == value]
+            self._catalog = catalog.reset_index(drop=True)
+        return self._catalog
+
+    def filesystem(self, querytime: str | Petdt) -> dict[str, str]:
+        """
+        Find the Zarr store of each variable.
+
+        Args:
+            querytime: Ignored, as a store holds all times.
+
+        Raises:
+            ValueError: If a variable has no store, or exists on several grids and no ``grid_label`` was given.
+
+        Returns:
+            Store URL of each variable, using the latest version in the catalog.
+        """
+        catalog = self._load_catalog()
+        stores = {}
+        for variable in self._variables:
+            rows = catalog[catalog["variable_id"] == variable]
+            if self._grid_label is not None:
+                rows = rows[rows["grid_label"] == self._grid_label]
+            if rows.empty:
+                raise ValueError(f"No Pangeo CMIP6 store for variable {variable!r} matching {self._query}")
+            grids = sorted(rows["grid_label"].unique())
+            if len(grids) > 1:
+                raise ValueError(
+                    f"Variable {variable!r} exists on several grids {grids}; pass grid_label to choose one."
+                )
+            stores[variable] = str(rows.sort_values("version").iloc[-1]["zstore"])
+        return stores
+
+    def load(self, files: dict[str, str], **kwargs) -> xarray.Dataset:
+        """
+        Open the store of each variable lazily, keep only that variable, and merge them.
+
+        Args:
+            files: Store URL of each variable, as found by ``filesystem``.
+            **kwargs: Passed to ``xarray.open_zarr``.
+
+        Returns:
+            Dataset holding all times of all the variables. Scalar coordinates that differ between the variables
+            are dropped, e.g. ``height``, which is 2 m for ``tas`` but 10 m for ``uas``.
+        """
+        open_args = {**self._OPEN_ARGS, **kwargs}
+        datasets = [xarray.open_zarr(url, **open_args)[[variable]] for variable, url in files.items()]
+        return xarray.merge(datasets, join="exact", compat="minimal")
+
+    @staticmethod
+    def catalog(catalog: str | os.PathLike = DEFAULT_CATALOG) -> pandas.DataFrame:
+        """Return the Pangeo CMIP6 catalogue, so the available data can be discovered.
+
+        Args:
+            catalog: CSV catalogue to read. Defaults to the public Pangeo CMIP6 catalogue.
+
+        Returns:
+            DataFrame with one row per variable of one model run, with the columns `source_id`,
+            `experiment_id`, `member_id`, `table_id`, `variable_id`, `grid_label`, `zstore` and `version`.
+        """
+        return pandas.read_csv(catalog, usecols=_CATALOG_COLUMNS)
+
+    def exists(self, querytime: str | Petdt) -> bool:
+        """Whether a store is found for every variable. Stores are URLs, so they are not checked on disk."""
+        try:
+            self.search(querytime)
+        except ValueError:
+            return False
+        return True
+
+    @property
+    def _desc_(self):
+        return {
+            "singleline": "Pangeo CMIP6 analysis-ready, cloud-optimised data accessor",
+            "Documentation": "https://pangeo-data.github.io/pangeo-cmip6-cloud/accessing_data.html",
+        }

packages/third_party/pangeo_site_archive/tests/data/mini_catalog.csv

@@ -0,0 +1,8 @@
+activity_id,institution_id,source_id,experiment_id,member_id,table_id,variable_id,grid_label,zstore,dcpp_init_year,version
+CMIP,TEST-INST,TEST-MODEL,historical,r1i1p1f1,Amon,tas,gn,gs://test-bucket/tas/gn/v20190101/,,20190101
+CMIP,TEST-INST,TEST-MODEL,historical,r1i1p1f1,Amon,tas,gn,gs://test-bucket/tas/gn/v20200101/,,20200101
+CMIP,TEST-INST,TEST-MODEL,historical,r1i1p1f1,Amon,pr,gn,gs://test-bucket/pr/gn/v20190101/,,20190101
+CMIP,TEST-INST,TEST-MODEL,historical,r1i1p1f1,Amon,psl,gn,gs://test-bucket/psl/gn/v20190101/,,20190101
+CMIP,TEST-INST,TEST-MODEL,historical,r1i1p1f1,Amon,psl,gr,gs://test-bucket/psl/gr/v20190101/,,20190101
+ScenarioMIP,TEST-INST,TEST-MODEL,ssp585,r1i1p1f1,Amon,tas,gn,gs://test-bucket/ssp/tas/gn/v20190101/,,20190101
+CMIP,TEST-INST,TEST-MODEL,historical,r1i1p1f1,Amon,uas,gn,gs://test-bucket/uas/gn/v20190101/,,20190101