Cross-array constraints for Xarray

[TomNicholas]

Hi, we met at SciPy, this package looks very cool! I'm impressed with the amount of thought that seems to have gone into the upstream LinkML work too.

My main question after skimming all the docs is: can this system specify constraints across multiple arrays? That's necessary to support Xarray correctly.

I see you already support Zarr & HDF5, but Xarray's data model is subtly different. Fr example it's possible to create a Zarr group which cannot be converted to an xarray.Dataset.

Concretely, if I want to for example check that the schema of a given zarr store is something that is compatible with xarray, I need to check that within a group, for every occurrence of any given dimension name, the length of that dimension is equal. This is a cross-array constraint - the correctness can't be determined by looking at any one array, as it depends on all the other arrays. (In xarray parlance we would say that the array's dimensions are aligned.)

You don't need this feature for a single xarray.Variable, but you do need it for both xarray.DataArray and xarray.Dataset. xarray.DataTree effectively has even more complicated constraints, that are now cross-group as well as cross-array.

Can these types of constraints be supported in your system?

[TomNicholas]

In other words, there doesn't seems to be an example on this page of how to say "a set of arrays with named dimensions such that all occurrences of the same name possess the same length".

[sneakers-the-rat]

good question :) Let's focus on the single-class cross-array validation for now, and then think about tree-like and cross-model cross-array validation after we get a basic sketch.

this:

I need to check that within a group, for every occurrence of any given dimension name, the length of that dimension is equal.

is hard to specify in a bunch of schema languages, and esp. relevant here is JSON-Schema because of the expectation that pydantic models have correct JSON Schema, but it doesn't support the necessary semantics - see https://github.com/orgs/json-schema-org/discussions/730

however we can certainly implement it here even without correct JSON-Schema generation.

Upstream in the linkml arrays schema, labeled/xarray-style array spec are still a work in progress, but we have the notion of labeled_by (though after seeing some of the developments in xarray last week, we probably want to revisit whether to call that indexed_by) that would imply that the lengths of labeling arrays/coords need to be equal.

What would be a syntax you would like for this?

The current syntax is a holdover from nptyping that is destined to be replaced by type annotations that are actually valid python types, that has the generic structure like

NDArray[Shape["{cardinality} {alias}, ..."], dtype]

This will likely be replaced by something like

NDArray[dtype, Dim(cardinality, alias), Dim(cardinality, alias), ...]

which should make it easier to make parameterizations like this - i am sure nobody loves string-based microformats.

So maybe a few ideas...

As a dependency

So we could have some of the dimension sizes be allowed to vary, and then specify dependencies between them...

string syntax

(purposely making the names sort of awkward to highlight separable parts of style):

class MyModel(BaseModel):
    index_a: NDArray[Shape["* x"]]
    index_b: NDArray[Shape["* y"]]
    array: NDArray[Shape["index_a.x xalias, index_b.y yalias"]]

Where one could specify that the cardinality of the first dimension xalias of array should be equal to the shape of the x dimension of index_a, and so on. In the case that the other arrays were 1D, we could add a simplifying syntax like this that doesn't require specifying the alias of the single dimension

class MyModel(BaseModel):
    index_a: NDArray
    index_b: NDArray
    array: NDArray[Shape["index_a xalias, index_b yalias"]]

which would fail with a validation error if index_a or index_b had more than 1 dimension because it would be ambiguous.

Dataclass syntax

In the other, as-yet-still-being-drafted style that doesn't rely on a very cursed string microsyntax, that might look like

class MyModel(BaseModel):
    index_a: NDArray[Dim(alias="x")]
    index_b: NDArray[Dim(alias="y")]
    array: NDArray[
        Dim(shape=Eq(field="index_a", dim="x")),
        Dim(shape=Eq(field="index_b", dim="y"))
    ]

where we use specific operators like Eq to say "shape should be equal to this other thing" because we may want to have "should be greater, less than, etc."

As a variable

I sort of think this would be a bit harder to implement, but e.g. if there isn't an obvious "root" item that would define the dimension cardinality, we could use implicit variables to link things. this is actually something already present in nptyping.

The reason i think this might be a bit more fragile is that it would be harder to handle the more complex cases in trees and etc - in what scope should a variable apply? How would we handle compositional behavior where i might be using a model from someone else, but they use an overlapping variable pattern? It also might be a bit less expressive, where before i was just giving the labeling arrays arbitrary * shapes, but say we want to constrain those to a specific cardinality range too - would one be able to use constraints as well as variables? It also invites expectation of an expression syntax - can i say "shape=SizeVar * 2"?

However the impl as a variable does have some desirable features, like being insensitive to the names of the other arrays, and maybe separating the notion of "these things should be the same" from "these things should have shape constraints" is good to do.

Basically i'm posing these to you as questions to see what would be best and what you would need, because i don't know xarray well enough to know which would be better.

String syntax

class MyModel(BaseModel):
    index_a: NDArray[Shape["Size1 x"]]
    index_b: NDArray[Shape["Size2 y"]]
    array: NDArray[Shape["Size1 xalias, Size2 yalias"]]

Dataclass syntax

class MyModel(BaseModel):
    index_a: NDArray[Dim(shape=Var("Size1"))]
    index_b: NDArray[Dim(shape=Var("Size2"))]
    array: NDArray[
        Dim(shape=Eq(Var("Size1"))),
        Dim(shape=Eq(Var("Size2"))),
    ]

From an implementation POV, we would need to split the validation step into different validators so that we could add an AfterValidator stage that would have access to all of the other fields within the model, but for the single-model, cross-array case it shouldn't be all that complicated to implement once we decide on a syntax.

cc @rly , @melonora re: labeled arrays discussion and implementation.

[sneakers-the-rat]

maybe a tl;dr is "where do the constraints come from" - is there any "directionality" between them that could be treated as a dependency kind of relationship, which is a simplifying idiom, or is it just free-for-all, can be specified in a ton of different ways, etc.

we totally can make a variable-like syntax if that fits xarray better, just want to scope out what we need before trying to write something :)

[sneakers-the-rat]

@TomNicholas not trying to give you homework or anything, but if you throw me a MWE of the kind of thing you need to do, i can get noodling on it :)