feat: Forecast wrapper for custom xarray datasets

[darothen]

EWB Pull Request

Description

This PR implements a new type of ForecastBase, an XarrayForecast which wraps around a previously-opened xr.Dataset. It can be used to interact with datasets which were manually constructed (e.g. by reading many different datasets and concatenating together), or as a way to read an Icechunk dataset which was previously opened.

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)

How Has This Been Tested?

Comprehensive unit tests. This code was adapted from real-world code used for testing icechunk integrations with EWB.

Checklist:

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• Any dependent changes have been merged and published in downstream modules

[darothen]

@aaTman this is working version of my "in-memory" forecast dataset I used for the test evals on our new icechunk-based MLWP archive. Given that icechunk has a few more complexities, it may be easier to open an icechunk archive separately and then use this pass into EWB.

[Copilot]

Pull request overview

This PR introduces a new XarrayForecast class that allows users to wrap pre-opened or manually constructed xarray datasets for use in forecast evaluations. This is particularly useful when working with datasets assembled from multiple sources or when integrating with alternative storage backends like Icechunk.

Changes:

• Added XarrayForecast class to src/extremeweatherbench/inputs.py that extends ForecastBase and stores an in-memory xarray dataset
• Comprehensive test suite with 39 test methods covering instantiation, preprocessing, variable mapping, integration with evaluation objects, and edge cases

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
src/extremeweatherbench/inputs.py	Implements the new XarrayForecast class with validation in __post_init__ and a simple _open_data_from_source method that returns the stored dataset
tests/test_inputs.py	Adds comprehensive test class TestXarrayForecast with 39 test methods covering all functionality and edge cases

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[aaTman]

Love this, simple and elegant. Just some thoughts on making this a bit nicer with upstream changes before merging.

Do you happen to have an example script on hand we can also include to show how its used?

[darothen]

@aaTman here's a hacky version of what my script looked like:

from extremeweatherbench.inputs import XarrayForecast
import icechunk
import xarray as xr
from extremeweatherbench import cases, evaluate

mlwp_model_variable_mapping = {
    "2m_temperature": "surface_air_temperature",
    "2m_dewpoint_temperature": "surface_dewpoint_temperature",
    "2m_relative_humidity": "surface_relative_humidity",
    "2m_wind_speed": "surface_wind_speed",
    "2m_wind_from_direction": "surface_wind_from_direction",
    "2m_wind_gust": "surface_wind_gust",
    "2m_wind_gust_direction": "surface_wind_gust_direction",
    "2m_wind_gust_speed": "surface_wind_gust_speed",
}

def open_icechunk_dataset(
    bucket: str = DEFAULT_ICECHUNK_BUCKET,
    prefix: str = DEFAULT_ICECHUNK_PREFIX,
    variable_mapping: dict[str, str] | None = None,
    chunks: str | dict | None = "auto",
    source_credentials_prefix: str = DEFAULT_SOURCE_CREDENTIALS_PREFIX,
) -> xr.Dataset:
    """Open a dataset from an Icechunk repository with preprocessing.

    The repository config already knows where the virtual chunks are located.
    We just need to provide credentials broad enough to cover that location.

    Args:
        bucket: GCS bucket containing the Icechunk repository.
        prefix: Prefix within the bucket for the repository.
        variable_mapping: Dictionary mapping source variable names to target names.
        chunks: Chunk specification for xarray (default: "auto").
        source_credentials_prefix: GCS prefix for virtual chunk credentials.
            Should be broad enough to cover wherever the source data lives.

    Returns:
        Preprocessed xarray Dataset ready for evaluation.
    """
    logger.info(f"Opening Icechunk repository at gs://{bucket}/{prefix}")

    # Set up storage
    storage = icechunk.gcs_storage(bucket=bucket, prefix=prefix)

    # Set up credentials for virtual chunks.
    # The repo config knows the exact location; we just provide credentials
    # broad enough to cover it.
    gcs_credentials = icechunk.gcs_from_env_credentials()
    virtual_credentials = icechunk.containers_credentials({source_credentials_prefix: gcs_credentials})

    # Open repository
    repo = icechunk.Repository.open(storage, authorize_virtual_chunk_access=virtual_credentials)
    session = repo.readonly_session("main")

    # Open dataset
    ds = xr.open_dataset(session.store, engine="zarr", chunks=chunks)

    # Apply variable renaming if specified
    if variable_mapping:
        rename_dict = {k: v for k, v in variable_mapping.items() if k in ds.data_vars}
        if rename_dict:
            ds = ds.rename(rename_dict)

    return ds


if __name__ == "__main__":

    ds = open_icechunk_dataset(
        bucket=icechunk_bucket,
        prefix=icechunk_prefix,
        variable_mapping=mlwp_model_variable_mapping,
        chunks="auto",
        source_credentials_prefix=source_prefix,
    )


    forecast = XarrayForecast(
        ds,
        name=f"{icechunk_bucket}/{icechunk_prefix}",
        # NOTE: we have to pass in the variables that will actually be used for this
        # metrics calculations. We can at least bypass the variable mapping by manually
        # processing the datasets for ourselves.
        variables=["surface_air_temperature"],
        variable_mapping=mlwp_model_variable_mapping,
    )
    target = ...
    evals = ...
    ewb = evaluate.ExtremeWeatherBench(
        case_metadata=cases.load_ewb_events_yaml_into_case_collection(),
        evaluation_objects=evals,
    )

[aaTman]

@aaTman here's a hacky version of what my script looked like:

from extremeweatherbench.inputs import XarrayForecast
import icechunk
import xarray as xr
from extremeweatherbench import cases, evaluate

mlwp_model_variable_mapping = {
    "2m_temperature": "surface_air_temperature",
    "2m_dewpoint_temperature": "surface_dewpoint_temperature",
    "2m_relative_humidity": "surface_relative_humidity",
    "2m_wind_speed": "surface_wind_speed",
    "2m_wind_from_direction": "surface_wind_from_direction",
    "2m_wind_gust": "surface_wind_gust",
    "2m_wind_gust_direction": "surface_wind_gust_direction",
    "2m_wind_gust_speed": "surface_wind_gust_speed",
}

def open_icechunk_dataset(
    bucket: str = DEFAULT_ICECHUNK_BUCKET,
    prefix: str = DEFAULT_ICECHUNK_PREFIX,
    variable_mapping: dict[str, str] | None = None,
    chunks: str | dict | None = "auto",
    source_credentials_prefix: str = DEFAULT_SOURCE_CREDENTIALS_PREFIX,
) -> xr.Dataset:
    """Open a dataset from an Icechunk repository with preprocessing.

    The repository config already knows where the virtual chunks are located.
    We just need to provide credentials broad enough to cover that location.

    Args:
        bucket: GCS bucket containing the Icechunk repository.
        prefix: Prefix within the bucket for the repository.
        variable_mapping: Dictionary mapping source variable names to target names.
        chunks: Chunk specification for xarray (default: "auto").
        source_credentials_prefix: GCS prefix for virtual chunk credentials.
            Should be broad enough to cover wherever the source data lives.

    Returns:
        Preprocessed xarray Dataset ready for evaluation.
    """
    logger.info(f"Opening Icechunk repository at gs://{bucket}/{prefix}")

    # Set up storage
    storage = icechunk.gcs_storage(bucket=bucket, prefix=prefix)

    # Set up credentials for virtual chunks.
    # The repo config knows the exact location; we just provide credentials
    # broad enough to cover it.
    gcs_credentials = icechunk.gcs_from_env_credentials()
    virtual_credentials = icechunk.containers_credentials({source_credentials_prefix: gcs_credentials})

    # Open repository
    repo = icechunk.Repository.open(storage, authorize_virtual_chunk_access=virtual_credentials)
    session = repo.readonly_session("main")

    # Open dataset
    ds = xr.open_dataset(session.store, engine="zarr", chunks=chunks)

    # Apply variable renaming if specified
    if variable_mapping:
        rename_dict = {k: v for k, v in variable_mapping.items() if k in ds.data_vars}
        if rename_dict:
            ds = ds.rename(rename_dict)

    return ds


if __name__ == "__main__":

    ds = open_icechunk_dataset(
        bucket=icechunk_bucket,
        prefix=icechunk_prefix,
        variable_mapping=mlwp_model_variable_mapping,
        chunks="auto",
        source_credentials_prefix=source_prefix,
    )


    forecast = XarrayForecast(
        ds,
        name=f"{icechunk_bucket}/{icechunk_prefix}",
        # NOTE: we have to pass in the variables that will actually be used for this
        # metrics calculations. We can at least bypass the variable mapping by manually
        # processing the datasets for ourselves.
        variables=["surface_air_temperature"],
        variable_mapping=mlwp_model_variable_mapping,
    )
    target = ...
    evals = ...
    ewb = evaluate.ExtremeWeatherBench(
        case_metadata=cases.load_ewb_events_yaml_into_case_collection(),
        evaluation_objects=evals,
    )

Awesome. I'd love to have some generic version of this in the documentation.