mxalign 0.1.0__py3-none-any.whl

mxalign/loaders/anemoi_inference.py

@@ -0,0 +1,103 @@
+from pathlib import Path
+import xarray as xr
+
+from .registry import register_loader
+from ..properties.properties import Space, Time, Uncertainty
+from .base import BaseLoader
+
+DEFAULTS_NETCDF = {"chunks": "auto", "engine": "h5netcdf", "parallel": True}
+
+DEFAULTS_ZARR = {
+    "chunks": "auto",
+    "storage_options": {"anon": True},
+}
+
+
+@register_loader
+class AnemoiInferenceLoader(BaseLoader):
+    name = "anemoi-inference"
+
+    space = Space.GRID
+    time = Time.FORECAST
+    uncertainty = Uncertainty.DETERMINISTIC
+
+    def _load(self):
+
+        kwargs = self.kwargs.copy()
+
+        if isinstance(self.files, str):
+            if Path(self.files).suffix.lower() == ".zarr":
+                files = self.files
+
+                for k, v in DEFAULTS_ZARR.items():
+                    kwargs[k] = self.kwargs.get(k, v)
+
+                loader = _open_zarr
+            else:
+                files = [self.files]
+
+                for k, v in DEFAULTS_NETCDF.items():
+                    kwargs[k] = self.kwargs.get(k, v)
+
+                loader = _open_mf_dataset
+        else:
+            files = self.files
+            if Path(files[0]).suffix.lower() == ".zarr":
+                for k, v in DEFAULTS_ZARR.items():
+                    kwargs[k] = self.kwargs.get(k, v)
+                kwargs["engine"] = "zarr"
+
+            else:
+                for k, v in DEFAULTS_NETCDF.items():
+                    kwargs[k] = self.kwargs.get(k, v)
+
+            loader = _open_mf_dataset
+
+        ds = loader(files, **kwargs)
+        return ds
+
+
+def _open_mf_dataset(files, **kwargs):
+
+    times = xr.open_dataset(files[0], engine=kwargs["engine"], chunks=kwargs["chunks"])[
+        "time"
+    ].values
+    lead_times = times - times[0]
+
+    ds = xr.open_mfdataset(files, preprocess=_preprocess, **kwargs)
+
+    ds_out = (
+        ds.assign_coords({"lead_time": ("time", lead_times)})
+        .rename_dims({"values": "grid_index"})
+        .swap_dims({"time": "lead_time"})
+    )
+
+    return ds_out
+
+
+def _open_zarr(files, **kwargs):
+
+    ds = xr.open_zarr(files, **kwargs)
+    times = ds["time"].values
+    lead_times = times - times[0]
+
+    ds_out = _preprocess(ds)
+
+    ds_out = (
+        ds_out.assign_coords({"lead_time": ("time", lead_times)})
+        .rename_dims({"values": "grid_index"})
+        .swap_dims({"time": "lead_time"})
+    )
+
+    return ds_out
+
+
+def _preprocess(ds):
+    ds_out = (
+        ds.set_coords(["longitude", "latitude"])
+        .expand_dims("reference_time")
+        .assign_coords({"reference_time": ("reference_time", [ds["time"].values[0]])})
+        .drop_vars("time")
+    )
+
+    return ds_out

mxalign/loaders/base.py

@@ -0,0 +1,103 @@
+from abc import ABC, abstractmethod
+
+from .registry import register_loader
+from ..properties.properties import Properties, Space, Time, Uncertainty
+from ..properties.validation import validate_dataset
+from ..properties.utils import properties_to_attrs
+
+
+class BaseLoader(ABC):
+    """Base class for all loaders."""
+
+    name: str = "base"
+
+    space: Space | None = None
+    time: Time | None = None
+    uncertainty: Uncertainty | None = None
+
+    def __init__(self, files, variables=None, grid_mapping=None, **kwargs):
+        self.files = files
+        self.variables = [variables] if isinstance(variables, str) else variables
+        self.grid_mapping = grid_mapping
+        self.kwargs = kwargs
+
+    def load(self):
+        ds = self._load()
+        if self.variables:
+            ds = self._select_variables(ds)
+
+        properties = self._get_properties(ds)
+        validate_dataset(ds, properties)
+
+        ds.attrs["properties"] = properties_to_attrs(properties)
+
+        if self.grid_mapping:
+            ds = self._add_grid_mapping(ds)
+
+        # Make sure all the coordinates are loaded
+        for coord in ds.coords:
+            ds[coord] = ds[coord].compute()
+
+        return ds
+
+    @abstractmethod
+    def _load(self): ...
+
+    def _select_variables(self, ds):
+        return ds[self.variables]
+
+    def _add_grid_mapping(self, ds):
+        ds = ds.space.add_crs(self.grid_mapping)
+        ds = ds.space.add_grid_mapping(self.grid_mapping)
+        return ds
+
+    def _get_properties(self, ds):
+        properties = Properties(
+            space=self.space, time=self.time, uncertainty=self.uncertainty
+        )
+        return properties
+
+
+@register_loader
+class MxAlignLoader(BaseLoader):
+    name = "mxalign"
+
+    space = None
+    time = None
+    uncertainty = None
+
+    def _load(self):
+        import xarray as xr
+
+        files = [self.files] if isinstance(self.files, str) else self.files
+
+        ds = xr.open_mfdataset(files, chunks="auto", **self.kwargs)
+        if "code" in ds.dims:
+            ds = ds.rename_dims({"code": "point_index"}).transpose(
+                "valid_time", "point_index"
+            )
+        return ds
+
+    def _get_properties(self, ds):
+        if "reference_time" in ds.dims and "lead_time" in ds.dims:
+            time = Time.FORECAST
+        elif "valid_time" in ds.dims:
+            time = Time.OBSERVATION
+        else:
+            raise ValueError("Unknown temporal dimensions")
+
+        if "grid_index" in ds.dims or "xc" in ds.dims or "latitude" in ds.dims:
+            space = Space.GRID
+        elif "point_index" in ds.dims:
+            space = Space.POINT
+        else:
+            raise ValueError("Unknown spatial dimensions")
+
+        if "member" in ds.dims:
+            uncertainty = Uncertainty.ENSEMBLE
+        elif "quantile" in ds.dims:
+            uncertainty = Uncertainty.QUANTILE
+        else:
+            uncertainty = Uncertainty.DETERMINISTIC
+
+        return Properties(space=space, time=time, uncertainty=uncertainty)

mxalign/loaders/harp_obstable.py

@@ -0,0 +1,81 @@
+import sqlite3
+import pandas as pd
+
+from .registry import register_loader
+from ..properties.properties import Space, Time, Uncertainty
+from .base import BaseLoader
+
+COORDS = {
+    "longitude": "lon",
+    "latitude": "lat",
+    "valid_time": "validdate",
+    "code": "SID",
+    "altitude": "elev",
+}
+
+
+@register_loader
+class ObstableLoader(BaseLoader):
+    name = "harp-obstable"
+
+    space = Space.POINT
+    time = Time.OBSERVATION
+    uncertainty = Uncertainty.DETERMINISTIC
+
+    def _load(self):
+        if isinstance(self.files, list) and len(self.files > 1):
+            raise NotImplementedError(
+                "Reading from multiple SQLite-files not implemented"
+            )
+
+        conn = sqlite3.connect(self.files)
+
+        if self.variables is None:
+            # Retrieve all variables
+            variables = [
+                var
+                for var in pd.read_sql_query(
+                    "SELECT * FROM SYNOP LIMIT 0", conn
+                ).columns
+                if var not in COORDS.values()
+            ]
+            print(variables)
+        else:
+            variables = self.variables
+
+        # Read the SIDs
+        codes = pd.read_sql(
+            "SELECT SID as code, MIN(lat) AS latitude, MIN(lon) AS longitude, elev as altitude FROM SYNOP GROUP BY SID",
+            conn,
+            index_col="code",
+        ).to_xarray()
+
+        print(codes)
+        # Read the data
+        query = f"""
+                SELECT SID as code, validdate as valid_time, {", ".join(variables)}
+                FROM SYNOP
+            """
+        print(query)
+        df = pd.read_sql(
+            query,
+            conn,
+            index_col=["code", "valid_time"],
+            parse_dates={"valid_time": {"unit": "s"}},
+        )
+        print(df)
+
+        ds = df.to_xarray()
+        lon_values = codes["longitude"].sel(code=ds["code"]).values
+        lat_values = codes["latitude"].sel(code=ds["code"]).values
+        alt_values = codes["altitude"].sel(code=ds["code"]).values
+
+        ds = ds.assign_coords(
+            longitude=("code", lon_values),
+            latitude=("code", lat_values),
+            altitude=("code", alt_values),
+        )
+
+        return ds.rename_dims({"code": "point_index"}).transpose(
+            "valid_time", "point_index"
+        )

mxalign/loaders/loader.py

@@ -0,0 +1,8 @@
+from .registry import get_loader
+
+
+def load(name, files, variables=None, grid_mapping=None, **kwargs):
+    loader_cls = get_loader(name)
+    loader = loader_cls(files, variables, grid_mapping, **kwargs)
+
+    return loader.load()

mxalign/loaders/registry.py

@@ -0,0 +1,17 @@
+_LOADERS = {}
+
+
+def register_loader(cls):
+    _LOADERS[cls.name] = cls
+    return cls
+
+
+def available_loaders():
+    return list(_LOADERS.keys())
+
+
+def get_loader(name):
+    try:
+        return _LOADERS[name]
+    except KeyError:
+        raise ValueError(f"Unknown loader: {name}")

mxalign/properties/properties.py

@@ -0,0 +1,25 @@
+from enum import Enum
+from dataclasses import dataclass
+
+
+class Space(str, Enum):
+    GRID = "grid"
+    POINT = "point"
+
+
+class Time(str, Enum):
+    FORECAST = "forecast"
+    OBSERVATION = "observation"
+
+
+class Uncertainty(str, Enum):
+    DETERMINISTIC = "deterministic"
+    ENSEMBLE = "ensemble"
+    QUANTILE = "quantile"
+
+
+@dataclass(frozen=True)
+class Properties:
+    space: Space
+    time: Time
+    uncertainty: Uncertainty = Uncertainty.DETERMINISTIC

mxalign/properties/specs.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass, field
+from typing import Callable
+from .properties import Space, Time, Uncertainty
+
+
+@dataclass
+class PropertySpec:
+    dim_variants: list[set[str]] = field(default_factory=list)
+    required_coords: set[str] = field(default_factory=set)
+    optional_dims: set[str] = field(default_factory=set)
+    optional_coords: set[str] = field(default_factory=set)
+    validators: list[Callable] = field(default_factory=list)
+
+
+SPACE_SPECS = {
+    Space.GRID: PropertySpec(
+        dim_variants=[
+            {"xc", "yc"},
+            {"grid_index"},
+            {"longitude", "latitude"},
+        ],
+        required_coords={"longitude", "latitude"},
+        optional_coords={"xc", "yc"},
+        optional_dims={"member"},
+    ),
+    Space.POINT: PropertySpec(
+        dim_variants=[
+            {"point_index"},
+        ],
+        required_coords={"longitude", "latitude"},
+        optional_coords={"code", "elevation", "name", "country"},
+    ),
+}
+TIME_SPECS = {
+    Time.FORECAST: PropertySpec(
+        dim_variants=[{"reference_time", "lead_time"}],
+        required_coords={"reference_time", "lead_time"},
+        optional_coords={"valid_time"},
+    ),
+    Time.OBSERVATION: PropertySpec(
+        dim_variants=[{"valid_time"}],
+        required_coords={"valid_time"},
+    ),
+}
+
+UNCERTAINTY_SPECS = {
+    Uncertainty.DETERMINISTIC: PropertySpec(),
+    Uncertainty.ENSEMBLE: PropertySpec(
+        dim_variants=[{"member"}], required_coords={"member"}
+    ),
+    Uncertainty.QUANTILE: PropertySpec(
+        dim_variants=[{"quantile"}], required_coords={"quantile"}
+    ),
+}

mxalign/properties/utils.py

@@ -0,0 +1,43 @@
+from .properties import Properties, Space, Time, Uncertainty
+from .validation import validate_time_dataset, validate_space_dataset
+
+
+def properties_to_attrs(prop: Properties) -> dict:
+    return {
+        "space": prop.space.value,
+        "time": prop.time.value,
+        "uncertainty": prop.uncertainty.value,
+    }
+
+
+def properties_from_attrs(ds) -> Properties:
+    attrs = ds.attrs.get("properties", {})
+    return Properties(
+        space=Space(attrs["space"]),
+        time=Time(attrs["time"]),
+        uncertainty=Uncertainty(attrs.get("uncertainty", Uncertainty.DETERMINISTIC)),
+    )
+
+
+def update_space_property(ds, prop: Space):
+    old_props = properties_from_attrs(ds)
+    new_props = Properties(
+        space=prop,
+        time=old_props.time,
+        uncertainty=old_props.uncertainty,
+    )
+    validate_space_dataset(ds, new_props)
+    ds.attrs["properties"] = properties_to_attrs(new_props)
+    return ds
+
+
+def update_time_property(ds, prop: Time):
+    old_props = properties_from_attrs(ds)
+    new_props = Properties(
+        space=old_props.space,
+        time=prop,
+        uncertainty=old_props.uncertainty,
+    )
+    validate_time_dataset(ds, new_props)
+    ds.attrs["properties"] = properties_to_attrs(new_props)
+    return ds

mxalign/properties/validation.py

@@ -0,0 +1,48 @@
+from .specs import SPACE_SPECS, TIME_SPECS, UNCERTAINTY_SPECS
+
+
+def _validate_dims(ds, variants):
+    if not variants:
+        return
+
+    ds_dims = set(ds.dims)
+
+    for variant in variants:
+        if variant.issubset(ds_dims):
+            return
+
+    raise ValueError(f"Dataset dims {ds_dims} do not match allowed variants {variants}")
+
+
+def _validate_coords(ds, required_coords, axis):
+    missing = required_coords - set(ds.coords)
+    if missing:
+        raise ValueError(f"{axis}: missing required coordinates {missing}")
+
+
+# TIME
+def validate_time_dataset(ds, properties):
+    time_spec = TIME_SPECS[properties.time.value]
+    _validate_dims(ds, time_spec.dim_variants)
+    _validate_coords(ds, time_spec.required_coords, "time")
+
+
+# SPACE
+def validate_space_dataset(ds, properties):
+    space_spec = SPACE_SPECS[properties.space.value]
+    _validate_dims(ds, space_spec.dim_variants)
+    _validate_coords(ds, space_spec.required_coords, "space")
+    validate_time_dataset(ds, properties)
+
+
+# UNCERTAINTY
+def validate_uncertainty_dataset(ds, properties):
+    uncertainty_spec = UNCERTAINTY_SPECS[properties.uncertainty.value]
+    _validate_dims(ds, uncertainty_spec.dim_variants)
+    _validate_coords(ds, uncertainty_spec.required_coords, "uncertainty")
+
+
+def validate_dataset(ds, properties):
+    validate_time_dataset(ds, properties)
+    validate_space_dataset(ds, properties)
+    validate_uncertainty_dataset(ds, properties)

mxalign/runner.py

@@ -0,0 +1,167 @@
+import os
+import xarray as xr
+
+from .utils.config import Config
+from .loaders.loader import load
+from .transformations.transform import transform
+from .align.time import align_time
+from .align.space import align_space
+from .align.nans import broadcast_nans
+from .utils.save import save_dataset, save_metrics
+from .verification import Metric
+
+
+class Runner:
+    def __init__(self, config: str | dict):
+        self.config = Config(config)
+        self.datasets = {}
+
+    def run(self):
+        # 1. Load the datasets
+        self.load_datasets()
+
+        # 2. Transform the datasets
+        self.transform_datasets()
+        self.align()
+        self.verify()
+
+    def load_datasets(self):
+        config = self.config["datasets"]
+        if config is None:
+            return ValueError("No datasets section in the config.")
+        for name, config_ds in config.items():
+            config_ds = config_ds.copy()
+            # Check if all the files exist
+            loader = config_ds.pop("loader")
+            variables = config_ds.pop("variables", None)
+            grid_mapping = config_ds.pop("grid_mapping", None)
+            files = []
+            # Check if all the files exist
+            for file in config_ds.pop("files"):
+                if os.path.exists(file):
+                    files.append(file)
+                else:
+                    print(f"File: {file} is missing, skipping.")
+            self.datasets[name] = load(
+                name=loader,
+                files=files,
+                variables=variables,
+                grid_mapping=grid_mapping,
+                **config_ds,
+            )
+
+    def transform_datasets(self):
+        config = self.config["transformations"]
+        if config is None:
+            pass
+        for transformation, config_trans in config.items():
+            config_trans = config_trans.copy()
+            # if no datasets specified, apply to all datasets
+            names_ds = config_trans.pop("datasets", self.datasets.keys())
+            for name in names_ds:
+                ds = self.datasets[name]
+                self.datasets[name] = transform(
+                    name=transformation, datasets=ds, **config_trans
+                )
+
+    def align(self):
+        config = self.config["alignment"]
+        reference = config.pop("reference")
+        brdcst_nans = config.pop("broadcast_nans", True)
+        config_align_time = config.get("time", None)
+        config_align_space = config.get("space", None)
+        config_align_save = config.get("save", None)
+
+        # align in time
+        if config_align_time:
+            self.align_time(config_align_time)
+        else:
+            print("Skipping temporal alignment")
+
+        # align in space
+        if config_align_space:
+            self.align_space(reference=reference, config=config_align_space)
+        else:
+            print("Skipping spatial alignment")
+
+        # broadcast NaNs
+        if brdcst_nans:
+            self.datasets = broadcast_nans(self.datasets)
+
+        # Save aligned datasets
+        if config_align_save:
+            config = config_align_save.copy()
+            method = config.pop("method")
+            datasets = config.pop("datasets", "all")
+            if datasets == "all":
+                for name, ds in self.datasets.items():
+                    save_dataset(method, name, ds, **config)
+            elif datasets == "merge":
+                ds = xr.concat(
+                    self.datasets.values(),
+                    dim=xr.Variable("model", list(self.datasets.keys())),
+                )
+                save_dataset(method, name, ds, **config)
+            else:
+                raise ValueError("Unknown option for dataset saving.")
+
+    def verify(self):
+        config = self.config["verification"]
+        reference = self.datasets[config["reference"]]
+        config_metrics = config.get("metrics", None)
+        config_save_metrics = config.get("save", None)
+
+        common_vars = set(reference.data_vars)
+        for ds in self.datasets.values():
+            common_vars.intersection_update(set(ds.data_vars))
+        common_vars = list(common_vars)
+
+        if config_metrics:
+            metrics = {}
+            for metric_name, config_metric in config["metrics"].items():
+                config_metric = config_metric.copy()
+                func_path = config_metric.pop("function")
+                inputs = config_metric.pop("inputs")
+
+                metric = Metric(
+                    name=metric_name,
+                    func_path=func_path,
+                    ds_ref=reference[common_vars],
+                    inputs=inputs,
+                    **config_metric,
+                )
+                models = {}
+                for ds_name, ds in self.datasets.items():
+                    if ds_name != config["reference"]:
+                        models[ds_name] = metric.compute(ds[common_vars])
+                models = xr.concat(
+                    models.values(), dim=xr.Variable("model", list(models.keys()))
+                )
+                metrics[metric.name] = models
+            metrics = xr.concat(
+                metrics.values(), dim=xr.Variable("metric", list(metrics.keys()))
+            )
+            self.metrics = metrics.transpose("model", "metric", ...).compute()
+
+        if config_save_metrics:
+            config = config_save_metrics.copy()
+            method = config.pop("method")
+            save_metrics(method, self.metrics, **config)
+
+    def align_time(self, config):
+        self.datasets = align_time(self.datasets, **config)
+
+    def align_space(self, reference, config):
+        ds_ref = self.datasets[reference]
+        for name, ds in self.datasets.items():
+            if name != reference:
+                options = config.get(get_spatial_alignment(ds, ds_ref), {})
+                self.datasets[name] = align_space(ds, ds_ref, **options)
+
+
+def get_spatial_alignment(ds, reference):
+    if reference.space.is_point() and ds.space.is_grid():
+        return "interpolation"
+    if reference.space.is_grid() and ds.space.is_grid():
+        return "regrid"
+    return "null"

mxalign/transformations/__init__.py

@@ -0,0 +1,7 @@
+from . import base
+from . import external
+
+__all__ = [
+    "base",
+    "external",
+]