Errors in xroms.open_mfnetcdf step

[tlsw231]

I am trying to run an xroms workflow I have used successfully in the past. I upgraded recently to dask=2022.7.1 and xarray=2022.6.0 after which this is the first time I am returning to the same workflow. I am trying something like this, very basic stuff and early in the workflow before any real computations:

xrargs = dict(parallel='True',data_vars='minimal',
              coords='minimal',compat='override',drop_variables=droplist)

chunks = 'auto'

# Just reading in two files from the full list of files
ds = xroms.open_mfnetcdf(files[0:2],chunks=chunks,xrargs=xrargs)

The xroms.open_mfnetcdf step is throwing this error:

NotImplementedError: Cannot chunk along a core dimension for a grid ufunc which has a signature which includes one of the axis positions ['inner', 'outer']

I am certain this step was going through before the upgrades. Here is the full error message:

    812 ds = xr.open_mfdataset(files, chunks=chunks, **xrargsin)
    814 # modify Dataset with useful ROMS z coords and make xgcm grid operations usable.
--> 815 ds, grid = roms_dataset(ds, Vtransform=Vtransform, add_verts=add_verts, proj=proj)
    817 return ds

File ./xroms/xroms.py:194, in roms_dataset(ds, Vtransform, add_verts, proj)
    187 z_w0.attrs = {
    188     "long_name": "depth of W-points",
    189     "field": "z_w0, scalar",
    190     "units": "m",
    191 }
    193 ds.coords["z_w"] = xroms.order(z_w)
--> 194 ds.coords["z_w_u"] = grid.interp(ds.z_w, "X")
    195 ds.coords["z_w_u"].attrs = {
    196     "long_name": "depth of U-points on vertical W grid",
    197     "time": "ocean_time",
    198     "field": "z_w_u, scalar, series",
    199     "units": "m",
    200 }
    201 ds.coords["z_w_v"] = grid.interp(ds.z_w, "Y")

File /software/conda/envs/dask_2022/lib/python3.10/site-packages/xgcm/grid.py:2041, in Grid.interp(self, da, axis, **kwargs)
   1989 def interp(self, da, axis, **kwargs):
   1990     """
   1991     Interpolate neighboring points to the intermediate grid point along
   1992     this axis.
   (...)
   2039     >>> grid.interp(da, ["X", "Y"], periodic={"X": True, "Y": False})
   2040     """
-> 2041     return self._1d_grid_ufunc_dispatch("interp", da, axis, **kwargs)

File /software/conda/envs/dask_2022/lib/python3.10/site-packages/xgcm/grid.py:1836, in Grid._1d_grid_ufunc_dispatch(self, funcname, data, axis, to, keep_coords, metric_weighted, other_component, **kwargs)
   1833 else:
   1834     map_overlap = False
-> 1836 array = grid_ufunc(
   1837     self,
   1838     array,
   1839     axis=[(ax_name,)],
   1840     keep_coords=keep_coords,
   1841     dask=dask,
   1842     map_overlap=map_overlap,
   1843     other_component=other_component,
   1844     **remaining_kwargs,
   1845 )
   1847 if ax_metric_weighted:
   1848     metric = self.get_metric(array, ax_metric_weighted)

File /software/conda/envs/dask_2022/lib/python3.10/site-packages/xgcm/grid_ufunc.py:460, in GridUFunc.__call__(self, grid, axis, *args, **kwargs)
    458 map_overlap = kwargs.pop("map_overlap", self.map_overlap)
    459 pad_before_func = kwargs.pop("pad_before_func", self.pad_before_func)
--> 460 return apply_as_grid_ufunc(
    461     self.ufunc,
    462     *args,
    463     axis=axis,
    464     grid=grid,
    465     signature=self.signature,
    466     boundary_width=self.boundary_width,
    467     boundary=boundary,
    468     dask=dask,
    469     map_overlap=map_overlap,
    470     pad_before_func=pad_before_func,
    471     **kwargs,
    472 )

File /software/conda/envs/dask_2022/lib/python3.10/site-packages/xgcm/grid_ufunc.py:745, in apply_as_grid_ufunc(func, axis, grid, signature, boundary_width, boundary, fill_value, keep_coords, dask, map_overlap, pad_before_func, other_component, *args, **kwargs)
    742 # Maybe map function over chunked core dims using dask.array.map_overlap
    743 if map_overlap:
    744     # Disallow situations where shifting axis position would cause chunk size to change
--> 745     _check_if_length_would_change(sig)
    747     mapped_func = _map_func_over_core_dims(
    748         func,
    749         args,
   (...)
    753         out_dtypes,
    754     )
    755 else:

File /software/conda/envs/dask_2022/lib/python3.10/site-packages/xgcm/grid_ufunc.py:1002, in _check_if_length_would_change(signature)
    996 all_ax_positions = set(
    997     p
    998     for arg_ps in signature.in_ax_positions + signature.out_ax_positions
    999     for p in arg_ps
   1000 )
   1001 if any(pos in DISALLOWED_OVERLAP_POSITIONS for pos in all_ax_positions):
-> 1002     raise NotImplementedError(
   1003         "Cannot chunk along a core dimension for a grid ufunc which has a signature which "
   1004         f"includes one of the axis positions {DISALLOWED_OVERLAP_POSITIONS}"
   1005     )

NotImplementedError: Cannot chunk along a core dimension for a grid ufunc which has a signature which includes one of the axis positions ['inner', 'outer']

Here is the output from xr.show_versions():

If anybody is facing similar issues, it would be great to hear how you resolved it.

[dcherian]

This is an upstream xgcm issue.

[tlsw231]

Ah okay. Thanks for the prompt reply! For the time being, I will switch to my older conda environment.