Creating xr.DataArray templates with constrained dims?

[tylerflex]

I want to create several templates that define xr.DataArray objects constrained to have specific dims that are defined at the class level.

For example, I might want to define an ImageDataArray that always has dims (x, y) or a TimeSeriesDataArray with only dims ('t'). These templates can then be used as type information in other models and if one tries to construct one with the wrong dims, they should fail gracefully.

What is the recommended way to handle this sort of behavior in xarray? There are two approaches I've tried, but each has some drawbacks.

A naive approach is to subclass xr.DataArray and provide some logic to check the dims.

import xarray as xr
class ImageDataArray(xr.DataArray):

    __slots__ = ('x', 'y')
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # do some checking to make sure that self.dims == __slots__?

However, subclassing xr.DataArray is not recommended, as explained here. And in my experience, it is indeed difficult to maintain.

Another option is to define classes that contain the dim information as well as a regular xr.DataArray attributes, and do a bit of checking. For example:

import xarray as xr

class ImageDataArray:
    dims = ('x', 'y')

    def __init__(self, data_array):
        assert data_array.dims == self.dims, "wrong dimensions"
        self.data_array = data_array

This works, but then introduces the need for many intermediate calls to .data_array as the ImageDataArray class itself is not very useful in my application.

Any thoughts or discussion appreciated, thanks!

[kmuehlbauer]

You might have a look at https://github.com/xarray-contrib/xarray-schema which provides Dataset/DataArray validation.

[tylerflex]

Looks interesting, so using this package, one would probably still need to define a separate template to store the xr.DataArray and its schema? For example

class ImageDataArray:
import xarray as xr

class ImageDataArray:
    schema = DataArraySchema(dtype=float',  dims=['x', 'y'])

    def __init__(self, data_array):
        self. schema.validate(data_array)
        self.data_array = data_array

similar to my 2nd example above?

[benbovy]

Maybe you don't need the ImageDataArray class? It is not very useful if it's just for saving the schema somewhere, you could simply define image_schema = DataArraySchema(...) and use image_schema.validate(...) where you need it in your application.

[tylerflex]

In the broader picture, I want to use these class templates to define specific types of DataArrays that live in pydantic / data class models. So basically I do want to have the some of the template information baked into the schema definition of my model. For example:

Model(pydantic.BaseModel):
    image: ImageDataArray
    time_series: TimeSeriesDataArray

And if someone got the schema of Model, it would clearly show the different types of data arrays. Also, supplying something with the wrong dims to image would error.

Now I could just specify these fields as raw xr.DataArray and then validate later on, but then the schema information is lost in my Model.

[benbovy]

Although I haven't tried it, perhaps you could add the schemas as class attributes and still declare fields as DataArrays with custom validators?

from typing import ClassVar

import xarray as xr
from pydantic import BaseModel, validator
from xarray_schema import DataArraySchema


class Model(BaseModel):
    image_schema: ClassVar[DataArraySchema] = DataArraySchema(dims=["x", "y"])
    image: xr.DataArray

    @validator("image")
    def validate_image(cls, value):
        cls.image_schema.validate(value)
        return value

(Like your snippet #7077 (reply in thread), but in a pydantic model class)

[benbovy]

And if someone got the schema of Model, it would clearly show the different types of data arrays.

Maybe you could use TypeVar for that?

from typing import TypeVar
...

T_DataArrayImage = TypeVar("T_DataArrayImage", bound=xr.DataArray)

class Model(BaseModel):
    ...
    image: T_DataArrayImage

    ...