mxalign 0.1.0__py3-none-any.whl

mxalign/__init__.py

@@ -0,0 +1,36 @@
+from .properties.properties import Properties, Time, Space, Uncertainty
+from .loaders.loader import load
+from .loaders.registry import available_loaders, register_loader
+from .transformations.transform import transform
+from .transformations.registry import available_transformations, register_transformation
+from .interpolations.interpolate import interpolate
+from .interpolations.registry import available_interpolations, register_interpolator
+from .align.time import align_time
+from .align.space import align_space
+
+from . import accessors
+from . import loaders
+from . import transformations
+from . import interpolations
+
+__all__ = [
+    "Properties",
+    "Time",
+    "Space",
+    "Uncertainty",
+    "load",
+    "available_loaders",
+    "register_loader",
+    "transform",
+    "available_transformations",
+    "register_transformation",
+    "interpolate",
+    "available_interpolations",
+    "register_interpolator",
+    "align_time",
+    "align_space",
+    "accessors",
+    "loaders",
+    "transformations",
+    "interpolations",
+]

mxalign/accessors/__init__.py

@@ -0,0 +1,7 @@
+from . import space
+from . import time
+
+__all__ = [
+    "space",
+    "time",
+]

mxalign/accessors/space.py

@@ -0,0 +1,205 @@
+import xarray as xr
+import cartopy.crs as ccrs
+import numpy as np
+
+from ..properties.properties import Space
+from ..properties.utils import properties_from_attrs
+
+from ..utils.projections import create_cartopy_crs, BUILTIN
+
+# Tolerance in degrees that the coordinates of two grids can differ while still being interpreted as the same grid.
+# 0.0001 degrees ~ 10m at 45 deg latitude
+COORD_TOLERANCE = 0.0001
+
+
+@xr.register_dataset_accessor("space")
+class SpaceAccessor:
+    def __init__(self, ds):
+        self._space = properties_from_attrs(ds).space
+        self._ds = ds
+
+    def is_grid(self):
+        return self._space == Space.GRID
+
+    def is_point(self):
+        return self._space == Space.POINT
+
+    def add_crs(self, crs):
+        if self.is_point():
+            raise ValueError("Cannot add CRS to a point dataset")
+        if isinstance(crs, str):
+            try:
+                crs = BUILTIN[crs.lower()]
+            except KeyError:
+                raise ValueError("crs: {crs} not found in supported projections")
+        if isinstance(crs, dict):
+            crs = create_cartopy_crs(
+                projection=crs["projection"],
+                kws_projection=crs["kws_projection"],
+                kws_globe=crs.get("kws_globe", None),
+            )
+        return self._ds.assign_attrs({"crs": crs})
+
+    def add_grid_mapping(self, grid_mapping: str | dict):
+        if self.is_point():
+            raise ValueError("Cannot add grid mapping to a point dataset")
+        if isinstance(grid_mapping, str):
+            try:
+                grid_mapping = BUILTIN[grid_mapping.lower()]["kws_grid"]
+            except KeyError:
+                raise ValueError(
+                    "grid mapping: {grid_mapping} not found in supported mappings"
+                )
+        return self._ds.assign_attrs({"grid_mapping": grid_mapping})
+
+    def add_xy(self, crs=None):
+        if crs is not None:
+            self._ds = self.add_crs(crs)
+
+        crs = self._ds.attrs.get("crs", None)
+
+        if crs is None:
+            raise ValueError("No CRS provided and no CRS found in dataset attributes")
+
+        if {"longitude", "latitude"}.issubset(self._ds.dims):
+            raise ValueError(
+                "Cannot add x/y coordinates to a GRID dataset that has longitude/latitude dimensions"
+            )
+        elif {"xc", "yc"}.issubset(self._ds.coords):
+            return self._ds
+        else:
+            xyz = crs.transform_points(
+                x=self._ds["longitude"].values,
+                y=self._ds["latitude"].values,
+                src_crs=ccrs.PlateCarree(),
+            )
+
+        if self.is_grid():
+            ds_out = self._ds.assign_coords(
+                xc=("grid_index", xyz[:, 0]), yc=("grid_index", xyz[:, 1])
+            )
+        elif self.is_point():
+            ds_out = self._ds.assign_coords(
+                xc=("point_index", xyz[:, 0]), yc=("point_index", xyz[:, 1])
+            )
+        else:
+            raise ValueError("Dataset does not have expected spatial properties")
+
+        return ds_out
+
+    def is_stacked(self):
+        if {"xc", "yc"}.issubset(self._ds.dims) or {"longitude", "latitude"}.issubset(
+            self._ds.dims
+        ):
+            return False
+        elif "grid_index" in self._ds.dims:
+            return True
+        else:
+            raise ValueError("Dataset does not have expected dimensions for GRID")
+
+    def stack(self):
+        if self.is_point():
+            raise ValueError("POINT datasets cannot be stacked")
+        if self.is_stacked():
+            return self._ds
+        else:
+            if {"xc", "yc"}.issubset(self._ds.dims):
+                dims_to_stack = ["yc", "xc"]
+            elif {"lat", "lon"}.issubset(self._ds.dims):
+                dims_to_stack = ["lat", "lon"]
+            else:
+                raise ValueError("Could not find correct dimensions to stack")
+        return self._ds.stack({"grid_index": dims_to_stack}).reset_index("grid_index")
+
+    def unstack(self, crs=None, **kwargs):
+        if self.is_point():
+            raise ValueError("POINT datasets cannot be unstacked")
+        if not self.is_stacked():
+            return self._ds
+        else:
+            if crs:
+                self.add_crs(crs)
+            kws_mindex = dict.fromkeys(["nx", "ny", "lon_ll", "lat_ll", "dx", "dy"])
+            for key in kws_mindex.keys():
+                value = kwargs.get(key, None)
+                if value is None:
+                    try:
+                        value = self._ds.attrs["grid_mapping"][key]
+                    except KeyError:
+                        raise KeyError(
+                            f"Did not find a value for {key} in the dataset attributes, please provide it as an argument"
+                        )
+                kws_mindex[key] = value
+
+            mindex = self._create_multiindex(**kws_mindex)
+            mcoords = xr.Coordinates.from_pandas_multiindex(mindex, "grid_index")
+            ds_mindex = self._ds.assign_coords(mcoords)
+            ds_mindex.attrs["grid_mapping"] = kws_mindex
+            return ds_mindex.unstack()
+
+    def _create_multiindex(self, nx, ny, lon_ll, lat_ll, dx, dy, **kwargs):
+        from pandas import MultiIndex
+
+        if self._ds.sizes["grid_index"] != nx * ny:
+            raise ValueError(
+                f"Size of grid_index ({self._ds.sizes['grid_index']}) does not match product of nx and ny ({nx * ny})"
+            )
+
+        crs = self._ds.attrs["crs"]
+        x_ll, y_ll = crs.transform_point(x=lon_ll, y=lat_ll, src_crs=ccrs.PlateCarree())
+
+        xc = x_ll + np.arange(nx) * dx
+        yc = y_ll + np.arange(ny) * dy
+
+        mindex = MultiIndex.from_product([yc, xc], names=["yc", "xc"])
+
+        return mindex
+
+    def align_with(self, ds, **kwargs):
+        if self.is_grid():
+            if ds.space.is_grid():
+                return _align_grid_grid(self._ds, ds, **kwargs)
+            elif ds.space.is_point():
+                return _align_grid_point(self._ds, ds, **kwargs)
+        elif self.is_point():
+            if ds.space.is_point():
+                return _align_point_point(self._ds, ds, **kwargs)
+            elif ds.space.is_grid():
+                return _align_point_grid(self._ds, ds, **kwargs)
+        else:
+            raise ValueError("Datasets do not have compatible spatial properties")
+
+
+def _align_grid_grid(ds1, ds2, **kwargs):
+    if np.array_equal(
+        ds1["longitude"].values, ds2["longitude"].values
+    ) and np.array_equal(ds1["latitude"].values, ds2["latitude"].values):
+        return ds1, ds2
+    elif np.allclose(
+        ds1["longitude"].values, ds2["longitude"].values, atol=COORD_TOLERANCE
+    ) and np.allclose(
+        ds1["latitude"].values, ds2["latitude"].values, atol=COORD_TOLERANCE
+    ):
+        print(
+            f"Some lat-lon coordinates differ. But the difference is smaller than {COORD_TOLERANCE} degrees, considering both grids as equal"
+        )
+        return ds1, ds2
+    else:
+        raise NotImplementedError("Regridding not implemented")
+
+
+def _align_grid_point(ds1, ds2, **kwargs):
+    from ..interpolations.interpolate import interpolate
+
+    method = kwargs.pop("method", "xarray")
+    ds1 = interpolate(ds1, ds2, method, **kwargs)
+
+    return ds1, ds2
+
+
+def _align_point_point(ds1, ds2, **kwargs):
+    raise NotImplementedError("Point selection not implemented")
+
+
+def _align_point_grid(ds1, ds2, **kwargs):
+    raise NotImplementedError("Gridding of Point datanot implemented")

mxalign/accessors/time.py

@@ -0,0 +1,180 @@
+import xarray as xr
+import numpy as np
+
+from ..properties.properties import Time
+from ..properties.utils import properties_from_attrs, update_time_property
+
+
+@xr.register_dataset_accessor("time")
+class TimeAccessor:
+    def __init__(self, ds):
+        self._time = properties_from_attrs(ds).time
+        self._ds = ds
+
+    def is_forecast(self):
+        return self._time == Time.FORECAST
+
+    def is_observation(self):
+        return self._time == Time.OBSERVATION
+
+    def add_valid_time(self):
+        if self.is_forecast():
+            valid_time = (
+                self._ds["reference_time"].values[:, np.newaxis]
+                + self._ds["lead_time"].values
+            )
+            ds_out = self._ds.assign_coords(
+                {"valid_time": (["reference_time", "lead_time"], valid_time)}
+            )
+        else:
+            ds_out = self._ds
+        return ds_out
+
+    def align_with(self, ds, **kwargs):
+        if self.is_forecast():
+            if ds.time.is_forecast():
+                return _align_forecast_forecast(self._ds, ds, **kwargs)
+            elif ds.time.is_observation():
+                return _align_forecast_observation(self._ds, ds, **kwargs)
+        elif self.is_observation():
+            if ds.time.is_observation():
+                return _align_observation_observation(self._ds, ds, **kwargs)
+            elif ds.time.is_forecast():
+                return _align_observation_forecast(self._ds, ds, **kwargs)
+        else:
+            raise ValueError("Datasets do not have compatible temporal properties")
+
+
+def _align_forecast_forecast(ds1, ds2, only_common=False):
+    # Align the reference times
+    common_reference_times = ds1.indexes["reference_time"].intersection(
+        ds2.indexes["reference_time"]
+    )
+    ds1_aligned = ds1.sel(reference_time=common_reference_times)
+    ds2_aligned = ds2.sel(reference_time=common_reference_times)
+
+    # Align the lead times
+    if only_common:
+        common_lead_times = ds1_aligned.indexes["lead_time"].intersection(
+            ds2_aligned.indexes["lead_time"]
+        )
+        ds1_aligned = ds1_aligned.sel(lead_time=common_lead_times)
+        ds2_aligned = ds2_aligned.sel(lead_time=common_lead_times)
+    else:
+        non_aligning_dims = (set(ds1.dims) | set(ds2.dims)) - set(["lead_time"])
+        ds1_aligned, ds2_aligned = xr.align(
+            ds1_aligned, ds2_aligned, join="outer", exclude=non_aligning_dims
+        )
+        ds1_aligned = ds1_aligned.time.add_valid_time()
+        ds2_aligned = ds2_aligned.time.add_valid_time()
+    return ds1_aligned, ds2_aligned
+
+
+def _align_forecast_observation(
+    ds_forecast, ds_observation, only_common=False, lead_time="start-min"
+):
+    ds_forecast = ds_forecast.time.add_valid_time()
+
+    # Check if reference_times are continuous
+    reference_time_diff = ds_forecast.reference_time.diff("reference_time").values
+    if not (reference_time_diff[0] == reference_time_diff).all():
+        raise NotImplementedError(
+            "Aligning a forecast with non-continuous reference times with an observation is not implemented."
+        )
+    if lead_time == "start-min":
+        min_diff = reference_time_diff[0]
+        ds_forecast_reduced = ds_forecast.where(
+            ds_forecast.lead_time < min_diff, drop=True
+        )
+    elif lead_time == "start-max":
+        max_diff = ds_forecast.lead_time.max().values
+        reference_times = np.arange(
+            ds_forecast.reference_time.min().values,
+            ds_forecast.reference_time.max().values,
+            max_diff,
+            dtype="datetime64[ns]",
+        )
+        ds_forecast_reduced = ds_forecast.sel(reference_time=reference_times)
+    else:
+        raise ValueError(
+            "Invalid value for lead_time. Expected 'start-min' or 'start-max'."
+        )
+
+    ds_forecast_stacked = (
+        ds_forecast_reduced.stack(time=["reference_time", "lead_time"])
+        .reset_index("time")
+        .swap_dims({"time": "valid_time"})
+        .transpose("valid_time", ...)
+    )
+    if only_common:
+        ds_forecast_aligned, ds_observation_aligned = xr.align(
+            ds_forecast_stacked,
+            ds_observation,
+            join="inner",
+            exclude=set(ds_forecast_stacked.coords)
+            | set(ds_observation.coords) - set(["valid_time"]),
+        )
+    else:
+        ds_forecast_aligned, ds_observation_aligned = xr.align(
+            ds_forecast_stacked,
+            ds_observation,
+            join="outer",
+            exclude=set(ds_forecast_stacked.coords)
+            | set(ds_observation.coords) - set(["valid_time"]),
+        )
+    ds_forecast_aligned = update_time_property(ds_forecast_aligned, Time.OBSERVATION)
+    return ds_forecast_aligned, ds_observation_aligned
+
+
+def _align_observation_observation(ds1, ds2, only_common=False):
+    exclude = (set(ds1.dims) | set(ds2.dims)) - set(["valid_time"])
+    if only_common:
+        ds1_aligned, ds2_aligned = xr.align(ds1, ds2, join="inner", exclude=exclude)
+    else:
+        ds1_aligned, ds2_aligned = xr.align(ds1, ds2, join="outer", exclude=exclude)
+    return ds1_aligned, ds2_aligned
+
+
+def _align_observation_forecast(ds_observation, ds_forecast, only_common=False):
+    ds_forecast_cut = ds_forecast.time.add_valid_time()
+    if (
+        ds_forecast_cut.reference_time.min().values
+        < ds_observation.valid_time.min().values
+    ):
+        ds_forecast_cut = ds_forecast_cut.sel(
+            reference_time=slice(ds_observation.valid_time.min().values, None)
+        )
+    if ds_forecast_cut.valid_time.max().values > ds_observation.valid_time.max().values:
+        # The forecast time-step/lead times might not always align with the maximum observation time
+        valid_diff = (
+            ds_forecast_cut["valid_time"] - (ds_observation["valid_time"].max())
+        ).isel(lead_time=-1)
+        last_valid_index = (
+            np.abs(valid_diff.where(valid_diff <= 0, drop=True)).argmin().values
+        )
+        max_reference_time = ds_forecast_cut.isel(reference_time=last_valid_index)[
+            "reference_time"
+        ].values
+
+        # max_reference_time = ds_observation.valid_time.max().values - (ds_forecast_cut.lead_time.max().values - shift)
+        ds_forecast_cut = ds_forecast_cut.sel(
+            reference_time=slice(None, max_reference_time)
+        )
+
+    ds_observation_aligned = ds_observation.sel(valid_time=ds_forecast_cut.valid_time)
+    ds_observation_aligned = ds_observation_aligned.transpose(
+        "reference_time", "lead_time", ...
+    )
+    ds_observation_aligned = update_time_property(ds_observation_aligned, Time.FORECAST)
+    if only_common:
+        return ds_observation_aligned, ds_forecast_cut
+    else:
+        ds_observation_aligned, ds_forecast_aligned = xr.align(
+            ds_observation_aligned,
+            ds_forecast.time.add_valid_time(),
+            join="outer",
+            exclude=(set(ds_observation_aligned.coords) | set(ds_forecast_cut.coords))
+            - set(["reference_time", "lead_time"]),
+        )
+        ds_observation_aligned["valid_time"] = ds_forecast_aligned["valid_time"]
+        return ds_observation_aligned, ds_forecast_aligned

mxalign/align/__init__.py

@@ -0,0 +1,7 @@
+from . import time
+from . import space
+
+__all__ = [
+    "time",
+    "space",
+]

mxalign/align/nans.py

@@ -0,0 +1,72 @@
+import xarray as xr
+import itertools
+
+
+def broadcast_nans(datasets: dict | list) -> None:
+    """
+    Broadcasts NaN values across a list of xarray Datasets by ensuring that if a value is NaN
+    in one dataset at a specific coordinate, it becomes NaN in all datasets at that coordinate.
+
+    Parameters
+    ----------
+    datasets : list[xr.Dataset] | dict[str, xr.Dataset]
+        A list of xarray Datasets to process. The datasets should share some common
+        coordinates and variables.
+
+    Returns
+    -------
+    list[xr.Dataset] | dict[str, xr.Dataset]
+
+
+    Notes
+    -----
+    - The function operates on pairs of datasets, comparing each dataset with every other dataset
+      in the list.
+    - Only coordinate values that exist in both datasets of a pair are considered.
+    - Only variables that exist in both datasets of a pair are processed.
+    - The NaN broadcasting is performed at the intersection of coordinates between each pair
+      of datasets.
+
+    Examples
+    --------
+    >>> ds1 = xr.Dataset(...)
+    >>> ds2 = xr.Dataset(...)
+    >>> ds3 = xr.Dataset(...)
+    >>> broadcast_nans([ds1, ds2, ds3])
+    """
+
+    if isinstance(datasets, xr.Dataset):
+        return datasets
+    elif isinstance(datasets, dict):
+        keys = list(datasets.keys())
+        working = [ds.copy(deep=True) for ds in datasets.values()]
+    else:
+        keys = None
+        working = [ds.copy(deep=True) for ds in datasets]
+
+    # Iterate over all pairs of datasets
+    for dsA, dsB in itertools.combinations(working, 2):
+        # Find the shared coordinates for all dimensions
+        common_coords = {
+            dim: sorted(set(dsA[dim].values) & set(dsB[dim].values)) for dim in dsA.dims
+        }
+
+        # Iterate over all variables
+        for var in dsA.data_vars:
+            if var in dsB:  # Ensure both datasets have the variable
+                # Select the data at common coordinates
+                selA = dsA[var].sel(**common_coords)
+                selB = dsB[var].sel(**common_coords)
+
+                # Compute NaN mask for shared coordinates
+                nan_mask = selA.isnull() | selB.isnull()
+
+                # Apply NaN mask back to both datasets
+                dsA[var].loc[common_coords] = (
+                    dsA[var].sel(**common_coords).where(~nan_mask)
+                )
+                dsB[var].loc[common_coords] = (
+                    dsB[var].sel(**common_coords).where(~nan_mask)
+                )
+
+    return dict(zip(keys, working)) if keys else working

mxalign/align/space.py

@@ -0,0 +1,21 @@
+import xarray as xr
+
+
+def align_space(datasets, reference, **kwargs):
+    if isinstance(datasets, (xr.Dataset, xr.DataArray)):
+        datasets = [datasets]
+    if isinstance(datasets, dict):
+        keys = datasets.keys()
+        datasets = datasets.items()
+    else:
+        keys = None
+
+    datasets = [ds.space.align_with(reference, **kwargs)[0] for ds in datasets]
+
+    if keys is None:
+        if len(datasets) == 1:
+            return datasets[0]
+        else:
+            return datasets
+    else:
+        return {key: value for (key, value) in zip(keys, datasets)}

mxalign/align/time.py

@@ -0,0 +1,62 @@
+import xarray as xr
+
+
+def align_time(
+    datasets: list[xr.Dataset] | dict[str, xr.Dataset], return_as: str = "forecast"
+):
+    if isinstance(datasets, (xr.Dataset, xr.DataArray)):
+        datasets = [datasets]
+    if isinstance(datasets, dict):
+        keys = datasets.keys()
+        datasets = datasets.values()
+    else:
+        keys = None
+
+    if return_as != "forecast":
+        NotImplementedError(
+            "Currently only temporal alignment return forecast structure is supported."
+        )
+
+    # Get the first forecast to start building the valid times
+    valid_times_fcst = None
+    valid_times_obs = None
+    first_fcst = True
+    first_obs = True
+    for ds in datasets:
+        if ds.time.is_forecast():
+            if first_fcst:
+                valid_times_fcst = ds.time.add_valid_time()["valid_time"].to_dataset(
+                    name="valid_times"
+                )
+                valid_times_fcst = valid_times_fcst.assign_attrs(ds.attrs)
+                first_fcst = False
+            else:
+                _ds = ds.time.add_valid_time()["valid_time"].to_dataset(
+                    name="valid_times"
+                )
+                _ds = _ds.assign_attrs(ds.attrs)
+                _, valid_times_fcst = _ds.time.align_with(valid_times_fcst)
+        elif ds.time.is_observation():
+            if first_obs:
+                valid_times_obs = ds["valid_time"].to_dataset(name="valid_times")
+                valid_times_obs = valid_times_obs.assign_attrs(ds.attrs)
+                first_obs = False
+            else:
+                _ds = ds["valid_time"].to_dataset(name="valid_times")
+                _ds = _ds.assign_attrs(ds.attrs)
+                _, valid_times_obs = _ds.time.align_with(valid_times_obs)
+
+    if (valid_times_obs is None) and (valid_times_fcst is None):
+        raise ValueError("No observations or forecasts found")
+    elif valid_times_fcst is None:
+        valid_times = valid_times_obs
+    elif valid_times_obs is None:
+        valid_times = valid_times_fcst
+    else:
+        _, valid_times = valid_times_obs.time.align_with(valid_times_fcst)
+
+    datasets = [ds.time.align_with(valid_times)[0] for ds in datasets]
+    if keys is None:
+        return datasets
+    else:
+        return {key: value for (key, value) in zip(keys, datasets)}