Xarray's explicit indexes

[Huite]

For examples, see:

https://github.com/martinfleis/xvec
https://github.com/dcherian/crsindex/blob/main/crsindex.ipynb

They can be set with:
https://docs.xarray.dev/en/stable/generated/xarray.DataArray.set_xindex.html

One issue is that in the UGRID conventions, the connectivity arrays are 2D. So the first dimension matches whatever's going on in the data (node, edge, face), but the second doesn't. This can be avoided with a trick using numpy structured arrays:

import numpy as np
import xarray as xr

ds = xr.open_dataset("Grevelingen_FM_grid_20190603_0001_net.nc")
faces2d = ds["NetElemNode"].values.astype(int) - 1  # face node_connectivity
triangle = np.dtype([("0", int), ("1", int), ("2", int)])
faces = np.empty(len(faces2d), dtype=triangle)
faces.view(int)[:] = faces2d.ravel()
ds = ds.assign_coords(faces=("nNetElem", faces))

This feels like a hack, and doesn't seem like a serious option.
An alternative approach is perhaps using a Pandas ExtensionArray. This seems like a lot of work: https://stackoverflow.com/questions/68893521/simple-example-of-pandas-extensionarray

Next, we can define a UgridIndex:

class Ugrid2dIndex(Index):
    def __init__(
        self,
        node_x,
        node_y,
        faces,
    ):
        self._ugrid = xu.Ugrid2d(node_x.values, node_y.values, -1, faces.values.view(int).reshape((-1, 3)))

    @classmethod
    def from_variables(cls, variables, options):
        return cls(
            variables["NetNode_x"],
            variables["NetNode_y"],
            variables["faces"],
        )

This seems to do the trick:

newds = ds.set_xindex(
    (
        "NetNode_x",
        "NetNode_y",
        "faces",
    ),
    Ugrid2dIndex,
)

Now we could start adding methods to the Ugrid2dIndex and make it possible to select values with sel. Or alternatively provide other operations via a .ugrid accessor.

However, it doesn't solve all problems that we encounter with UGRID topologies (?):

• Node, edge, and face dimension are fully separated according to xarray's data model. Selecting one but not the other will generally invalidate the data <-> topology link. (So solve via a .ugrid.sel and .ugrid.isel instead.)
• A full UGRID topology requires possibly connectivities with faces, edges, nodes, and possible interrelations. As far I can tell, the index is not carried around for every variable. (Solve by using multiple indexes internally containing full grid? One for faces, one for edges (with edge node connectivity). But seems impractical for data on nodes?)

I think the fundamental question is whether a UgridDataset is properly considered as an xarray Dataset: I don't think so, as it breaks a premise of fully separated, orthogonal dimensions. From that perspective, it may be better to accept this, and special case the UGRID topology dimensions.

[Huite]

More notes:

These explicit indexes are currently insufficiently mature (e.g. tries to convert to PandasIndex now when accessing via ds.indexes), but once mature enough:

• The UGRID topology could be fully contained within an explicit index. Building the index still requires a numpy structured hack or pandas ExtensionArray to accept the connectivity arrays as coordinates, or requires changes to set_xindex.
• The indexes cannot be relied upon to persist through many operations, as they may generate invalid topologies. This is currently also the case if (node, face, edge) dimensions are given indexes, but there's no upfront warnings or errors with standard Datasets.

Without changes to the explicit indexes:

• No UgridDataArray exists, since there's multiple variables required retain the topology information, which isn't split off.
• Specific operations occur through the .ugrid accessor, generally mentioning the variable name. uds["water_depth"].ugrid.plot() would become uds.ugrid.plot("water_depth").
• That might get a little cumbersome for operations on multiple Ugrid objects? E.g.: uds.ugrid.func(uds2, "var1", "var2")
• All operations basically require mutation of the dataset.
• I.e. ergonomics seem worse in general.

The largest benefit is that it's a good deal simpler. No need to juggle between xarray <-> xugrid types. No need to forward calls, which still feels a bit haphazard to me as it currently stands. Also no custom type which trip up checks in other packages that do understand xarray objects.

[Huite]

Alternatively, maybe a topology and its data should be grouped together in a completely different way, see: https://xarray-datatree.readthedocs.io/en/latest/

This also has a similar drawback, which is that the API looks quite different when working with unstructured viz a viz structured data, which I'd really like to avoid.

[dcherian]

Node, edge, and face dimension are fully separated according to xarray's data model.

I think you can handle this in Index.sel which will be called by Dataset.sel

One issue is that in the UGRID conventions, the connectivity arrays are 2D. So the first dimension matches whatever's going on in the data (node, edge, face), but the second doesn't.

The same problem exists with handling CF style bounds arrays and IntervalIndex, so I think we have enough motivation to construct a proper solution in Xarray. I'm not entirely sure what that solution looks like (both structured and extension arrays have come up in the discussion).

[Huite]

I think you can handle this in Index.sel which will be called by Dataset.sel

Yes, I think so too now. That would be a very comfortable solution.

The same problem exists with handling CF style bounds arrays and IntervalIndex, so I think we have enough motivation to construct a proper solution in Xarray.

Ah, right -- that's good news!

Ideally, the same index could be associated with all three dimensions somehow (... I think). I think that's the only exceptional part of the UGRID topology versus the other uses of the explicit indexes.

I'm itching to find out what this entails, hopefully I can free up some time to investigate.

[dcherian]

, the same index could be associated with all three dimensions somehow

This is already possible with Datasets today. In https://github.com/dcherian/xindexes/blob/main/crsindex.ipynb, CRSIndex is associated with the x dimension, y dimension, and spatial_ref scalar variable.

[Huite]

With the merging of pydata/xarray#10137

# %%

import xugrid as xu
import xarray as xr

from xarray import Index

# %%

class Ugrid2dIndex(Index):
    def __init__(
        self,
        node_x,
        node_y,
        faces,
    ):
        face_nodes = faces.fillna(-1).to_numpy().astype(int)
        self._ugrid = xu.Ugrid2d(node_x.values, node_y.values, -1, face_nodes)

    @classmethod
    def from_variables(cls, variables, options):
        return cls(
            variables["mesh2d_node_x"],
            variables["mesh2d_node_y"],
            variables["mesh2d_face_nodes"],
        )

    def should_add_coord_to_array(self, name, var, dims):
        return True


# %%

ds = xu.data.elevation_nl(xarray=True)
ds2 = ds.set_coords("mesh2d_face_nodes")

newds = ds2.set_xindex(("mesh2d_node_x", "mesh2d_node_y", "mesh2d_face_nodes"), Ugrid2dIndex)
print(newds["elevation"])

<xarray.DataArray 'elevation' (mesh2d_nFaces: 5248)> Size: 21kB
[5248 values with dtype=float32]
Coordinates:
    mesh2d_face_x      (mesh2d_nFaces) float64 42kB ...
    mesh2d_face_y      (mesh2d_nFaces) float64 42kB ...
  * mesh2d_node_x      (mesh2d_nNodes) float64 22kB 9.926e+03 ... 3.217e+04
  * mesh2d_node_y      (mesh2d_nNodes) float64 22kB 3.847e+05 ... 3.974e+05
  * mesh2d_face_nodes  (mesh2d_nFaces, nmax_face) float64 126kB 138.0 ... 1.9...
Dimensions without coordinates: mesh2d_nFaces
Indexes:
  ┌ mesh2d_node_x      Ugrid2dIndex
  │ mesh2d_node_y
  └ mesh2d_face_nodes
Attributes:
    unit:     m NAP

This now succesfully propagates the indexes and the mesh2d_face_nodes coordinates to the DataArray, which is great news!

I noticed .indexes does still try to cast to PandasIndex, but I think this is enough to start prototyping!

[wpbonelli]

@Huite I think .indexes will stay that way since it is deprecated in favor of .xindexes, if I am reading this comment right? pydata/xarray#10257 (comment)

Excited to see where this goes.

[dcherian]

^ Yes exactly, use .xindexes instead.

One issue is that in the UGRID conventions, the connectivity arrays are 2D. So the first dimension matches whatever's going on in the data (node, edge, face), but the second doesn't.

We can now defer to the index to decide which coord vars to include with a DataArray : pydata/xarray#10137. This is released on v2025.06.0

[Huite]

My idea of combining with an accessor seems to work as expected. This ad hoc plotting method runs:

from xugrid.plot.plot import pcolormesh

@xr.register_dataarray_accessor("ugrid")
class UgridAccessor:
    def __init__(self, xarray_obj):
        self._obj = xarray_obj

    def plot(self, ax=None, **kwargs):
        da = self._obj
        indexes = list(da.xindexes.values())
        grids = {index._ugrid for index in indexes if isinstance(index, Ugrid2dIndex)}
        assert len(grids) == 1
        grid = grids.pop()
        return pcolormesh(grid, da, ax, **kwargs)

ds = xu.data.elevation_nl(xarray=True)
ds2 = ds.set_coords("mesh2d_face_nodes")
newds = ds2.set_xindex(("mesh2d_node_x", "mesh2d_node_y", "mesh2d_face_nodes", "mesh2d_face_x", "mesh2d_face_y"), Ugrid2dIndex)
uda = newds["elevation"]
uda.ugrid.plot(vmin=-20, vmax=50)

Currently spatial selections occur via the ugrid accessor as well, we could probably directly support via the Index as well:

...
    def sel(self, labels):
        grid = self._ugrid
        x = labels["mesh2d_face_x"]  # hardcoded for demonstration
        y = labels["mesh2d_face_y"]
        xy = np.column_stack((x, y))
        face_index = self._ugrid.locate_points(xy)
        return IndexSelResult(
            {grid.face_dimension: face_index},
            drop_coords=["mesh2d_node_x", "mesh2d_node_y", "mesh2d_face_nodes"],
        )
...


uda.sel(mesh2d_face_x=[100_000.0, 150_000.0, 200_000.0], mesh2d_face_y=[400_000.0, 450_000.0, 500_000.0])

The result still seems to contain the face_x and face_y coords despite explicitly dropping them.

Anyway, this keeps looking very good!

Changes

By my reckoning now, we can basically start deprecating the UgridDataArray and UgridDataset classes already with current xarray capabilities.

• The topology objects (Ugrid1d, Ugrid2d) may remain unchanged. They could possibly be pared down a bit, in particular with regards to the selection methods, but I see no immediate reason to prioritize this.
• The UgridDataArrayAccessor and UgridDatasetAccessor can be converted into ordinary xarray accessors.
• The UgridRolesAccessor (which is already an ordinary accessor) needs no changes. I see no harm in having more than one accessor.
• For backwards compatibility, we could keep the UgridDataArray and UgridDataset constructors around, but they would emit a warning and return ordinary xarray objects.
• We do need something else that does something similar: i.e. recognize the UGRID variables, create the relevant coordinates, and initialize the topologies and indexes. Currently thinking of something along the lines of:

# IO methods can be easily maintained for convenience, are not confusing:
xugrid.open_dataset(path)
xugrid.open_dataarray(path)

# From existing objects:
uds = ds.ugrid.initialize()

# With an optional argument for providing a topology (e.g. as the result from mesh generation) and an xarray object that doesn't yet hold UGRID vars:
uds = ds.ugrid.initialize(grids=[...])
uda = da.ugrid.initialize(grid)

There are three alternative constructors: from_data (from a numpy array and a grid), from_geodataframe (geopandas) and from_structured2d.

The last one could be added to the .ugrid accessor, although that feels a bit confusing. Easiest to have them as free functions:

uda = xugrid.from_data(values, grid)
uds = xugrid.from_geodataframe(gdf)
uda = xugrid.from_structured2d(da)

Which seems manageable.

Finally, we won't be able to rely on isinstance(UgridDataArray) anymore. I think we can use da.ugrid.is_ugrid instead as mentioned in #75. This then basically just checks whether theres's a UgridIndex involved somewhere in the .xindexes.

We can also rid ourselves of wrapped functions like full_like, concat, merge, etc., since they will just work out of the box (again should be given warnings).

We can also warn for typechecks with a metaclass:

import warnings

class DeprecatedMeta(type):
    def __call__(cls, *args, **kwargs):
        warnings.warn(
            f"{cls.__name__} is deprecated, use NewClass instead",
            DeprecationWarning,
            stacklevel=2
        )
        return super().__call__(*args, **kwargs)
    
    def __instancecheck__(cls, instance):
        warnings.warn(
            f"isinstance() check against {cls.__name__} is deprecated",
            DeprecationWarning,
            stacklevel=2
        )
        return super().__instancecheck__(instance)
    
    def __subclasscheck__(cls, subclass):
        warnings.warn(
            f"issubclass() check against {cls.__name__} is deprecated",
            DeprecationWarning,
            stacklevel=2
        )
        return super().__subclasscheck__(subclass)

class OldClass(metaclass=DeprecatedMeta):
    def __init__(self):
        self.value = "old"

Additional changes

The changes above will make xugrid quite a bit better, since we can throw away the wrapping code and close all of several issues that are caused by incomplete wrapping.

In the first iteration, I think we wouldn't add much to the UgridIndex. Everything stays under the .ugrid accessor for now, we just use the Index to hold the topology object.

The xarray Index base class mentions the following methods:

• init/from_variables/should_add_coord_to_array: obviously always needed
• concat: Not sure this makes sense for unstructured topologies. In principle, we could join topologies if they have no overlap. Regardless, we'd need to add it as Ugrid2d method first.
• copy: makes sense, should be implemented.
• create_variables: I don't think the topology can be re-used as coordinates.
• equals: should be implemented
• isel: standard indexing behavior suffices (to select a face or node, etc.). But I don't think it makes sense to allow selecting a single mesh2d_node_x or something.
• join: same as concat
• reindex_like: arguably yes, since Ugrid2d has such a method: https://deltares.github.io/xugrid/api/xugrid.Ugrid2d.reindex_like.html#xugrid.Ugrid2d.reindex_like; this should probably be called reorder_like instead.... Reindexing might be possible, maybe as a method wrapper around the regridding classes. But in that case, I feel like we might as well refer to the regridders directly.
• rename: I think yes, since the topologies have names too.
• roll: doesn't make sense
• sel: probably the most interesting one. This would work on the actual coordinate names (rather than the current .ugrid.sel which just takes x and y arguments). This is also an opportunity to evaluate the selection semantics and make sure they align with xarray as much as possible. But no immediate changes required.
• stack: doesn't make sense
• to_pandas_index: won't work
• unstack: doesn't make sense

So in terms of priorities:

High: copy, equals, rename
Low: sel, concat, join, reindex_like (probably in that order)

[deltamarnix]

Since we are still on a 0.x version, I wouldn't mind to not deprecate, but directly change the API and make this the first 1.0 release. The ideas seem good and solve a bunch of issues.
This is the way we want to go forward forward and removing the old classes gives good errors to the users. Adding a bunch of warnings is often not so obvious, because people think warnings are not to be looked at.

[Huite]

Fair enough; an error instead of a warning may be better.

In my mind, the benefit of catching it explicitly allows us to suggest the non-obvious solution immediately, e.g. .ugrid.is_ugrid instead of isinstance.

[JoerivanEngelen]

Since we are still on a 0.x version, I wouldn't mind to not deprecate, but directly change the API and make this the first 1.0 release. The ideas seem good and solve a bunch of issues. This is the way we want to go forward forward and removing the old classes gives good errors to the users. Adding a bunch of warnings is often not so obvious, because people think warnings are not to be looked at.

I see the removal of two very commonly used classes (UgridDataArray and UgridDataset) as quite a severe API change. I think if we want to push this to the 1.0 release and somewhat control damage for users, we need to think of the following as well:

• Have separate documentation for each xugrid release and a version switcher (e.g. like pandas docs have). There will be packages depending on xugrid (e.g. iMOD Python) that will not update instantly to the 1.0 release, as the changing API of xugrid also affects the internals of these packages. Users of these packages are forced to use an older version of xugrid for a while, and thus need documentation of this older version as well. The 1.0 release docs will not be representative for them, given the API changes.
• Docs how to migrate xugrid 0.x code to 1.x code.

[dcherian]

Hello, we are looking to make a big splash at SciPy. I am currently working on building a gallery of examples for this stuff.

So I was curious about the state of your experiments with the indexes.

Also, we are happy to help! Please reach out with questions if you have any.