Misleading error messages in the plotting functions

I've been having some problems with the plot function in pyicon, with the error messages being given not really pointing to the problem. I've detailed some of these errors below. I'll submit a pull request to try and improve the error messages later this week.

Producing the errors

Problem 1

import xarray as xr
import pyicon as pyic

ds = xr.open_mfdataset(
    "/home/m/m301014/work.me/exp.TropAtl/data/raw/exp.TropAtl.r2b6/exp.TropAtl.r2b6_production_1955????T000000Z.nc",
    decode_cf=False,
)

ds["vort"].isel(time=10, depth=0).pyic.plot()

File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_accessor.py:26, in pyiconDataArray.plot(self, **kwargs)
23 @functools.wraps(pyic.plot, assigned=("__doc__", "__anotations__"))
24 def plot(self, **kwargs):
25     da = self._obj
---> 26     ax, hm = pyic.plot(da, **kwargs)
27     return ax, hmFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_plotting.py:1966, in plot(data, Plot, ax, cax, asp, fig_size_fac, mask_data, logplot, lon_reg, lat_reg, central_longitude, clim, cmap, conts, contfs, clevs, use_pcol_or_contf, contcolor, cincr, clabel, cbticks, xlabel, ylabel, xticks, yticks, template, cbar_str, cbar_pos, title_right, title_left, title_center, projection, coastlines_color, land_facecolor, axes_facecolor, noland, do_plot_settings, do_xyticks, do_gridlines, gname, fpath_tgrid, plot_method, grid_type, res, fpath_ckdtree, coordinates, lonlat_for_mask)
1963   fpath_ckdtree = f'{path_grid}/{gname}/ckdtree/rectgrids/{gname}_res{res:3.2f}_180W-180E_90S-90N.nc'
1965 if grid_type=='auto':
-> 1966   if Dgrid["name"].startswith("healpix"):
1967     grid_type = 'healpix'
1968   else:UnboundLocalError: cannot access local variable 'Dgrid' where it is not associated with a value

This is fixed by including grid_type="native" in the function call. But then we get a new error....

Problem 2

ds["vort"].isel(time=10, depth=0).pyic.plot(
    grid_type="native",
)

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/file_manager.py:211, in CachingFileManager._acquire_with_cache_info(self, needs_lock)
210 try:
--> 211     file = self._cache[self._key]
212 except KeyError:File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/lru_cache.py:56, in LRUCache.__getitem__(self, key)
55 with self._lock:
---> 56     value = self._cache[key]
57     self._cache.move_to_end(key)KeyError: [<class 'netCDF4._netCDF4.Dataset'>, ('/work/mh0033/m300602/icon/grids/none/ckdtree/rectgrids/none_res0.30_180W-180E_90S-90N.nc',), 'r', (('clobber', True), ('diskless', False), ('format', 'NETCDF4'), ('persist', False)), 'ab78d15b-6e6b-4165-a718-d702fcc4afe4']During handling of the above exception, another exception occurred:FileNotFoundError                         Traceback (most recent call last)
Cell In[242], line 1
----> 1 ds["vort"].isel(time=10, depth=0).pyic.plot(
2     grid_type="native"
3 )File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_accessor.py:26, in pyiconDataArray.plot(self, **kwargs)
23 @functools.wraps(pyic.plot, assigned=("__doc__", "__anotations__"))
24 def plot(self, **kwargs):
25     da = self._obj
---> 26     ax, hm = pyic.plot(da, **kwargs)
27     return ax, hmFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_plotting.py:1994, in plot(data, Plot, ax, cax, asp, fig_size_fac, mask_data, logplot, lon_reg, lat_reg, central_longitude, clim, cmap, conts, contfs, clevs, use_pcol_or_contf, contcolor, cincr, clabel, cbticks, xlabel, ylabel, xticks, yticks, template, cbar_str, cbar_pos, title_right, title_left, title_center, projection, coastlines_color, land_facecolor, axes_facecolor, noland, do_plot_settings, do_xyticks, do_gridlines, gname, fpath_tgrid, plot_method, grid_type, res, fpath_ckdtree, coordinates, lonlat_for_mask)
1989 #if 'grid_mapping' in list(data.attrs):
1990 #  plot_method = 'healpix'
1991
1992 # --- interpolate and cut to region
1993 if grid_type=='native' and plot_method=='nn':
-> 1994   datai = interp_to_rectgrid_xr(data.compute(), fpath_ckdtree, lon_reg=lon_reg, lat_reg=lat_reg, coordinates=coordinates)
1995   lon = datai.lon
1996   lat = datai.latFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_tb.py:241, in interp_to_rectgrid_xr(arr, fpath_ckdtree, lon_reg, lat_reg, coordinates, radius_of_influence, compute, mask_out_of_range, mask_out_of_range_before)
238   print(fpath_ckdtree)
240 # --- load interpolation indices
--> 241 ds_ckdt = xr.open_dataset(fpath_ckdtree)
242 if ('clon' in coordinates) or (coordinates==''):
243   inds = ds_ckdt.ickdtree_cFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/api.py:572, in open_dataset(filename_or_obj, engine, chunks, cache, decode_cf, mask_and_scale, decode_times, decode_timedelta, use_cftime, concat_characters, decode_coords, drop_variables, inline_array, chunked_array_type, from_array_kwargs, backend_kwargs, **kwargs)
560 decoders = _resolve_decoders_kwargs(
561     decode_cf,
562     open_backend_dataset_parameters=backend.open_dataset_parameters,
(...)
568     decode_coords=decode_coords,
569 )
571 overwrite_encoded_chunks = kwargs.pop("overwrite_encoded_chunks", None)
--> 572 backend_ds = backend.open_dataset(
573     filename_or_obj,
574     drop_variables=drop_variables,
575     **decoders,
576     **kwargs,
577 )
578 ds = _dataset_from_backend_dataset(
579     backend_ds,
580     filename_or_obj,
(...)
590     **kwargs,
591 )
592 return dsFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/netCDF4_.py:593, in NetCDF4BackendEntrypoint.open_dataset(self, filename_or_obj, mask_and_scale, decode_times, concat_characters, decode_coords, drop_variables, use_cftime, decode_timedelta, group, mode, format, clobber, diskless, persist, lock, autoclose)
572 def open_dataset(  # type: ignore[override]  # allow LSP violation, not supporting **kwargs
573     self,
574     filename_or_obj: str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore,
(...)
590     autoclose=False,
591 ) -> Dataset:
592     filename_or_obj = _normalize_path(filename_or_obj)
--> 593     store = NetCDF4DataStore.open(
594         filename_or_obj,
595         mode=mode,
596         format=format,
597         group=group,
598         clobber=clobber,
599         diskless=diskless,
600         persist=persist,
601         lock=lock,
602         autoclose=autoclose,
603     )
605     store_entrypoint = StoreBackendEntrypoint()
606     with close_on_error(store):File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/netCDF4_.py:400, in NetCDF4DataStore.open(cls, filename, mode, format, group, clobber, diskless, persist, lock, lock_maker, autoclose)
394 kwargs = dict(
395     clobber=clobber, diskless=diskless, persist=persist, format=format
396 )
397 manager = CachingFileManager(
398     netCDF4.Dataset, filename, mode=mode, kwargs=kwargs
399 )
--> 400 return cls(manager, group=group, mode=mode, lock=lock, autoclose=autoclose)File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/netCDF4_.py:347, in NetCDF4DataStore.__init__(self, manager, group, mode, lock, autoclose)
345 self._group = group
346 self._mode = mode
--> 347 self.format = self.ds.data_model
348 self._filename = self.ds.filepath()
349 self.is_remote = is_remote_uri(self._filename)File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/netCDF4_.py:409, in NetCDF4DataStore.ds(self)
407 @property
408 def ds(self):
--> 409     return self._acquire()File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/netCDF4_.py:403, in NetCDF4DataStore._acquire(self, needs_lock)
402 def _acquire(self, needs_lock=True):
--> 403     with self._manager.acquire_context(needs_lock) as root:
404         ds = _nc4_require_group(root, self._group, self._mode)
405     return dsFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/contextlib.py:137, in _GeneratorContextManager.__enter__(self)
135 del self.args, self.kwds, self.func
136 try:
--> 137     return next(self.gen)
138 except StopIteration:
139     raise RuntimeError("generator didn't yield") from NoneFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/file_manager.py:199, in CachingFileManager.acquire_context(self, needs_lock)
196 @contextlib.contextmanager
197 def acquire_context(self, needs_lock=True):
198     """Context manager for acquiring a file."""
--> 199     file, cached = self._acquire_with_cache_info(needs_lock)
200     try:
201         yield fileFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/xarray/backends/file_manager.py:217, in CachingFileManager._acquire_with_cache_info(self, needs_lock)
215     kwargs = kwargs.copy()
216     kwargs["mode"] = self._mode
--> 217 file = self._opener(*self._args, **kwargs)
218 if self._mode == "w":
219     # ensure file doesn't get overridden when opened again
220     self._mode = "a"File src/netCDF4/_netCDF4.pyx:2469, in netCDF4._netCDF4.Dataset.__init__()File src/netCDF4/_netCDF4.pyx:2028, in netCDF4._netCDF4._ensure_nc_success()FileNotFoundError: [Errno 2] No such file or directory: '/work/mh0033/m300602/icon/grids/none/ckdtree/rectgrids/none_res0.30_180W-180E_90S-90N.nc'

This is fixed by gname="r2b6_oce_r0004" but leads to problem 3... I think we can fix this by raising an error when gname is set to None.

Problem 3

ds["vort"].isel(time=10, depth=0).pyic.plot(
grid_type="native",
gname="r2b6_oce_r0004",
)

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[243], line 1
----> 1 ds["vort"].isel(time=10, depth=0).pyic.plot(
2     grid_type="native",
3     gname="r2b6_oce_r0004",
4 )File /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_accessor.py:26, in pyiconDataArray.plot(self, **kwargs)
23 @functools.wraps(pyic.plot, assigned=("__doc__", "__anotations__"))
24 def plot(self, **kwargs):
25     da = self._obj
---> 26     ax, hm = pyic.plot(da, **kwargs)
27     return ax, hmFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_plotting.py:1994, in plot(data, Plot, ax, cax, asp, fig_size_fac, mask_data, logplot, lon_reg, lat_reg, central_longitude, clim, cmap, conts, contfs, clevs, use_pcol_or_contf, contcolor, cincr, clabel, cbticks, xlabel, ylabel, xticks, yticks, template, cbar_str, cbar_pos, title_right, title_left, title_center, projection, coastlines_color, land_facecolor, axes_facecolor, noland, do_plot_settings, do_xyticks, do_gridlines, gname, fpath_tgrid, plot_method, grid_type, res, fpath_ckdtree, coordinates, lonlat_for_mask)
1989 #if 'grid_mapping' in list(data.attrs):
1990 #  plot_method = 'healpix'
1991
1992 # --- interpolate and cut to region
1993 if grid_type=='native' and plot_method=='nn':
-> 1994   datai = interp_to_rectgrid_xr(data.compute(), fpath_ckdtree, lon_reg=lon_reg, lat_reg=lat_reg, coordinates=coordinates)
1995   lon = datai.lon
1996   lat = datai.latFile /work/mh0256/m301014/.conda/envs/age-tracer/lib/python3.11/site-packages/pyicon/pyicon_tb.py:285, in interp_to_rectgrid_xr(arr, fpath_ckdtree, lon_reg, lat_reg, coordinates, radius_of_influence, compute, mask_out_of_range, mask_out_of_range_before)
283 # --- interpolate by nearest neighbor
284 inds = inds.flatten()
--> 285 arr_interp = xr.DataArray(arr.data[inds].reshape(lat.size, lon.size), dims=['lat', 'lon'], coords=dict(lat=lat, lon=lon))
287 # --- mask values where nearest neighbor is too far away
288 # (doing this after compute seems to be faster) FIXME check that!
289 if mask_out_of_range_before:IndexError: index 203212 is out of bounds for axis 0 with size 120947

This is fixed by coordinates="vlon vlat" and I suspect we can get pyicon to automatically guess how to fix this type of problem.

So the final working piece of code is

ds["vort"].isel(time=10, depth=0).pyic.plot(
grid_type="native",
gname="r2b6_oce_r0004",
coordinates="vlon vlat",
)