UXarray Redesign Thoughts and Options

[erogluorhan]

⚠️ If you haven't done so, you can start with reading Dataset class for handling Grid and Data files as background information.

We are in the middle of a UXarray API redesign work based on user feedback and developers’ discussions recently. We’d like to come up with an API design that satisfies both the usage requirements & specifications as well as complies with the Xarray-extension best practices.

I’ll describe 4 potential design options (including the current and the former UXarray API), provide concrete examples to give an idea of how UXarray usage would look like with each of them, and discuss pros & cons with each.

For simplicity and voting purposes, each option will be provided in a separate cell below. Please provide your comments under which option they correspond to, or propose a new design in a new cell.

I can also try providing some UML-like sketches if needed.

For the sample codes, I’ll consider a case where:

grid_path = Path or URL to the unstructured grid topology definition file
⚠️ (This is one of the essential differences of unstructured grid analysis workflow than the structured grids; unstructured data sets/variables need this separate file that defines the grid topology)

data_path = Path or URL to the dataset file that contains the actual data, i.e. variable(s) of interest, e.g. vortex, which conforms to the unstructured grid topology given in grid_path

[erogluorhan]

1. Current API (e.g. xgcm)

Description: The API backbone is the uxarray.Grid class. It takes care of unstructured grid recognition, grid attributes, metadata, and grid-specific functionality. Dataset read-in is left to Xarray itself in this API.

Example:

# Open grid file and dataset file with Xarray
# So, `ds_grid` and `ds` are of type `xarray.Dataset`
import xarray as xr
ds_grid = xr.open_dataset(grid_path)
ds = xr.open_dataset(data_path)

# Construct an `uxarray.Grid` object out of `ds_grid`
import uxarray as ux
uxgrid = ux.Grid(ds_grid)

# A grid-wise analysis function call looks like
uxgrid.calculate_total_face_area()

# A variable-specific analysis function call looks like (i.e. would consume a `xarray.DataArray` variable from `xr_data_ds`):
uxgrid.integrate(ds.var1)

# An Xarray function call that does not require grid information, e.g. `min`, on dataset and a data variable looks like:
ds.min()
ds.var1.min()

Pros:

• UXarray does not deal with data loading at all, i.e. focuses only on handling unstructured grid-aware functions on well-known xarray.Dataset and xarray.DataArray

• Cear distinction between Xarray and UXarray namespaces/signatures and hence workflows

Cons:

• Unstructured grid analysis may require user to have full awareness of they the difference between structured and unstructured grid data and a pretty different workflow than Xarray workflows

[rljacob]

I redid this example to make the variables as terse as in examples 3 and 4.

[paullric]

I disagree. Why should users need to care about the structure of the underlying data? Structured grids (2d coords) can be easily recast as unstructured grids (1d coords), so there's no reason why a general set of operators couldn't work on structured grids similarly to how they work on unstructured grids.

[whannah1]

@paullric if I'm understanding you correctly, you are suggesting that the uxarray.Grid class should be able to handle structured data automatically? I'm imagining that a user could have the same block of code work seamlessly for both unstructured and structured cases (as long as the user remembers to recreate the uxgrid object for each dataset).

[rljacob]

There is certainly A set of operators that work on both structured and unstructured but I contend that is subset of all xarray operators/operations. Possibly a small subset.

[erogluorhan]

All, please see the update on the team decision below.

[erogluorhan]

2. Former API

Description: The API backbone was still the uxarray.Grid class. However, UXarray not only took care of unstructured grid recognition, grid attributes, metadata, and grid-specific functionality, but also handled the dataset read-in via its own Xarray-like dataset module that returned uxarray.Grid object when files were opened.

Example:

# Open grid file and dataset file together with UXarray
# Note in this case that `ux_grid` below is of type `uxarray.Grid` and it had both grid topology and data together
# Also note that `ux.open_dataset()` below could take more than one data files
import uxarray as ux
ux_grid = ux.open_dataset(grid_path, data_path)

# A grid-wise analysis function call looks like
ux_grid.calculate_total_face_area()

# A variable-specific analysis function call looks like (i.e. would consume a `xarray.DataArray` variable from `xr_data_ds`):
ux_grid.integrate(ux_grid.var1)

# An Xarray function call that does not require grid information, e.g. `min`, on dataset and a data variable looks like:
ux_grid.ds.min()
ux_grid.ds.var1.min()

Pros:

• Well-defined uxarray.Grid object with its own Xarray-like data loading features as well as data and grid topology both together.

• Clear distinction between Xarray and UXarray namespaces/signatures and hence workflows

Cons:

• Unstructured grid analysis may require user to have full awareness of the difference between structured and unstructured grid data and a pretty different workflow than Xarray workflows

• The fact that ux.open_dataset returning ux.Grid might be confusing to the Xarray-experienced users

[xylar]

It would be important for this API to also have an API similar to the first option, that is an option to use xarray.Dataset objects to construct an uxarray.Grid, because it wouldn't always be the case that they are written to files. For example, it might be necessary or desirable to manipulate an xarray dataset before it is ready to be in input to uxarray.Grid. The xarray datasets could, of course, be written to temporary files but that seems needless.

ux_grid = ux.Grid(ds_grid, ds)

[erogluorhan]

All, please see the update on the team decision below.

[erogluorhan]

3. Registering Grid object as an accessor for xarray.Dataset (e.g. xcdat)

Description: uxarray.Grid class can be renamed to ux.GridAccessor and can be decorated as follows:

@xr.register_dataset_accessor("uxgrid")
class GridAccessor:
...

Thus, with any xarray.Dataset instantiation, an ux.Grid object can be added to xr.Dataset as a custom property (i.e. actually a custom namespace under). Via this namespace, UXarray can take care of unstructured grid recognition, grid attributes, metadata, and grid-specific functionality through the actual xr.Dataset. We can also implement ux.open_dataset that would tweak opening and returning a single xr.Dataset object from grid topology and data files with uxgrid property.

We could even define xr.DataArray accessors to execute unstructuresd grid-aware functions on Xarray variables in a Xarray-like signature (shown at the very end of the example below)

Example:

# Open grid file and dataset file together with UXarray
# In this case, we will implement `ux.open_dataset()` and `ux.mfopen_dataset()` to return `xr.Dataset` (i.e. `ds` below) that has `uxgrid` in it.
import uxarray as ux
ds = ux.open_dataset(data_path, grid_path, **kwargs)

# A grid-wise analysis function call looks like
ds.uxgrid.calculate_total_face_area()

# A variable-specific analysis function call looks like (i.e. would consume a `xarray.DataArray` variable from `xr_data_ds`):
ds.uxgrid.integrate(ds.var1)

# `xr.DataArray` accessor would give us a function execution like:
ds.var1.uxgrid.integrate()

# An Xarray function call that does not require grid information, e.g. `min`, on dataset and a data variable looks like:
ds.min()
ds.var1.min()

Pros:

• Quote from Extending Xarray: “Using an accessor provides an implicit namespace for your custom functionality (UXarray in this case) that clearly identifies it as separate from built-in xarray methods.”

• An unstructured grids workflow closer to Xarray-based workflows compared to options (1) and (2)

Cons:

• By design, any operation that creates a new xr.DataSet will also re-construct a ux.GridAccessor, which would lead to loss of the state of the existing uxgrid` (Re Stateful user-defined accessors)

• I don’t know for now if the above issue can be worked around. Assuming yes for now, we’d need to avoid computationally-complex function calls in GridAccessor.init()` because of this reconstruction.

[erogluorhan]

I have been prototyping approaches discussed here for a while, and I have some thoughts and a draft PR (#216 ) up. I'd like to share some details about those and ask further feedback from you all.

First of all, xarray's native dataset accessor seems unapplicable to UXarray's case because ustructured workflows consist of two file types (one grid topology file, and the other, maybe multiple files, that house data sets, i.e. actual variables of interest). There is no problem when the dataset accessor is triggered by xarray.Dataset instantiation through grid topology files, and our Grid object gets created without any issues. However, at the times when xarray.Dataset gets instantiated through dataset files, a Grid object cannot be created due to the lack of grid topology variables in such dataset files. Since we are trying to open dataset files and register grid object to them as accessor, this goes nowhere.

Based on my readings on xarray's custom indexes, I am not sure if we should make UXarray rely upon that because at this point, efforts that use it does not seem like more than experiments and trials. Changing the design and relying upon a WIP sounds risky to me.

Because of these, even though I was hesitant about subclassing xarray, I based my prototype still on an inheritance scheme; please see PR #216 if you have time. I feel like the current architecture and what we are planning to do with it and how we will expand it would not give use too many unexpected drawbacks (and actually each approach has its own drawbacks).

I's highly appreciate any perspectives/feedback on this prototype and views on potential flaws/risks with it.

[benbovy]

A couple of comments (although I haven't looked at #216 in detail yet and I'm not very familiar with ugrids):

First of all, xarray's native dataset accessor seems unapplicable to UXarray's case because ustructured workflows consist of two file types (one grid topology file, and the other, maybe multiple files, that house data sets, i.e. actual variables of interest).

I guess it could be handled nicely with an Xarray custom IO backend? It would allow something like this:

import xarray as xr

ds = xr.open_dataset(
    {"grid": grid_filename_or_obj, "data": filename_or_obj},
    engine="uxarray",
    gridspec=...,
    islatlon=...,
) 

This would be equivalent to the open_dataset function implemented in #216, but with the benefit of using Xarray's API directly.

Based on my readings on xarray's custom indexes, I am not sure if we should make UXarray rely upon that because at this point, efforts that use it does not seem like more than experiments and trials. Changing the design and relying upon a WIP sounds risky to me.

I'm obviously biased here, but even though Xarray custom index is an experimental feature, Xarray's internal PandasIndex and PandasMultiIndex already use the new data model since the 2-3 last Xarray releases, with no major issue reported so far.

In its most basic version, you could have a custom "meta" index that simply groups the ugrid coordinates together and that forwards API calls to the underlying PandasIndex objects of those coordinates (see, e.g., pydata/xarray#7182). Then you can attach a uxgrid object to this meta-index and do whatever you want with it. Actually, it wouldn't be much different than attaching it to UxDataset like in #216, with the great advantage that subclassing xarray.Dataset is not needed.

I don't think this index approach is more risky than subclassing xarray.Dataset. The main risk is that xarray.indexes.Index API will change in the future, but it is mitigated if the custom Index just forwards the arguments to its encapsulated Xarray built-in index(es). Also the Index API is public (we will carefully consider any change that would break it) and closely follows the Dataset / DataArray API which makes it already pretty stable IMO.

Like @dcherian says, one of the big challenges I guess is propagating the ugrid information or object across various operations, but I don't see how this would be more or less challenging using either the custom index or Dataset subclassing approach.

Since we are trying to open dataset files and register grid object to them as accessor, this goes nowhere.

Having the grid object attached to a custom index rather than an accessor, this wouldn't be a problem. You could then provide an accessor that accesses / consumes the grid object via the index. This is the approach used in https://github.com/martinfleis/xvec.

[erogluorhan]

Hi @benbovy , thank you very much for these great inputs. There are many things I can learn and probably employ here; e.g. I wasn't familiar with the custom IO backend and its applicability here. I will go over each of them as soon as I can.

[dcherian]

Here's a prototype XgcmIndex: https://github.com/dcherian/xindexes/blob/main/xgcmindex.ipynb

It somewhat works (I only implemented very simple things, selecting out a single cell center along X). The cool thing is that we can use the .sel machinery to return a consistent grid topology. However

1. AFAICT there is no mechanism to do the same thing with .isel (i.e. if the user provides desired positions for cell centers only; the index cannot rewrite that to include consistent positions for the faces; cc @benbovy)
2. The other missing piece is propagating the Index when extracting a DataArray (Propagating multi-variable indexes in dataarrays pydata/xarray#7538)

Now that the Grid object is in the Dataset structure (through the Index) an accessor can use it to provide functionality. This doesn't work now because of (2) [xgcm can only operate on dataarrays]

[erogluorhan]

All, please see the update on the team decision below.

[erogluorhan]

4. Inheriting (i.e. subclassing) xarray.Dataset and xr.DataArray

I will skip describing this and giving an example as I suppose we’d know what inheritance would look like (I can still provide if needed though). I will share a few Pros/Cons and say we should avoid due to the cons/costs with it.

Pros:

• This would give us the most OOP way for our own ux.DataSet and ux.DataArray so that we could have very Xarray-like function calls, thus workflows. Something like:

  import uxarray as ux
  ds = ux.open_dataset(...)   # ds of type ux.Dataset
  ds.calculate_total_face_area()
  ds.var1.integrate()    # ds.var1 of type ux.DataArray
  
  
  # An Xarray function call that does not require grid information, e.g. `min`, on dataset and a data variable looks like:   
  ds.min()
  ds.var1.min()
  

Cons:

• Xarray upgrades would turn into a severe headache on our code (Re: Extending Xarray)

• We’d need to deal with native xarray objects returned by builtin xarray methods. Taking into account the size of Xarray API, it’d be too much even if it was only converting them to uxarray objects (Re: Subclassing Xarray Dataset/DataArray and Extending Xarray).

[clyne]

From a user perspective, inheritance may provide the simplest API for people who are already Xarray users or who want to perform analysis that combines unstructured and structured data. I'd suggest considering a few more use cases. For example, what would a script look like that read in a collection of structured and unstructured data sets, and plotted the same variable from each? Ideally one would want to loop over the the data sets and minimize special handling (if grid is structured do.. else ..). What if the coordinates of the maximum value of a variable was wanted for each data set? And so on.

Maybe the handling of builtin methods returning Xarray objects can be automated in some way?

[whannah1]

This approach is how I always imagined uxarray would work, but now that I'm reading about the potential complications it doesn't seem like a good idea. I haven't seen anything I don't like in the other examples, but I need to think about them more.

[TomNicholas]

I wouldn't recommend subclassing xarray objects. They are not (currently) designed to be subclassed, and all sorts of weird things might happen. Fixing those unexpected behaviours upstream in xarray would also be slow and technical.

Maybe one day we will make this easier to do though, and if you want to work on it then by all means come and chat to us!

[rljacob]

A prototype implementation of option 4 is in PR #216. Please comment.

[erogluorhan]

A prototype implementation of option 4 is in PR #216. Please comment.

I think it'd be fair to say it's a combination of options (3) and (4). To note, option (3) is implemented not through xarray's dataset accessor regsitering (due to the drawback I explained above in the option 3 cell) but via manually setting Grid object as a property in UxDataset.

[clyne]

@erogluorhan could you include in your examples how an Xarray function that does not require grid information would be invoked? For example, returning the minimum value of the field. Thanks

[erogluorhan]

Great point, thanks! Added now to the end of examples in each.

[erogluorhan]

All, please see the update on the team decision below.

[rajeeja]

I tend to gravitate towards approach 4 as it gives the most OOP which is super critical for any forward looking design. Having designed/developed the current or former API I like both of those approaches :) Former approach is much simpler and efficient.

3rd approach looked good at first, but, later after reading @dcherian comment "will limit what other ecosystem packages you can use” I’m not very sure about switching to approach 3.

[erogluorhan]

I believe the comment "will limit what other ecosystem packages you can use” was in particular about the xugrid's approach of coming up with its own dataarray dataset structures, not directly about option (3) - accessors. @dcherian can clarify if otherwise.

Regarding the option (4), I am quite hesitant about going with that, given my own research and the red flags raised in the discussion at all.

@rljacob , any thoughts?

[rajeeja]

I reread the discussion again and it seems that you're right about the reference specifically to xugrid's approach of coming up with their own dataarray.

This has been a very long (time and length wise:) discussion. IMO, we should proceed with option 3.

[rljacob]

Yes I am fine with option 3.

[philipc2]

While I really like option 3 and believe that it should be implemented in one way or another, I think it's important to consider how we could build upon it, especially for supporting visualization functions. I'm concerned about how restricted we would be to Xarray data structures xr.Dataset and xr.DataArray. For polygon visualization with Datashader, the geometry of our grid is represented as a collection of Shapely Polygon objects, which are then stored in a spatialpandas.GeoDataFrame With a stateless accessor, we would need to re-construct this data structure often.

This could be resolved by using custom indexes (see vector cube). I'm a huge fan of this implementation, but it may be best to stick to one of the others for now until we can dedicate more time to it.

The following examples introduce the uxplot accessor or class, depending on if we can cache other data structures (GeoDataFrames) through something like a wrapper or ux.Dataset / ux.DataArray

Example 1 (uxplot Accessor)

If we can cache our geometry data structures, we can use a uxplot accessor to wrap Holoviews/Datashader plotting routines

import uxarray as ux
from cartopy import crs as ccrs
ds = ux.open_dataset(data_path, grid_path, **kwargs)

# uxgrid accessor would still exist 
ds.uxgrid.calculate_total_face_area()

# First visualization call computes and caches the geometry 
ds['var1'].uxplot.polygons(projection=ccrs.Robinson()) 

# Second visualization call uses cached geometry
ds['var2'].uxplot.polygons(projection=ccrs.Robinson())

# Visualizing with a different projection requires re-computing the geometry
ds['var1'].uxplot.polygons(projection=ccrs.PlateCarree()) 

Example 2 (uxplot Class)

Assuming a stateless accessor with no ux.Dataset or ux.DataArray, the best alternative would be having a function for converting to a GeoDataFrame with visualizations being done using that data structure, external from any xr.Dataset or xr.DataArray

import uxarray as ux
import hvplot.pandas
from cartopy import crs as ccrs

ds = ux.open_dataset(data_path, grid_path, **kwargs)

# GeoDataFrame needs to be stored outside of our Dataset
gdf = ds.uxgrid.to_gdf(projection=ccrs.Robinson()) 

# Plot with hvPlot 
gdf.hvplot.polygons(c='var1')

# Create datashader/holoviews wrappers for Uxarray
ux.uxplot.polygons(gdf, c='var2')

# Plotting a new projection would require new computations 
gdf = ds.uxgrid.to_gdf(projection=ccrs.PlateCarree()) 
ux.uxplot.polygons(gdf, c='var2')

Disclaimer

The topic of visualization merits it own discussion. Once we wrap up our re-design talks, I'll open up a separate discussion to dive deeper into the design and implementation

[rajeeja]

This could be resolved by using custom indexes (see vector cube). I'm a huge fan of this implementation, but it may be best to stick to one of the others for now until we can dedicate more time to it.

I tried to test this, had to upgrade shapely 2.0 etc. It is indeed important to consider the developments in this direction and evaluate our ability to be in sync with this capability, it should certainly be a priority during our design discussion phase. I like the uxplot accessor example, it'd require some maintainence on our part, but make life easier for users.

All that said, we should finalize things soon to make progress with planned new feature additions to the uxarray package.

[rljacob]

To all following this discussion, we are planning to make a decision on our Feb 28th telecon.

[erogluorhan]

Dear all, we have made a decision to subclass Xarray (i.e. option #4) as already proposed in the Redesign prototype PR #216. The whole conversation here and our further discussions with Xarray developers have been greatly helpful for this decision; thanks a lot!

To get this new design merged to the current code structure, we have

1. Formed a UXarray redesign working group (@anissa111, @erogluorhan, @philipc2, @rajeeja, @rljacob)
2. Put together a public redesign document to keep track of this work. Please feel free to contribute to this in any way.

We are hoping to get this new design released once we address all the items in the above linked document.