Fieldset.from_xarray_dataset with moving vertical grid

[apatlpo]

We are running simulations with a moving vertical grid. It works fine with netcdf files but fails when we attempt to create the fieldset from an xarray dataset (see error message at the bottom). Looking at the code, it is pretty clear this has not been implemented yet. Being able to feed xarray dataset would be very convenient for us and I was thus wondering whether there was a fundamental reason for it not being implemented or if it was simply that nobody had needed to do that thus far?

variables = {'U': 'u',
             'V': 'v',
             'W': 'w',
             'depth_psi': 'z_psi',
            }

dimensions = {'U':{'lon': 'lon_psi', 'lat': 'lat_psi', 'depth':'not_yet_set', 'time': 'time_counter'},
              'V':{'lon': 'lon_psi', 'lat': 'lat_psi', 'depth':'not_yet_set', 'time': 'time_counter'},
              'W':{'lon': 'lon_psi', 'lat': 'lat_psi', 'depth':'not_yet_set', 'time': 'time_counter'},
              'depth_psi':{'lon': 'lon_psi', 'lat': 'lat_psi', 'depth':'not_yet_set', 'time': 'time_counter'},
             }

fieldset = FieldSet.from_xarray_dataset(ds,variables,dimensions,
                                allow_time_extrapolation=True,
                                interp_method='cgrid_velocity',
                                       )

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/.miniconda3/envs/equinox/lib/python3.8/site-packages/xarray/core/dataarray.py in _getitem_coord(self, key)
    730         try:
--> 731             var = self._coords[key]
    732         except KeyError:

KeyError: 'not_yet_set'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
/dev/shm/pbs.8777261.datarmor0/ipykernel_26609/1988345921.py in <module>
----> 1 fieldset = FieldSet.from_xarray_dataset(ds,variables,dimensions,
      2                                 allow_time_extrapolation=True,
      3                                 interp_method='cgrid_velocity',
      4                                        )

~/.miniconda3/envs/equinox/lib/python3.8/site-packages/parcels/fieldset.py in from_xarray_dataset(cls, ds, variables, dimensions, mesh, allow_time_extrapolation, time_periodic, **kwargs)
    926             cls.checkvaliddimensionsdict(dims)
    927 
--> 928             fields[var] = Field.from_xarray(ds[name], var, dims, mesh=mesh, allow_time_extrapolation=allow_time_extrapolation,
    929                                             time_periodic=time_periodic, **kwargs)
    930         u = fields.pop('U', None)

~/.miniconda3/envs/equinox/lib/python3.8/site-packages/parcels/field.py in from_xarray(cls, da, name, dimensions, mesh, allow_time_extrapolation, time_periodic, **kwargs)
    479 
    480         time = da[dimensions['time']].values if 'time' in dimensions else np.array([0])
--> 481         depth = da[dimensions['depth']].values if 'depth' in dimensions else np.array([0])
    482         lon = da[dimensions['lon']].values
    483         lat = da[dimensions['lat']].values

~/.miniconda3/envs/equinox/lib/python3.8/site-packages/xarray/core/dataarray.py in __getitem__(self, key)
    740     def __getitem__(self, key: Any) -> "DataArray":
    741         if isinstance(key, str):
--> 742             return self._getitem_coord(key)
    743         else:
    744             # xarray-style array indexing

~/.miniconda3/envs/equinox/lib/python3.8/site-packages/xarray/core/dataarray.py in _getitem_coord(self, key)
    732         except KeyError:
    733             dim_sizes = dict(zip(self.dims, self.shape))
--> 734             _, key, var = _get_virtual_variable(
    735                 self._coords, key, self._level_coords, dim_sizes
    736             )

~/.miniconda3/envs/equinox/lib/python3.8/site-packages/xarray/core/dataset.py in _get_virtual_variable(variables, key, level_vars, dim_sizes)
    168         ref_var = dim_var.to_index_variable().get_level_variable(ref_name)
    169     else:
--> 170         ref_var = variables[ref_name]
    171 
    172     if var_name is None:

KeyError: 'not_yet_set'

[erikvansebille]

Thanks for reporting, @apatlpo. I don't think there is a fundamental reason for this not being implemented for xarray.

When we developed this feature of time-varying depth grids, we only had a use-case for netcdf files. And because our development philosophy is to only develop features for which we have use-cases (and hence test-cases), we never got around to implementing it for xarray; or even a proper NotImplementedError for that matter...

Perhaps you can give a go at implementing this for xarray? See the grid.z4d boolean that controls whether the depth is time-varying

[apatlpo]

Ok, I can give it a look. If you have any other pointers, don't hesitate to push them.

[apatlpo]

I give it a first shot but it leads to an out of bounds error at the moment. Discussion should probably move to #1131