[Bug]: xgcm vertical regridder throws error ValueError: dimension temp_unique on 0th function argument to apply_ufunc with dask='parallelized' consists of multiple chunks, but is also a core dimension.

[tomvothecoder]

What happened?

I am copying the example from "Remap cloud fraction from model hybrid coordinate to pressure levels" to the SciPy notebook (PR #658) to serve as an example of vertical regridding.

I download the datasets from ESGF via http and opened them using xc.open_dataset()/xc.open_mfdataset(). regridder.vertical() raises ValueError: dimension temp_unique on 0th function argument to apply_ufunc with dask='parallelized' consists of multiple chunks, but is also a core dimension..

This error is not raised when opening the datasets from ESGF via OPeNDAP (refer to MVCE).

What did you expect to happen? Are there are possible answers you came across?

Vertically regridding local .nc datasets should work the same as using the remote OPeNDAP dataset.

I'm not sure why there is a difference in behavior.

Minimal Complete Verifiable Example (MVCE)

# %%
import numpy as np

import xcdat as xc


# Build hybrid pressure coordinate
def hybrid_coordinate(p0, a, b, ps, **kwargs):
    return a * p0 + b * ps


# %%
## 1. Using dataset downloaded from ESGF via http -- raises ValueError
# Download link: https://esgf-data2.llnl.gov/thredds/fileServer/user_pub_work/CMIP6/CMIP/E3SM-Project/E3SM-2-0/historical/r1i1p1f1/Amon/cl/gr/v20220830/cl_Amon_E3SM-2-0_historical_r1i1p1f1_gr_185001-189912.nc
LOCAL_URL = "cl_Amon_E3SM-2-0_historical_r1i1p1f1_gr_185001-189912.nc"
ds_cl = xc.open_dataset(LOCAL_URL, chunks={"time": 4})
output_grid = xc.create_grid(z=xc.create_axis("lev", np.linspace(100000, 1, 13)))
pressure = hybrid_coordinate(**ds_cl.data_vars)
new_pressure_grid = xc.create_grid(z=xc.create_axis("lev", np.linspace(100000, 1, 13)))

output_cl = ds_cl.regridder.vertical(
    "cl", new_pressure_grid, method="linear", target_data=pressure
)

output_cl.compute()

# %%
## 2. Using dataset from ESGF OPeNDAP -- works fine.
REMOTE_URL = "https://esgf-data2.llnl.gov/thredds/dodsC/user_pub_work/CMIP6/CMIP/E3SM-Project/E3SM-2-0/historical/r1i1p1f1/Amon/cl/gr/v20220830/cl_Amon_E3SM-2-0_historical_r1i1p1f1_gr_185001-189912.nc"

ds_cl = xc.open_dataset(REMOTE_URL, chunks={"time": 4})
output_grid = xc.create_grid(z=xc.create_axis("lev", np.linspace(100000, 1, 13)))
pressure = hybrid_coordinate(**ds_cl.data_vars)
new_pressure_grid = xc.create_grid(z=xc.create_axis("lev", np.linspace(100000, 1, 13)))

output_cl = ds_cl.regridder.vertical(
    "cl", new_pressure_grid, method="linear", target_data=pressure
)
output_cl.compute()

Relevant log output

ValueError                                Traceback (most recent call last)
Cell In[31], line 1
----> 1 ds_vr = ds_3d.regridder.vertical(
      2     "cl", new_pressure_grid, method="linear", target_data=pressure_data
      3 )

File ~/repositories/xcdat/xcdat/regridder/accessor.py:281, in RegridderAccessor.vertical(self, data_var, output_grid, tool, **options)
    273 input_grid = _get_input_grid(
    274     self._ds,
    275     data_var,
   (...)
    278     ],
    279 )
    280 regridder = regrid_tool(input_grid, output_grid, **options)
--> 281 output_ds = regridder.vertical(data_var, self._ds)
    283 return output_ds

File ~/repositories/xcdat/xcdat/regridder/xgcm.py:179, in XGCMRegridder.vertical(self, data_var, ds)
    174     target_data = None
    176 # assumes new verical coordinate has been calculated and stored as `pressure`
    177 # TODO: auto calculate pressure reference http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#parametric-v-coord
    178 #       cf_xarray only supports two ocean s-coordinate and ocean_sigma_coordinate
--> 179 output_da = grid.transform(
    180     ds[data_var],
    181     "Z",
    182     output_coord_z,
    183     target_data=target_data,
    184     method=self._method,
    185     **self._extra_options,
    186 )
    188 # when the order of dimensions are mismatched, the output data will be
    189 # transposed to match the input dimension order
    190 if output_da.dims != ds[data_var].dims:

File /opt/miniconda3/envs/xcdat_scipy_2024/lib/python3.11/site-packages/xgcm/grid.py:2692, in Grid.transform(self, da, axis, target, **kwargs)
   2608 """Convert an array of data to new 1D-coordinates along `axis`.
   2609 The method takes a multidimensional array of data `da` and
   2610 transforms it onto another data_array `target_data` in the
   (...)
   2688 
   2689 """
   2691 ax = self.axes[axis]
-> 2692 return ax.transform(da, target, **kwargs)

File /opt/miniconda3/envs/xcdat_scipy_2024/lib/python3.11/site-packages/xgcm/grid.py:1080, in Axis.transform(self, da, target, target_data, method, mask_edges, bypass_checks, suffix)
   1076 if method == "linear" or method == "log":
   1077     target, target_dim, target_data = _parse_target(
   1078         target, target_dim, dim, target_data
   1079     )
-> 1080     out = linear_interpolation(
   1081         da,
   1082         target_data,
   1083         target,
   1084         dim,
   1085         dim,  # in this case the dimension of phi and theta are the same
   1086         target_dim,
   1087         mask_edges=mask_edges,
   1088         bypass_checks=bypass_checks,
   1089         logarithmic=(method == "log"),
   1090     )
   1091 elif method == "conservative":
   1092     # the conservative method requires `target_data` to be on the `outer` coordinate.
   1093     # If that is not the case (a very common use case like transformation on any tracer),
   1094     # we need to infer the boundary values (using the interp logic)
   1095     # for this method we need the `outer` position. Error out if its not there.
   1096     try:

File /opt/miniconda3/envs/xcdat_scipy_2024/lib/python3.11/site-packages/xgcm/transform.py:209, in input_handling.<locals>.wrapper_input_handling(*args, **kwargs)
    207 # Execute function with temporary names
    208 args = (phi, theta, target_theta_levels, temp_dim2, theta_dim, temp_dim)
--> 209 value = func(*args, **kwargs)
    211 # rename back to original name
    212 value = value.rename({temp_dim: target_dim})

File /opt/miniconda3/envs/xcdat_scipy_2024/lib/python3.11/site-packages/xgcm/transform.py:227, in linear_interpolation(phi, theta, target_theta_levels, phi_dim, theta_dim, target_dim, **kwargs)
    223 @input_handling
    224 def linear_interpolation(
    225     phi, theta, target_theta_levels, phi_dim, theta_dim, target_dim, **kwargs
    226 ):
--> 227     out = xr.apply_ufunc(
    228         interp_1d_linear,
    229         phi,
    230         theta,
    231         target_theta_levels,
    232         kwargs=kwargs,
    233         input_core_dims=[[phi_dim], [theta_dim], [target_dim]],
    234         output_core_dims=[[target_dim]],
    235         exclude_dims=set((phi_dim, theta_dim)),
    236         dask="parallelized",
    237         output_dtypes=[phi.dtype],
    238     )
    239     return out

File /opt/miniconda3/envs/xcdat_scipy_2024/lib/python3.11/site-packages/xarray/core/computation.py:1266, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, on_missing_core_dim, *args)
...
    775                     )
    776     dask_gufunc_kwargs["allow_rechunk"] = True
    778 output_sizes = dask_gufunc_kwargs.pop("output_sizes", {})

ValueError: dimension temp_unique on 0th function argument to apply_ufunc with dask='parallelized' consists of multiple chunks, but is also a core dimension. To fix, either rechunk into a single array chunk along this dimension, i.e., ``.chunk(dict(temp_unique=-1))``, or pass ``allow_rechunk=True`` in ``dask_gufunc_kwargs`` but beware that this may significantly increase memory usage.

Anything else we need to know?

No response

Environment

xcdat=0.7.1 and main
xarray=2024.5.0
xgcm=0.8.1

[tomvothecoder]

Hey @jasonb5 can you take a look at this by our next meeting on 6/17 (Wed)?

[jasonb5]

@tomvothecoder Just got around to this issue, I've identified the problem, but not sure of the cause. Basically calling open_dataset with a remote URL and local path cause the file to be chunked differently, even when chunks={"time": 4} is used. Remotely the file is chunked as (4, 72, 180, 360) which is correct, but locally the file is chunked as (4, 36, 90, 180). Since it's chunked along lev, lat, and lon, xgcm errors because those dimensions cannot be chunked as they are "core" or dependent. The simple fix was to add ds_cl = ds_cl.chunk({"time": 4, "lev": 72, "lat": 180, "lon": 360}) causing the file to be chunked in a way xgcm can handle.

I'm going to look at open_dataset and determine where the issue is there, might be an xarray issue though. I'll update once I have more.

[jasonb5]

@tomvothecoder Found the cause of the issue. Basically xarray implemented preferred_chunks for the netcdf4 backend. For non-contiguous files the the encoded chunking is used and merged with the desired chunking. In this case the encoded chunking is (1, 36, 90, 180) and the desired chunking was (4, , , ) when this was merged it resulted in (4, 36, 90, 180) for the chunk size. Seems the best way if you don't know whether the source is contiguous or not is to just specify all dimensions in the chunks argument.

The reason the OPeNDAP url works, it's a completely different backend that utilizes pydap. This backend doesn't supported preferred_chunks and considers the file one whole chunk and the updates this with the desired chunking.

[tomvothecoder]

Thank you @jasonb5. It's great to know how Xarray behaves with pre-chunked datasets and specifying an incompatible chunks argument.

The datasets used in our notebooks will be updated in PR #705, so this won't be an issue anymore.