Transpose ufunc parameters

[aazuspan]

The current ufunc API roughly follows a simplified version of xr.apply_ufunc and dask.apply_gufunc, i.e. parameters are specified as lists and/or dicts that are mapped positionally to outputs. For example, a ufunc that returns three single-band outputs is defined as (using the new FeaturewiseUfunc from #103):

ufunc = FeaturewiseUfunc(
    calculate_ogsi,
    output_dims=[["band"], ["band"], ["band"]],
    output_dtypes=[np.float32, np.uint8, np.uint8],
    output_sizes={"band": 1},
    output_coords=[{"band": ["OGSI"]}, {"band": ["LTQ"]}, {"band": ["OGSI_CLS"]}],
)

While this is reasonably familiar for Xarray ufunc users, it requires reading positionally across parameters to understand each output (i.e. the first output is composed of the first output dimension, the first dtype, and the first set of output coordinates). It's also easy to specify incorrectly by nesting improperly or omitting an element, and difficult to type check.

@grovduck and I discussed a more intuitive API here which would transpose from parameters as groups of outputs to outputs as groups of parameters. For example, the same function could be defined using some new dataclasses as:

@dataclass
class Dimension:
    name: str
    size: int | None = None
    coords: list[str | int] | None = None

@dataclass
class Output:
    dims: list[Dimension]
    dtype: type[np.generic] | None = None
    

ufunc = FeaturewiseUfunc(
    calculate_ogsi,
    outputs=[
        Output(dims=[Dimension(name="band", size=1, coords=["OGSI"])], dtype=np.float32),
        Output(dims=[Dimension(name="band", size=1, coords=["LTQ"])], dtype=np.uint8),
        Output(dims=[Dimension(name="band", size=1, coords=["OGSI_CLS"])], dtype=np.uint8),
    ]
)

This change means that each output is defined independently, making them easier to validate, type, and reuse.

One drawback (aside from losing the familiarity of the Xarray API) is that Dask ufuncs don't support specifying a different size for the same dimension with multiple outputs, so this would require a posthoc validation for that case.

[aazuspan]

Summarizing a short discussion with @grovduck from #103 so that we can continue it here.

NoData output specification

With #103 and the changes proposed above, ufunc parameters are divided into constructor parameters that define the ufunc (the function and expected output dimensions, sizes, dtypes, and coordinates), and call parameters that modify a given run (whether to skip NoData, allow casting, limit threads, etc).

nodata_output is currently a call parameter because it modifies the value stored as NoData in outputs, independent of the ufunc implementation, but @grovduck raised a good point here that with the modified API above, it would be very intuitive to define a NoData value per output, e.g.

ufunc = FeaturewiseUfunc(
    calculate_ogsi,
    outputs=[
        Output(dims=[Dimension(name="band", size=1, coords=["OGSI"])], dtype=np.float32, nodata=np.nan),
        Output(dims=[Dimension(name="band", size=1, coords=["LTQ"])], dtype=np.uint8, nodata=255),
        Output(dims=[Dimension(name="band", size=1, coords=["OGSI_CLS"])], dtype=np.uint8, nodata=255),
    ]
)

This is somewhat less flexible/reusable since you now have to redefine a ufunc to change the NoData output encoding, but substantially more convenient since NoData values are defined once and are visibly tied to the outputs. A third option would be to allow both, defining a default NoData in the ufunc specification which could be overridden during the call, at the cost of some additional complexity.