Representing vector quantities in xarray

[jthielen]

In both following along with the development of xr.cross (#5365) and having some recent discussions downstream in MetPy (Unidata/MetPy#2079), I wanted to raise (or re-raise) the idea/suggestion of having some explicit form of vector data object in xarray (or a suitable xarray-contrib extension package).

I would imagine it behaving like a Dataset, but with additional constraints (e.g., identical (or at least coupled) dimensions and chunking across contained variables) and additional operations (e.g., tuple unpacking, presently DataArray-only things like xr.dot and xr.cross working as expected). In this way, it could fit naturally as a subclass of/wrapper around the DataNode proposed in #4118 (comment), as this would ideally both:

• fit inside a Dataset (so can't just be a Dataset as it currently exists)
• have each component with its own attrs (so can't just be a DataArray with a stack dimension)

This would allow for some rather streamlined code when doing vector analysis. For example, with a dataset/datatree like

DataNode("root")
|-- VectorDataNode("wind")
|   |-- DataArray("u")
|   |-- DataArray("v")
|-- DataArray("temperature")

horizontal advection can be expressed as

xr.dot(data['wind'], gradient(data['temperature']))

where gradient is like that of MetPy or a simple vector-based wrap of da.differentiate().

What are everyone's thoughts on this?

xref Unidata/MetPy#537

[keewis]

It doesn't answer the question on how to best represent vector quantities, but:

have each component with its own attrs (so can't just be a DataArray with a stack dimension)

I think this might be possible by choosing a appropriate attrs handler for .to_array / .to_dataset: if the one for .to_array puts each variable's attrs into a dict of name → attrs value, then the one for .to_dataset can restore the attrs for each component. That might require thinking about how to hook into just a specific operation if we don't want to require every attrs handler to be universal.

[rabernat]

Just saw this issue. I wanted to note that it also comes up in xgcm, where we have support for finite-volume calculus.

Where we need to keep track of vectors, we put each component in a dictionary.

https://xgcm.readthedocs.io/en/latest/api.html#xgcm.Grid.diff_2d_vector

This is not very satisfactory.

[keewis]

we discussed this in a meeting a few weeks ago, but back then we concluded that it would probably be fine to use the definition from xr.cross (require a DataArray with a special dimension) and have a external library do everything else. That could, for example, include the attrs handler from #5775 (comment) and a combination of variables based on certain attributes.

There was nobody with a lot of experience with vector data, though, so we might have missed something.

[jthielen]

we discussed this in a meeting a few weeks ago, but back then we concluded that it would probably be fine to use the definition from xr.cross (require a DataArray with a special dimension) and have a external library do everything else. That could, for example, include the attrs handler from #5775 (comment) and a combination of variables based on certain attributes.

I agree that this looks to be the best that can be done for now, and can mostly see how such an external library could function for basic use cases (e.g., stacking/unstacking for IO, forcing the component special dimension to be leading to allow tuple unpacking, custom dict-like objects as the attributes for nicer reprs, etc.). However, one uncertainty: if the differing metadata across vector components is to be stored as a single object, would this not require that indexing/selecting operations provide some hook so that the attrs can also be split?

That being said, if/when DataNodes become an option, I'd still be in favor of leveraging that for vectors instead. First, there a couple nice things the "DataArray with fancy attrs" can't do:

• staggered grids (the components are on different, but coupled, dimensions)
• components having different Pint Quantity units

Second, I'm not a big fan of the semantics of the DataArray approach, as I feel like each component of the vector should act like a data variable in its own right (e.g., I want to be able to say ds["u"] instead of ds["wind"].sel(component="x"), or when exploring the vector repr, I'd want each components' metadata separated). Also, it would take a lot more behind-the-scenes magic to make IO work smoothly (rather than mapping neatly to NetCDF-style groups like the DataNode approach would).

[keewis]

do you have any example datasets for staggered grids (reprs might be enough)? I never worked with those so I don't really understand how they're different.

it would take a lot more behind-the-scenes magic to make IO work smoothly

here, too, an example would be great because the datasets I know have everything in a single group so open_dataset. Or did you mean that converting variables of a Dataset to a datatree seems less work than adding support for combining / separating the attrs?

The most important disadvantage I think is that some algorithms (like np.cross, but also the matrix multiplication used for rotating vectors with transformation matrices) assume that the data is stacked, so I believe you will have to stack at some point, even if you use DataNode objects.

Actually, wouldn't it be better to only convert when needed? With your example, you'd use ds["u"] for component wise access, but if you need the vector, you'd call something like ds.vector[["u", "v"]] (or ds.vector["wind"], where "wind" is a virtual variable based on the attrs).

[rabernat]

Staggered grids are the main focus of the xgcm package. You can read more about them here: https://xgcm.readthedocs.io/en/latest/grids.html

I like to keep reminding people of xgcm whenever staggered grids come up because we have put a lot of work into it! Several years ago, we concluded that staggered grids couldn't be supported by Xarray's limited data model, so we created an external package. That situation may change with the introduction of flexible indexes. However, xgcm does things besides just indexing--it supports "grid-aware" finite difference, cumsum, and interpolation.

[jthielen]

here, too, an example would be great because the datasets I know have everything in a single group so open_dataset. Or did you mean that converting variables of a Dataset to a datatree seems less work than adding support for combining / separating the attrs?

Kind of both! On the input/read side, I was indeed meaning the conversion to datatree being easier, as most source datasets I'm used to indeed just have a single root group. On the write/output side, it was that writing a datatree to netcdf/zarr is easier (should "just work" with the built-in to_* methods, rather than the custom split attr handling that would require every write operation to be wrapped).

The most important disadvantage I think is that some algorithms (like np.cross, but also the matrix multiplication used for rotating vectors with transformation matrices) assume that the data is stacked, so I believe you will have to stack at some point, even if you use DataNode objects.

I agree! I've been under the assumption that stacking and unstacking are not generally expensive operations, so given that I see roughly equal weight on needing the arrays themselves to be stacked or not, in addressing the problem I've been prioritizing what structure makes the most sense conceptually and gives the easier management of metadata (with the actual array stacking and/or unstacking happening behind-the-scenes). But, if that assumption is mistaken, then I'd have to re-evaluate!

Actually, wouldn't it be better to only convert when needed? With your example, you'd use ds["u"] for component wise access, but if you need the vector, you'd do ds.vector[["u", "v"]] (or ds.vector["wind"], where "wind" is a virtual variable based on the attrs).

I hadn't thought of that possibility; that could work! Though, assuming that the return type of this VectorAccessor.__getitem__ is a stacked DataArray (as it being a Dataset doesn't give it much benefit over just subsetting), staggered grids are still problematic. That could be worked around, however, by allowing Dataset returns for staggered grids, and then placing vector methods behind that same assessor (as they'd need special versions anyway to handle the staggering/inability to stack).

I guess in all this discussion, it's becoming clear that I'm just wishing for a data object that blends Dataset and DataArray and transparently behaves like the right one depending on the interaction. Functionally, that vector accessor or convert-when-needed approach seems like it could do the trick if fitting in to the existing data model is a priority. So, I think that makes this problem a more philosophical decision between:

• adding another type to the data model (and thereby hide/encapsulate the details of stacked versus unstacked-but-grouped arrays, allowing just one "vector API")
• use accessors and codified metadata (and thereby keeping stacked [DataArray-like] and unstacked-but-grouped [Dataset-like] arrays clearly separate, but still providing matching APIs in an isolated namespace)

Either way, it seems like an extension package (xvector?) would be the home for this? As much as I'd love to take the lead on something like this, my prior record on being able to consistently commit to extension package development has been poor. So, that being said, if someone wants to spin such a package up, I'd be glad to help brainstorm and contribute where able. If not, this will go on my open-source wish-list to hopefully get around to working on at some point!

[rafa-guedes]

What about using complex data types?

[keewis]

could you elaborate? Do you mean the pandas extension arrays or numpy's (experimental) dtype feature, or the complex dtype?

[jthielen]

I don't think this would work, as it would leave the attrs/metadata problem the same as in the stack dimension case and be less flexible than an arbitrary stack dimension.

[rafa-guedes]

I had the numpy complex dtype in mind. I thought complex numbers would be a nice way to handle it in a natural way - but yes, it does not solve the metadata problem. There are also issues when saving it in netcdf (although zarr seems to deal with that well).

[SimonHeybrock]

@jthielen Late to the party, and not having fully understood all the requirements, but isn't this something that could be handled using a duck-array? Say we need an array of 3-D vectors, one could imagine something like:

class Components:

    def __init__(self, obj):
        self._obj = obj

    def __getitem__(self, key):
        return self._obj[self._components.index(key)]


class VectorArray:

    def __init__(self, values: np.ndarray, components: List[str]):
        # Assuming a structure-of-array implementation
        assert len(values) == len(components)
        self._values = values
        self._components = components

    @property
    def shape(self):
        return self._values.shape[1:]

    def __getitem__(self, index):
        return VectorArray(self._values[:, index], self._components)

    @property
    def components(self):
        return Components(self)

    def __array_ufunc__(self):
        """TODO"""

    def __array_function__(self):
        """TODO"""

What is missing above is an interface ensures that gives us components with named dimensions. This can be solved using a custom xarray.DataArray accessors, but I don't feel that is a pretty solution.

This is a problem I am experiencing in other places as well --- duck arrays and the Array API specification lack named dimensions, but this would be extremely valuable when writing a duck-array. One could use xarray.Variable, but this comes with extra stuff not needed in such cases, and more critically, leads to chicken-egg problems when putting the duck array into and xarray.DataArray. I think what I am saying is that I am dreaming of a lightweight Array API "standard" with dimension labels.

[SimonHeybrock]

I had to lookup what you mean by staggered, if I understand correctly this refers to storing value on bin/cell boundaries instead of bin/cell centers? The xgcm example in #6331 seems to imply this?

Interestingly, this links into another experiment I started last week: Trying to use a duck-array for better representation of bin-edges in Xarray. Basically, I am wrapping an np.ndarray of shape=(N+1,) in a BinEdgeArray that gives shape=(N,) and can thus be used with an xr.DataArray of the latter shape. BinEdgeArray could provide properties such as left and right to return the views to the corresponding values in the underlying ndarray. The concept could be generalized to 2-D or 3-D for cell edges or faces.

One could then combine this with something like the VectorArray outlined above.

I don't know yet how well this would work in practice, but it is an attempt to abstract those independent concepts and minimize complications in higher-level code (i.e., in Xarray).

[jthielen]

I had to lookup what you mean by staggered, if I understand correctly this refers to storing value on bin/cell boundaries instead of bin/cell centers? The xgcm example in #6331 seems to imply this?

Yep!

Interestingly, this links into another experiment I started last week: Trying to use a duck-array for better representation of bin-edges in Xarray. Basically, I am wrapping an np.ndarray of shape=(N+1,) in a BinEdgeArray that gives shape=(N,) and can thus be used with an xr.DataArray of the latter shape. BinEdgeArray could provide properties such as left and right to return the views to the corresponding values in the underlying ndarray. The concept could be generalized to 2-D or 3-D for cell edges or faces.

One could then combine this with something like the VectorArray outlined above.

I don't know yet how well this would work in practice, but it is an attempt to abstract those independent concepts and minimize complications in higher-level code (i.e., in Xarray).

That's certainly a novel design approach! As someone whose mental model of representing multidimensional data is very much informed by NetCDF/Unidata CDM/CF conventions/xarray, this breaks some of the assumptions I've had about the problem in regards to what data object properties represent what and how they relate. So, I'd definitely be interested in hearing about where those experiments lead if they offer a more promising approach given the existing state of the ecosystem. I'd also be concerned that adding more duck array layers to the picture will add more complications rather than minimize them, but that could certainly be proven otherwise as well.

[SimonHeybrock]

I'll try to make a tiny prototype.

As you point out, layering duck-arrays can be awkward, I don't expect users (including myself) would be happy with accessing things as da.data.left.fields['u'] (or was it da.data.fields['u'].left? ). I noticed that pint.Quantity tries to forward to attributes of the wrapped duck array, but I am not sure that approach scales well to multiple layers of duck arrays?

[jthielen]

I'll try to make a tiny prototype.

Sounds good!

As you point out, layering duck-arrays can be awkward, I don't expect users (including myself) would be happy with accessing things as da.data.left.fields['u'] (or was it da.data.fields['u'].left? ). I noticed that pint.Quantity tries to forward to attributes of the wrapped duck array, but I am not sure that approach scales well to multiple layers of duck arrays?

Indeed, this is something that hasn't been decided on by the community (xref pydata/duck-array-discussion#6 for general topic of "carrying through" these sort of methods and attributes).

[SimonHeybrock]

My work on a prototype escalated into a potential solution for two of the duck-array discussions you linked to. See pydata/duck-array-discussion#6 (comment) for more!

[dcherian]

Trying to use a duck-array for better representation of bin-edges in Xarray. Basically, I am wrapping an np.ndarray of shape=(N+1,) in a BinEdgeArray that gives shape=(N,) and can thus be used with an xr.DataArray of the latter shape. BinEdgeArray could provide properties such as left and right to return the views to the corresponding values in the underlying ndarray.

I tried this too by wrapping a pandas ExtensionArray (IntervalArray). I had trouble getting it all to work well with Xarray though. Really I want IntervalIndex support in Xarray, so that was my motivation to try IntervalArray