xroms.open_mfnetcdf fails with xgcm 0.8.0

[dylanschlichting]

I recently upgraded xgcm to ver 0.8.0 from 0.5.2 and xroms.open_mfnetcdf fails. If running 0.8.0 with the below code, the following error returns when horizontally interpolating the vertical coordinate. Not sure which xgcm update resulted in the error appearing.

import glob
import xroms 

path = glob.glob('/d1/shared/TXLA_ROMS/numerical_mixing/non-nest/ver1/1hr/ocean_avg_0000*.nc')
ds_avg = xroms.open_mfnetcdf(path, chunks = {'ocean_time': 1})
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Input In [1], in <cell line: 5>()
      2 import xroms 
      4 path = glob.glob('/d1/shared/TXLA_ROMS/numerical_mixing/non-nest/ver1/1hr/ocean_avg_0000*.nc')
----> 5 ds_avg = xroms.open_mfnetcdf(path, chunks = {'ocean_time': 1})

File ~/.local/lib/python3.9/site-packages/xroms-0.0.0-py3.9.egg/xroms/xroms.py:815, in open_mfnetcdf(files, chunks, xrargs, Vtransform, add_verts, proj)
    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 ~/.local/lib/python3.9/site-packages/xroms-0.0.0-py3.9.egg/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 ~/.conda/envs/xroms/lib/python3.9/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 ~/.conda/envs/xroms/lib/python3.9/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 ~/.conda/envs/xroms/lib/python3.9/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 ~/.conda/envs/xroms/lib/python3.9/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 ~/.conda/envs/xroms/lib/python3.9/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']

If run with 0.5.2, the following output is returned, which works as expected:

xarray.Dataset
Dimensions:
tracer: 5, s_rho: 30, s_w: 31, eta_rho: 191, xi_rho: 671, xi_u: 670, eta_v: 190, ocean_time: 1080
Coordinates: (31)
Data variables: (165)
Attributes: (32)

[kthyng]

@dylanschlichting thanks for reporting this. I'm glad you have a workaround in the meantime — I need to update the code to use xgcm's grid_ufuncs, from what I can tell.

[kthyng]

Hi @dylanschlichting! I have made a bunch of updates and this should work now. I've recently put out a new version to PyPI (v0.2.4, still coming through on conda-forge). Please check to see if this issue is still present in the new version so that over time we can work to address these. Thanks!

[dylanschlichting]

Hi @kthyng. I created a new conda environment, reinstalled xroms per the instructions, and tested this again with the exact same code as shown above. Everything works as expected with xarray ver. 2023.5.0 and xgcm 0.8.1, so I'll close this issue.

[kthyng]

Excellent, thank you!