Machine-readable schema & validator for xarray.Dataset

[JackKelly]

Detailed Description

If we can find an off-the-shelf schema & validator for xarray.Dataset then we can, hopefully, combine the best of pydantic.BaseModel and xarray.Dataset. The ultimate aims are:

• Have a single source of truth for the precise structure of the data that flows through nowcasting_datatset. This can be used for:
  • Humans to understand the structure of the data
  • Machines to automatically validate the data

Context

@peterdudfield has done excellent work in pull request #195 using Pydantic to define schemas for our data. Inspired by, and building on @peterdudfield's great work, it's possible that we can get the same advantages by using something like pandera, but with less effort on our part :) (I'm lazy!)

This is also related to #209

Related

• Pandera provides machine-readable schemas for Pandas and, as of version 0.7, panderas has decoupled pandera and pandas types to make pandera more useful for things like xarray.
• Twitter discussion
• An xarray github issue discussing schemas (which was closed four years ago without a conclusion)

I'll look into this this morning :)

[JackKelly]

Hmm, I'm no longer confidence that panderas is capable of supporting n-dimensional numpy arrays (and hence can't support n-dimensional xr.DataArrays). I've asked the panderas folks.

But it looks like Pydantic might be able to. See this Pydantic issue and PR. Also, take a look at how SQLModel combines Pydantic and SQLAlchemy (thans to Benoît Bovy for suggesting this on twitter!)

[JackKelly]

Let me flesh out what I hope is possible (but I haven't tested yet!) This is adapted from pydantic/pydantic#667, and inspired by @peterdudfield's work in PR #195. (Also see Pydantic's docs on custom data types)

class PydanticXArrayDataset(xr.Dataset):
    """Abstract base class for validating xr.Dataset objects."""
    # From https://github.com/samuelcolvin/pydantic/issues/667

    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v: Any) -> str:
        """Validate data.  Must be overridden by child classes."""
        raise NotImplementedError()


class Satellite(PydanticXArrayDataset):
    @classmethod
    def validate(cls, v: Any) -> str:
        # validate Satellite data...
        return v

The above code is all that's required (I think) when pre-preparing on-disk batches, because, after #202 is implemented, the individual modalities wouldn't be squished together into a single batch object: Instead each modality would pass through nowcasting_dataset independently, and be written to disk independently.

When we load the batches of each modality from disk, then we could squish them into a Pydantic model like the code below, but I'd be a little worried about hurting performance, especially when all the pre-prepared batches should have been validated when they were created!

class Example(pydantic.BaseModel):
    satellite: Satellite

[JackKelly]

OK, here's a functional, but very rough example of using xarray with pydantic.

This code validates a few things. But isn't ideal as a human-readable specification of the structure of xr.DataArrays or xr.Datasets.

[JackKelly]

To quote @cosmicBboy from this comment:

once pandera schemas can be used as valid pydantic types, #453 is supported, the solution you outline [two comments above] would be pretty straightforward to port over to pandera

[JackKelly]

I think this is implemented in PR #229 (thanks @peterdudfield!)