Is there a robust way to set properties on a DataArray when opening a Dataset from file?

[frizwi]

DataArray's are self contained objects, for netcdf variables they include information about coordinate variable.

For example, consider the following netcdf file:

	record = UNLIMITED ; // (1 currently)
	j_centre = 120 ;
	i_centre = 175 ;
	j_grid = 121 ;
	i_grid = 176 ;
variables:
	float eta(record, j_centre, i_centre) ;
		eta:units = "metre" ;
		eta:long_name = "Surface Elevation" ;
		eta:coordinates = "t, x_centre, y_centre" ;
	double t(record) ;
		t:units = "days since 1990-01-01 00:00:00 +10" ;
		t:long_name = "Time" ;
		t:coordinate_type = "time" ;
	double x_centre(j_centre, i_centre) ;
		x_centre:long_name = "Longitude at cell centre" ;
		x_centre:coordinate_type = "longitude" ;
		x_centre:units = "degrees_east" ;
		x_centre:projection = "geographic" ;
	double x_grid(j_grid, i_grid) ;
		x_grid:long_name = "Longitude at grid corners" ;
		x_grid:coordinate_type = "longitude" ;
		x_grid:units = "degrees_east" ;
		x_grid:projection = "geographic" ;
	double y_centre(j_centre, i_centre) ;
		y_centre:long_name = "Latitude at cell centre" ;
		y_centre:coordinate_type = "latitude" ;
		y_centre:units = "degrees_north" ;
		y_centre:projection = "geographic" ;
	double y_grid(j_grid, i_grid) ;
		y_grid:long_name = "Latitude at grid corners" ;
		y_grid:coordinate_type = "latitude" ;
		y_grid:units = "degrees_north" ;
		y_grid:projection = "geographic" ;

The eta variable has "centre" dimensions and is on a 2D curvilinear quad grid. However, the reference to the coordinates of the grid corners, i.e. the variables x_grd and y_grid, are lost when the DataArray for eta is constructed.

We're currently extending xarray via accessors, to use the grid corners to properly plot variables like eta, but are forced to make our custom plot method on the Dataset class rather than the DataArray.

So, this works as the Dataset has all the information we need:

import xarray as xr
ds = xr.open_dataset("<filename>")
eta = ds['eta'].isel(record=0)
ds.cem.plot(eta) # where our dataset accessor is named cem

But we would prefer to make this work:

eta.cem.plot()

As the syntax is more natural, and consistent with other packages like GeoViews, and even xarray itself.

Is there a way to do this?

We've looked at using the attrs property to store the Dataset object on the DataArray variables via ds.cem.open_dataset(..) method but this is problematic in a range of cases, not least of which breaks to_netcdf()

[benbovy]

I don't know if there's a good way to do this currently. However this makes me wondering whether, in the context of the (ongoing) explicit indexes refactor, we might want to return all the relevant indexed coordinates when constructing a DataArray from a Dataset variable, possibly including coordinates that don't share common dimensions with the variable itself.

In your example, the indexes refactor will eventually allow setting a custom index based on the 4 coordinates x_centre, y_centre, x_grid and y_grid for enabling grid-aware (and geographic-aware) selection, alignment, etc. As we probably want to perform the same kind of operations on the eta variable accessed as a DataArray, it would make sense if the latter could still have this index (and as a side effect all of its 4 corresponding coordinates).

That said, it would mean that the set of all dimensions of the coordinates of a DataArray is not guaranteed to be always a subset of the dimensions of the DataArray itself, and I'm not sure about all the implications that such a big change in the DataArray data model would have.

[dcherian]

custom index based on the 4 coordinates

I think proposed XgcmIndex would propagate all grid info when selecting out DataArrays.

[frizwi]

Having access to all coordinates would certainly help as there is no way to currently get at this information from the DataArray. But for performance reasons, ideally, we'd be looking at attaching an object (by reference) to the DataArray. We create Shapely geometry for each "grid kind" (centre's, edges etc) that are shared across many variables so the workflow would be to open the dataset, create objects and attach to each relevant variable. It need not persist so for any downstream method that are looking to serialize it (eg. saving to disk), then it may safe be ignored.

[benbovy]

Then accessing all grid info through the index (without necessarily exposing all this info as coordinates -- if my understanding of XgcmIndex is correct @dcherian ?) may be the right way to go. Indexes are the kind of objects attached to Datasets (and propagated to selected out DataArrays along with the coordinates) that are essentially living in-memory.

We create Shapely geometry for each "grid kind" (centre's, edges etc)

As a side note, I have some rough plans experimenting later this year with Xarray-compatible geospatial or geographic indexes backed by pygeos (shapely) or s2geometry Python bindings (geopandas/community#10 (comment)), and I'd love to know more about your use case!

[frizwi]

@benbovy Our use case is that while the Xarray Dataset defines the "mesh", either structured or unstructured, the Xarray Dataarray lose that link. So for example, if you want to plot (2D) the Surface Elevation, then the "eta" variable only has the "cell" values and the indices for them, but not the polygons that define the mesh. For this we need both the Dataset (for the mesh) and the associated Dataarry