update prepare_variable (#119)

* update prepare_variable

* add test_table_id

* clean up tests

* added test for units convert

* added tmp table creation

* API allows dicts as tables

* docstring updates

* added ntimes_passed to cmor write

* fixed for fx data

* added resample_kwargs, defaults to right

* fixed mapping_table check

* added flox method and checks

* update of dataset attributes and coords

* env updates

* Update whats_new.rst

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Lars Buntemeyer <lars.buntemeyer@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

ci/requirements/environment.yml

@@ -20,6 +20,7 @@ dependencies:
   - python-cdo
   - cmor
   - pint-xarray
+  - flox
 # for testing
   - pytest
   - pytest-cov

cordex/cmor/__init__.py

@@ -1,5 +1,5 @@
 from .cmor import cmorize_variable, prepare_variable
-from .config import set_options
+from .config import options, set_options
 from .utils import (
     mid_of_month,
     mid_of_season,
@@ -30,4 +30,5 @@
     "season_bounds",
     "to_cftime",
     "set_options",
+    "options",
 ]

cordex/cmor/cmor.py

@@ -1,5 +1,5 @@
-import json
 import os
+from os import path as op
 from warnings import warn
 
 import cf_xarray as cfxr
@@ -31,14 +31,18 @@
     _get_cfvarinfo,
     _get_cordex_pole,
     _get_pole,
+    _read_table,
     _strip_time_cell_method,
+    _tmp_table,
     mid_of_month,
     month_bounds,
     time_bounds_name,
 )
 
 xr.set_options(keep_attrs=True)
 
+flox_method = "blockwise"
+
 
 def resample_both_closed(ds, hfreq, op, **kwargs):
     rolling = getattr(ds.rolling(time=hfreq + 1, center=True), op)()
@@ -67,17 +71,18 @@ def _resample(
     ds, time, time_cell_method="point", label="left", time_offset=True, **kwargs
 ):
     """Resample a variable."""
-    # freq = "{}H".format(hfreq)
     if time_cell_method == "point":
-        return ds.resample(
-            time=time, label=label, **kwargs
-        ).nearest()  # .interpolate("nearest") # use as_freq?
+        return ds.resample(time=time, label=label, **kwargs).nearest(
+            method=flox_method
+        )  # .interpolate("nearest") # use as_freq?
     elif time_cell_method == "mean":
         if time_offset is True:
             loffset = _get_loffset(time)
         else:
             loffset = None
-        return ds.resample(time=time, label=label, loffset=loffset, **kwargs).mean()
+        return ds.resample(time=time, label=label, loffset=loffset, **kwargs).mean(
+            method=flox_method
+        )
     else:
         raise Exception("unknown time_cell_method: {}".format(time_cell_method))
 
@@ -130,7 +135,7 @@ def _get_time_axis_name(time_cell_method):
 
 def _define_axes(ds, table_id):
     if "CORDEX_domain" in ds.attrs:
-        grid = cordex_domain(ds.attrs["CORDEX_domain"], add_vertices=True)
+        grid = cordex_domain(ds.attrs["CORDEX_domain"], bounds=True)
         lon_vertices = grid.lon_vertices.to_numpy()
         lat_vertices = grid.lat_vertices.to_numpy()
     else:
@@ -222,7 +227,11 @@ def _cmor_write(da, table_id, cmorTime, cmorGrid, file_name=True):
     else:
         coords = [cmorTime, cmorGrid]
     cmor_var = cmor.variable(da.name, da.units, coords)
-    cmor.write(cmor_var, da.values)
+    if "time" in da.coords:
+        ntimes_passed = da.time.size
+    else:
+        ntimes_passed = None
+    cmor.write(cmor_var, da.to_numpy(), ntimes_passed=ntimes_passed)
     return cmor.close(cmor_var, file_name=file_name)
 
 
@@ -237,15 +246,14 @@ def _units_convert(da, cf_units, format=None):
     return da
 
 
-def _cf_units_convert(da, table_file, mapping_table={}):
+def _cf_units_convert(da, table, mapping_table={}):
     """Convert units.
 
     Convert units according to the rules in units_convert_rules dict.
     Maybe metpy can do this also: https://unidata.github.io/MetPy/latest/tutorials/unit_tutorial.html
 
     """
-    with open(table_file) as f:
-        table = json.load(f)
+
     if da.name in mapping_table:
         map_units = mapping_table[da.name].get("units")
         atr_units = da.attrs.get("units")
@@ -365,10 +373,13 @@ def _add_time_bounds(ds, cf_freq):
         ds = _add_month_bounds(ds)
     else:
         try:
-            ds = ds.convert_calendar(ds.time.dt.calendar).cf.add_bounds("time")
+            ds = ds.convert_calendar(
+                ds.time.dt.calendar, use_cftime=False
+            ).cf.add_bounds("time")
         except Exception:
             # wait for cftime arithemtics in xarry here:
             warn("could not add time bounds.")
+
     ds[time_bounds_name].encoding = ds.time.encoding
     ds.time.attrs.update({"bounds": time_bounds_name})
     return ds
@@ -383,7 +394,9 @@ def _adjust_frequency(ds, cf_freq, input_freq=None, time_cell_method=None):
     pd_freq = freq_map[cf_freq]
     if pd_freq != input_freq:
         warn("resampling input data from {} to {}".format(input_freq, pd_freq))
-        resample = _resample(ds, pd_freq, time_cell_method=time_cell_method)
+        resample = _resample(
+            ds, pd_freq, time_cell_method=time_cell_method, **options["resample_kwargs"]
+        )
         return resample
     return ds
 
@@ -395,6 +408,14 @@ def cmorize_cmor(
     if inpath is None:
         inpath = os.path.dirname(cmor_table)
 
+    dataset_table_json = dataset_table
+    cmor_table_json = cmor_table
+
+    if isinstance(dataset_table, dict):
+        dataset_table_json = _tmp_table(dataset_table)
+    if isinstance(cmor_table, dict):
+        cmor_table_json = _tmp_table(cmor_table)
+
     cfvarinfo = _get_cfvarinfo(out_name, cmor_table)
 
     if cfvarinfo is None:
@@ -403,7 +424,7 @@ def cmorize_cmor(
     time_cell_method = _strip_time_cell_method(cfvarinfo)
 
     table_ids = _setup(
-        dataset_table, cmor_table, grids_table=grids_table, inpath=inpath
+        dataset_table_json, cmor_table_json, grids_table=grids_table, inpath=inpath
     )
 
     cmorTime, cmorGrid = _define_grid(ds, table_ids, time_cell_method=time_cell_method)
@@ -422,17 +443,29 @@ def prepare_variable(
     input_freq=None,
     CORDEX_domain=None,
     time_units=None,
-    time_cell_method=None,
-    cf_freq=None,
     rewrite_time_axis=False,
+    use_cftime=False,
     squeeze=True,
 ):
     """prepares a variable for cmorization."""
+
+    ds = ds.copy(deep=False)
+
+    if isinstance(cmor_table, str):
+        cmor_table = _read_table(cmor_table)
+    cfvarinfo = _get_cfvarinfo(out_name, cmor_table)
+
+    cf_freq = cfvarinfo["frequency"]
+    time_cell_method = _strip_time_cell_method(cfvarinfo)
+
     if isinstance(ds, xr.DataArray):
         ds = ds.to_dataset()
 
+    # ensure that we propagate everything
+    # ds = xr.decode_cf(ds, decode_coords="all")
+
     # no mapping table provided, we assume datasets has already correct out_names and units.
-    if mapping_table is None:
+    if out_name not in mapping_table:
         try:
             var_ds = ds[[out_name]]
         except Exception:
@@ -448,11 +481,15 @@ def prepare_variable(
     if squeeze is True:
         var_ds = var_ds.squeeze(drop=True)
     if CORDEX_domain is not None:
+        var_ds.attrs["CORDEX_domain"] = CORDEX_domain
         var_ds = _crop_to_cordex_domain(var_ds, CORDEX_domain)
     if replace_coords is True:
-        grid = cordex_domain(CORDEX_domain)
+        grid = cordex_domain(CORDEX_domain, bounds=True)
         var_ds = var_ds.assign_coords(rlon=grid.rlon, rlat=grid.rlat)
         var_ds = var_ds.assign_coords(lon=grid.lon, lat=grid.lat)
+        var_ds = var_ds.assign_coords(
+            lon_vertices=grid.lon_vertices, lat_vertices=grid.lat_vertices
+        )
 
     if "time" in var_ds:
         # ensure cftime
@@ -464,6 +501,8 @@ def prepare_variable(
         if "time" not in ds.cf.bounds and time_cell_method == "mean":
             warn("adding time bounds")
             var_ds = _add_time_bounds(var_ds, cf_freq)
+        if use_cftime is False:
+            var_ds = var_ds.convert_calendar(ds.time.dt.calendar, use_cftime=False)
         var_ds = _set_time_encoding(var_ds, time_units, ds)
 
     if allow_units_convert is True:
@@ -477,7 +516,10 @@ def prepare_variable(
         warn("adding pole from archive specs: {}".format(CORDEX_domain))
         mapping = _get_cordex_pole(CORDEX_domain)
 
-    var_ds = xr.merge([var_ds, mapping])
+    if "time" in mapping.coords:
+        raise Exception("grid_mapping variable should have no time coordinate!")
+
+    var_ds[mapping.name] = mapping
 
     return var_ds
 
@@ -509,10 +551,10 @@ def cmorize_variable(
     out_name: str
         CF out_name of the variable that should be cmorized. The corresponding variable name
         in the dataset is looked up from the mapping_table if provided.
-    cmor_table : str
-        Filepath to cmor table.
-    dataset_table: str
-        Filepath to dataset cmor table.
+    cmor_table : str or dict
+        Cmor table dict of filepath to cmor table (json).
+    dataset_table: str or dict
+        Dataset table dict of filepath to dataset cmor table (json).
     mapping_table: dict
         Mapping of input variable names and meta data to CF out_name. Required if
         the variable name in the input dataset is not equal to out_name.
@@ -548,7 +590,7 @@ def cmorize_variable(
         Rewrite the time axis to CF compliant timestamps.
     outpath: str
         Root directory for output (can be either a relative or full path). This will override
-        the outpath defined in the dataset cmor input table.
+        the outpath defined in the dataset cmor input table (``dataset_table``).
     **kwargs:
         Argumets passed to prepare_variable.
 
@@ -573,14 +615,11 @@ def cmorize_variable(
     if inpath is None:
         inpath = os.path.dirname(cmor_table)
 
-    # get meta info from cmor table
-    cfvarinfo = _get_cfvarinfo(out_name, cmor_table)
-
-    if cfvarinfo is None:
-        raise Exception("{} not found in {}".format(out_name, cmor_table))
+    if op.isfile(dataset_table):
+        dataset_table = _read_table(dataset_table)
 
-    cf_freq = cfvarinfo["frequency"]
-    time_cell_method = _strip_time_cell_method(cfvarinfo)
+    if outpath:
+        dataset_table["outpath"] = outpath
 
     ds_prep = prepare_variable(
         ds,
@@ -590,8 +629,6 @@ def cmorize_variable(
         mapping_table=mapping_table,
         replace_coords=replace_coords,
         input_freq=input_freq,
-        cf_freq=cf_freq,
-        time_cell_method=time_cell_method,
         rewrite_time_axis=rewrite_time_axis,
         time_units=time_units,
         allow_resample=allow_resample,

cordex/cmor/config.py

@@ -2,7 +2,11 @@
 
 import numpy as np
 
-options = {"table_prefix": "CORDEX-CMIP6", "exit_control": "CMOR_NORMAL"}
+options = {
+    "table_prefix": "CORDEX-CMIP6",
+    "exit_control": "CMOR_NORMAL",
+    "resample_kwargs": {"closed": "left"},
+}
 
 
 # time offsets relative to left labeling for resampling.

cordex/cmor/utils.py

@@ -2,6 +2,7 @@
 """
 import datetime as dt
 import json
+import tempfile
 from warnings import warn
 
 import cftime as cfdt
@@ -290,7 +291,7 @@ def _encode_time(time):
     return xr.conventions.encode_cf_variable(time)
 
 
-def _read_cmor_table(table):
+def _read_table(table):
     return _read_json_file(table)
 
 
@@ -300,9 +301,41 @@ def _read_json_file(filename):
     return data
 
 
-def _get_cfvarinfo(cf_varname, table):
-    data = _read_cmor_table(table)
-    return data["variable_entry"].get(cf_varname, None)
+def _write_json_file(filename, data):
+    with open(filename, "w") as fp:
+        json.dump(data, fp, indent=4)
+    return filename
+
+
+def _get_cfvarinfo(out_name, table):
+    """Returns variable entry from cmor table"""
+    if isinstance(table, str):
+        table = _read_table(table)
+    info = table["variable_entry"].get(out_name, None)
+    if info is None:
+        raise Exception(
+            "{} not found in table {}".format(out_name, get_table_id(table))
+        )
+    return info
+
+
+def get_table_id(table):
+    """parse the table_id from a cmor table header"""
+    separator = " "
+    table_id = table["Header"].get("table_id", None)
+    if table_id is None:
+        raise Exception("no table_id in Header")
+    if separator in table_id:
+        return table_id.split(separator)[1]
+    return table_id
+
+
+def _tmp_table(table, format="json"):
+    """creates a temporay table json file"""
+    _, filename = tempfile.mkstemp()
+    warn(f"writing temporary table to {filename}")
+    if format == "json":
+        return _write_json_file(filename, table)
 
 
 def _get_time_cell_method(cf_varname, table):

docs/whats_new.rst

@@ -8,6 +8,15 @@ What's New
 
     import cordex
 
+v0.6.0 (Unreleased)
+-------------------
+
+Internal Changes
+~~~~~~~~~~~~~~~~
+
+- CMOR updates, including fixing of time step warnings and resampling options, includes options for using `flox <https://flox.readthedocs.io>`_ in resampling operations (:pull:`116`).
+
+
 v0.5.1 (2 March 2023)
 ---------------------
 

tests/__init__.py

@@ -20,3 +20,4 @@ def _importorskip(modname):
 has_regionmask, requires_regionmask = _importorskip("regionmask")
 has_xesmf, requires_xesmf = _importorskip("xesmf")
 has_geopandas, requires_geopandas = _importorskip("geopandas")
+has_pint_xarray, requires_pint_xarray = _importorskip("pint_xarray")