Mean clim unified workflow

conda-env/dev.yml

@@ -15,7 +15,7 @@ dependencies:
     - matplotlib >=3.7.1
     - eofs >=2.0.0
     - seaborn >=0.12.2
-    #- enso_metrics >=1.1.5
+    # - enso_metrics >=1.1.5
     - xcdat >=0.10.0
     - xmltodict >=0.13.0
     - setuptools >=67.7.2

doc/obs_info/catalogue/build_obs4MIPsClims-pcmdi.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python
 
+# Code History
+# ------------
 # 2025-01-22 Peter J Glckler: Created
 # 2025-02-13 Jiwoo Lee: Cleaned up
 # 2025-02-18 Jiwoo Lee: default dict to record default reference dataset per variable

pcmdi_metrics/io/__init__.py

@@ -1,5 +1,5 @@
 # init for pcmdi_metrics.io
-from .xcdat_openxml import xcdat_open  # noqa  # isort:skip
+from .xcdat_open import xcdat_open  # noqa  # isort:skip
 from .string_constructor import StringConstructor, fill_template  # noqa  # isort:skip
 from . import base  # noqa
 from . import cdp_io  # noqa

pcmdi_metrics/io/xcdat_dataset_io.py

@@ -539,6 +539,45 @@ def get_grid(
     return d[[lat_key, lon_key, lat_bnds_key, lon_bnds_key]]
 
 
+def get_grid_shape(
+    d: Union[xr.Dataset, xr.DataArray],
+) -> tuple:
+    """
+    Get the shape of the grid in terms of longitude and latitude dimensions.
+
+    Parameters
+    ----------
+    d : Union[xr.Dataset, xr.DataArray]
+        Input xarray dataset or data array containing longitude and latitude dimensions.
+
+    Returns
+    -------
+    tuple
+        A tuple representing the size of the longitude and latitude dimensions
+        in the form (longitude_size, latitude_size).
+
+    Notes
+    -----
+    The function retrieves the keys for longitude and latitude using the
+    `get_longitude_key` and `get_latitude_key` functions, then uses these keys to
+    obtain the sizes of the respective dimensions.
+
+    Examples
+    --------
+    >>> ds = xr.Dataset(...)
+    >>> get_grid_shape(ds)
+    (180, 90)
+
+    See Also
+    --------
+    get_longitude_key, get_latitude_key
+    """
+    lon_key = get_longitude_key(d)
+    lat_key = get_latitude_key(d)
+
+    return tuple(d.sizes[dim] for dim in [lon_key, lat_key])
+
+
 def get_calendar(d: Union[xr.Dataset, xr.DataArray]) -> str:
     """
     Get the calendar type from an xarray Dataset or DataArray.

pcmdi_metrics/io/xcdat_openxml.py → pcmdi_metrics/io/xcdat_open.py

@@ -6,6 +6,7 @@
 import xarray as xr
 import xcdat as xc
 import xmltodict
+import yaml
 
 from pcmdi_metrics.io.xcdat_dataset_io import get_calendar
 
@@ -22,7 +23,7 @@ def xcdat_open(
     Parameters
     ----------
     infile : Union[str, list]
-        list of string, or string, for path of file(s) to open using xcdat.
+        list of string, or string, for path of file(s) to open using xcdat. Can include wildcards if a string.
     data_var : str, optional
         key of the non-bounds data variable to keep in the Dataset, alongside any existing bounds data variables.
         By default None, which loads all data variables.
@@ -51,43 +52,35 @@ def xcdat_open(
     # Open an xml file
     >>> ds = xcdat_open('mydata.xml')
     """
-    if isinstance(infile, list) or "*" in infile:
+
+    def _open_with_fallback(open_func, infile, data_var, decode_times, chunks):
+        """Helper to try opening with decode_times, fall back if ValueError."""
         try:
-            ds = xc.open_mfdataset(
+            ds = open_func(
                 infile, data_var=data_var, decode_times=decode_times, chunks=chunks
             )
-        except (
-            ValueError
-        ):  # Could be due to non-cf-compliant calendar or other attribute
-            ds = xc.open_mfdataset(
-                infile, data_var=data_var, decode_times=False, chunks=chunks
-            )
+        except ValueError:
+            ds = open_func(infile, data_var=data_var, decode_times=False, chunks=chunks)
             ds = fix_noncompliant_attr(ds)
+        return ds
+
+    if isinstance(infile, list) or "*" in infile:
+        ds = _open_with_fallback(
+            xc.open_mfdataset, infile, data_var, decode_times, chunks
+        )
     else:
-        if infile.split(".")[-1].lower() == "xml":
-            try:
-                ds = _xcdat_openxml(
-                    infile, data_var=data_var, decode_times=decode_times, chunks=chunks
-                )
-            except (
-                ValueError
-            ):  # Could be due to non-cf-compliant calendar or other attribute
-                ds = _xcdat_openxml(
-                    infile, data_var=data_var, decode_times=False, chunks=chunks
-                )
-                ds = fix_noncompliant_attr(ds)
+        if infile.endswith(".xml"):
+            ds = _open_with_fallback(
+                _xcdat_openxml, infile, data_var, decode_times, chunks
+            )
+        elif infile.endswith(".yml") or infile.endswith(".yaml"):
+            ds = _open_with_fallback(
+                _xcdat_openyml, infile, data_var, decode_times, chunks
+            )
         else:
-            try:
-                ds = xc.open_dataset(
-                    infile, data_var=data_var, decode_times=decode_times, chunks=chunks
-                )
-            except (
-                ValueError
-            ):  # Could be due to non-cf-compliant calendar or other attribute
-                ds = xc.open_dataset(
-                    infile, data_var=data_var, decode_times=False, chunks=chunks
-                )
-                ds = fix_noncompliant_attr(ds)
+            ds = _open_with_fallback(
+                xc.open_dataset, infile, data_var, decode_times, chunks
+            )
 
     return ds.bounds.add_missing_bounds()
 
@@ -125,7 +118,8 @@ def _xcdat_openxml(
     data_var: str, optional
         key of the non-bounds data variable to keep in the Dataset, alongside any existing bounds data variables, by default None, which loads all data variables
     decode_times : bool, optional
-        If True, attempt to decode times encoded in the standard NetCDF datetime format into cftime.datetime objects. Otherwise, leave them encoded as numbers. This keyword may not be supported by all the backends, by default True
+        If True, attempt to decode times encoded in the standard NetCDF datetime format into cftime.datetime objects.
+        Otherwise, leave them encoded as numbers. This keyword may not be supported by all the backends, by default True
 
     Returns
     -------
@@ -150,3 +144,44 @@ def _xcdat_openxml(
         )
 
     return ds
+
+
+def _xcdat_openyml(
+    ymlfile: str, data_var: str = None, decode_times: bool = True, chunks=None
+) -> xr.Dataset:
+    """Open input file (yaml generated by PMP data curator)
+
+    Parameters
+    ----------
+    infile: str
+        path of yaml file to open using xcdat
+    data_var: str, optional
+        key of the non-bounds data variable to keep in the Dataset, alongside any existing bounds data variables, by default None, which loads all data variables
+    decode_times : bool, optional
+        If True, attempt to decode times encoded in the standard NetCDF datetime format into cftime.datetime objects.
+        Otherwise, leave them encoded as numbers. This keyword may not be supported by all the backends, by default True
+
+    Returns
+    -------
+    xr.Dataset
+        xarray dataset opened via xcdat
+    """
+    if not os.path.exists(ymlfile):
+        sys.exit(f"ERROR: File not exist: {ymlfile}")
+
+    # read yml file to python dictionary
+    with open(ymlfile, "r") as f:
+        yml_dict = yaml.safe_load(f)
+    directory_path = yml_dict["directory_path"]
+    netcdf_files = yml_dict["netcdf_files"]
+    files_to_open = [os.path.join(directory_path, f) for f in netcdf_files]
+
+    # check if files exist
+    for f in files_to_open:
+        if not os.path.isfile(f):
+            raise FileNotFoundError(f"File not found: {f}")
+
+    ds = xc.open_mfdataset(
+        files_to_open, data_var=data_var, decode_times=decode_times, chunks=chunks
+    )
+    return ds

pcmdi_metrics/mean_climate/lib/__init__.py

@@ -1,4 +1,4 @@
-from .calculate_climatology import calculate_climatology  # noqa
+from .calculate_climatology import calculate_climatology, is_4d_variable  # noqa
 from .compute_metrics import compute_metrics  # noqa
 from .compute_statistics import (  # noqa
     annual_mean,
@@ -19,7 +19,7 @@
 from .plot_clim_maps import plot_climatology, plot_climatology_diff
 
 from .create_mean_climate_parser import create_mean_climate_parser  # noqa
-from .load_and_regrid import load_and_regrid, extract_level  # noqa
+from .load_and_regrid import load_and_regrid, extract_level, extract_levels  # noqa
 from .mean_climate_metrics_to_json import mean_climate_metrics_to_json  # noqa
 from .pmp_parser import PMPParser  # noqa