Allowing for dataset derived traits

[leifdenby]

The current python api is like this:

# mymodule.myloader
import xarray as xr

TIME_PROFILE = "observation"
SPACE_PROFILE = "grid"
UNCERTAINTY_PROFILE = "deterministic"

def load_dataset(paths: list[str], **kwargs) -> xr.Dataset:
    ds = xr.open_mfdataset(paths, combine="by_coords", **kwargs)
    return ds

# mlwp_data_loaders.cli
from mlwp_data_loaders import load_dataset
from mlwp_data_specs import validate_dataset

# 1. Load the dataset and extract the trait profiles defined by the loader
ds, dataset_traits = load_dataset(
    [
        "/path/to/file1.nc",
        "/path/to/file2.nc",
    ],
    loader="mymodule.myloader",
    return_dataset_traits=True,
)

# 2. Get a detailed validation report by passing the extracted traits
report = validate_dataset(
    ds,
    time=dataset_traits.get("time_profile"),
    space=dataset_traits.get("space_profile"),
    uncertainty=dataset_traits.get("uncertainty_profile"),
)

In this design the load_dataset implementation in mlwp-data-loaders reads the statically defined trait properties from the loader module provided, and return the loaded dataset and the traits, separately.

But this has the issue that the traits that a loader imposes is hard-coded in the loader module, but @mpvginde has pointed out we might want to support allowing the loader to define the traits of a dataset at runtime based on the contents what the loader is reading from disk. For example the current "loader" implementation in mxalign https://github.com/mlwp-tools/mxalign/blob/e2232d93275c7508897a7ddb0cce8b508665f24c/src/mxalign/loaders/base.py#L81-L107

If we instead want to loader to infer the dataset traits from what it reads from disk then:

1. the traits for a specific dataset need to be defined at runtime
2. the traits need to be somehow return from the loader

I will use this issue to outline a few approaches to achieving this

[leifdenby]

An implementation that would allow for dynamic dataset-derived traits could be expecting the user-provided load_dataset() function to return the dataset and the traits, e.g.:

# mymodule.myloader
import xarray as xr

def load_dataset(paths: list[str], **kwargs): -> tuple(xr.Dataset, Dict[str,str])
    ds = xr.open_mfdataset(paths, combine="by_coords", **kwargs)
    traits = dict(time="observation", space="grid", uncertainty="deterministic")
    return ds, traits

# mlwp_data_loaders.cli
from mlwp_data_loaders import load_dataset
from mlwp_data_specs import validate_dataset

# 1. Load the dataset and extract the trait profiles defined by the loader
ds, dataset_traits = load_dataset(
    [
        "/path/to/file1.nc",
        "/path/to/file2.nc",
    ],
    loader="mymodule.myloader",
    return_dataset_traits=True,
)

# 2. Get a detailed validation report by passing the extracted traits
report = validate_dataset(
    ds,
    time=dataset_traits.get("time_profile"),
    space=dataset_traits.get("space_profile"),
    uncertainty=dataset_traits.get("uncertainty_profile"),
)

Here mlwp_data_loaders.load_dataset(...) instead expects the user-provided loader module (where their load_dataset() function is defined) to return the traits as a separate part.

[leifdenby]

Two alternatives would be to:

1. expect the dataset loader to set the traits as global attributes on the returned dataset
2. alter mlwm_data_specs.validate_dataset(ds, ...) to look for these attributes by default, unless the user provides a specific trait value to mlwm_data_specs.validate_dataset(ds, time="analysis") in which case we should probably raise an exception (otherwise what is on the validated dataset won't match the content of the dataset, and so downstream for example mxalign can't use these attributes to choose alignment methods)

The benefit of using global dataset attributes is that downstream tools (e.g. mxalign) could directly from the dataset attributes decide what alignment operations are needed, otherwise we need to pass these traits around separately everywhere.

With both of those changes the implementation could look like this:

# mymodule.myloader
import xarray as xr

def load_dataset(paths: list[str], **kwargs): -> xr.Dataset
    ds = xr.open_mfdataset(paths, combine="by_coords", **kwargs)
    ds.attrs["mlwp_time_trait"] = "observation"
    ds.attrs["mlwp_space_trait"] = "grid"
    ds.attrs["mlwp_uncertainty_trait"] = "deterministic"
    return ds

# mlwp_data_loaders.cli
from mlwp_data_loaders import load_dataset
from mlwp_data_specs import validate_dataset

# 1. Load the dataset with a given loader, expecting the loader to set the trait attributes (prefixed with `mlwp_`)
ds = load_dataset(
    [
        "/path/to/file1.nc",
        "/path/to/file2.nc",
    ],
    loader="mymodule.myloader",
)

# 2. Get a detailed validation report (will look for traits as global attributes by default)
## validate dataset using attributes only
report = validate_dataset(ds)

## validate while overriding a trait already set via a dataset attribute
## (in which case we should maybe raise an exception since otherwise it doesn't 
## make sense for say `mxalign` to use this attribute to determine how to make alignment
## or maybe override the trait attribute and warn about doing this)
report = validate_dataset(
    ds,
    uncertainty="ensemble",
)

Or instead of raising an exception we should maybe override the traits attributes on the dataset and warn if the trait that is being set has a different value than the existing one? That would allow us to have a multiple overlapping traits (maybe we will add regular_grid while keeping gridded for example) and easily switch the traits. Or set them in cases where the loader hasn't set the trait attributes. Or maybe we need to enforce that the loader always sets all three of time, space and uncertainty traits?

[mpvginde]

Hi Leif, thanks for putting this together. I'm in favour of the structure in your last proposal.
Some thoughts:

• I would enforce that a loader always sets all three traits
• But allow for overriding, however I would do this in the (api-)loader part. The validator should best remain passive function.
• I would personally add the validation inside the core load_dataset. That way you always get a validated dataset when loading. (But from the code, that is already how it is done now.)

[leifdenby]

I would enforce that a loader always sets all three traits

ok, that is then assuming that all three traits must always be defined, but maybe that is reasonable, I'm not entirely sure. But we could start with this.

• But allow for overriding, however I would do this in the (api-)loader part. The validator should best remain passive function.

I'm not sure what you mean by "do this in the (api-)loader part", can you give a code example? If we allow for overriding the arguments like I suggested above:

# mymodule.myloader
import xarray as xr

def load_dataset(paths: list[str], **kwargs): -> xr.Dataset
    ds = xr.open_mfdataset(paths, combine="by_coords", **kwargs)
    ds.attrs["mlwp_time_trait"] = "observation"
    ds.attrs["mlwp_space_trait"] = "grid"
    ds.attrs["mlwp_uncertainty_trait"] = "deterministic"
    return ds

# mlwp_data_loaders.cli
ds = load_dataset(
    [
        "/path/to/file1.nc",
        "/path/to/file2.nc",
    ],
    loader="mymodule.myloader",
)

report = validate_dataset(
    ds,
    uncertainty="ensemble",
)

where I am changing the uncertainty trait to have the value ensemble. If the dataset validates successfully does it really make sense to have the returned dataset have the old value for ds.mlwp_uncertainty_trait (i.e. this case the value would be deterministic)? That would mean that a second run with the validator, without providing arguments to validate_dataset(...) to override a specific trait value, would then fail the validator (since the trait value is unchanged). Maybe what we should say that if the validation passes for a given trait then the trait attribute is updated?

• I would personally add the validation inside the core load_dataset. That way you always get a validated dataset when loading. (But from the code, that is already how it is done now.)

I can see why this would be tempting, but I don't think this is a good idea for three reasons:

1. We would then be expecting the user to always call validate_dataset inside any load_dataset() function they implement. They might forget to do this.
2. Adding a validation call into every load_dataset() implementation leads to many duplicate calls, instead mxalign can simply call validate_dataset() once with any dataset returned by a loader
3. I think this break separation of concerns. Instead of only loading a dataset the user-defined loader functions now will actually load-and-validate a dataset. This blurs the line between what piece of code is responsible for what. To me it makes more sense that the responsibility of the "loader" is only to load a dataset, the job of ensuring that the dataset is "valid" is the responsibility of mxalign and will achieve this by calling the validator.

[leifdenby]

After discussion today we have decided the API will be so that mlwp-data-loaders both loads (maybe from a user-provided loader) and validates a dataset, so that the API looks like this:

# mymodule.myloader
import xarray as xr

def load_dataset(paths: list[str], **kwargs): -> xr.Dataset
    ds = xr.open_mfdataset(paths, combine="by_coords", **kwargs)
    ds.attrs["mlwp_time_trait"] = "observation"
    ds.attrs["mlwp_space_trait"] = "grid"
    ds.attrs["mlwp_uncertainty_trait"] = "deterministic"
    return ds

# inside the user's script/notebook
from mlwp_data_loaders import load_and_validate_dataset
from mlwp_data_specs import validate_dataset

# Load the dataset with a given loader, expecting the loader to set the trait attributes (prefixed with `mlwp_`)
# will raise exception on validation error, if the report has validation errors
ds = load_and_validate_dataset(
    [
        "/path/to/file1.nc",
        "/path/to/file2.nc",
    ],
    loader="mymodule.myloader",
)

# will not raise an exception and instead returns the report (the user must check the report contents for errors)
ds, validation_report = load_and_validate_dataset(
    [
        "/path/to/file1.nc",
        "/path/to/file2.nc",
    ],
    loader="mymodule.myloader",
    return_validation_report=True
)

To ensure that if the user changes a validated xr.Dataset that the dataset still conforms to the mlwp_ global attributes it has we will run the validation before (and after, as it is now) every mxalign operation.

[leifdenby]

I implemented the dynamic trait from dataset functionality in #6 and this is now available on main :D