astropenguin / xarray-dataclasses

:zap: xarray data creation by data classes
https://astropenguin.github.io/xarray-dataclasses/
MIT License
71 stars 4 forks source link

Add an optional data field #223

Open JamieMcMillan opened 5 months ago

JamieMcMillan commented 5 months ago

In the same vein as the RGB image example, I would like to create a dataclass that can accommodate one or more optional channels. I expect to replace missing channels with (where downstream methods would operate) arrays of the same size as the given array with NaNs.

Provided below is an MWE. Please advise if I am misunderstanding how this module should be used or some other obvious solution. I am happy to add back to the documentation with a solution :)

from dataclasses import dataclass
from typing import Literal
from xarray_dataclasses import AsDataset, Data, Dataof, Name

X = Literal["x"]
Y = Literal["y"]

@dataclass
class Red:
    data: Data[tuple[X, Y], float]
    name: Name[str] = "Red image"

@dataclass
class Green:
    data: Data[tuple[X, Y], float]
    name: Name[str] = "Green image"

@dataclass
class Blue:
    data: Data[tuple[X, Y], float]
    name: Name[str] = "Blue image"

@dataclass
class ColorImage(AsDataset):
    """2D color image as Dataset."""

    red: Dataof[Red]  # Always provided
    green: Dataof[Green]  # Optional channel
    blue: Dataof[Blue]  # Optional channel

def main() -> None:
    ColorImage.new(  # successful image creation
        [[0, 0], [0, 0]],
        [[1, 1], [1, 1]],
        [[2, 2], [2, 2]],
    )

    ColorImage.new(  # failing image creation
        [[0, 0], [0, 0]],
        [[1, 1], [1, 1]],
        None,
    )

if __name__ == "__main__":
    main()