diff --git a/docs/source/mask/mask_xarray_migration_plan.md b/docs/source/mask/mask_xarray_migration_plan.md new file mode 100644 index 00000000..56ee8247 --- /dev/null +++ b/docs/source/mask/mask_xarray_migration_plan.md @@ -0,0 +1,180 @@ +# Mask-Without-Cutout Migration Plan + +## Goal + +Replace Cutout-dependent masking with a direct xarray-based workflow: + +- `datasets -> models -> masking -> analysis` + +The new masking flow should work on model output (`xarray.Dataset` / `xarray.DataArray`) directly, while reusing current `mask.py` code as much as possible. + +## What Changes, What Stays + +- Keep: + - `Mask` object for raster/shapefile mask creation and persistence. + - Existing layer operations in `src/geodata/mask.py` (`add_layer`, `filter_layer`, `merge_layer`, `extract_shapes`, `save_mask`, `from_name`). + - Existing geospatial utilities (`ras_to_xarr`, `calc_grid_area` logic from `cutout.py`, coordinate formatting helpers). +- Remove dependency on: + - `Cutout.add_mask(...)` + - `Cutout.add_grid_area(...)` + - `Cutout.mask(...)` +- Add: + - A new xarray-focused masking adapter class/module (proposed below). + +## Proposed Target API + +Create a dedicated class (example name: `XarrayMask`) that only deals with xarray data: + +1. **Creation / loading** + - `XarrayMask.from_mask(mask: Mask, grid: xr.Dataset | xr.DataArray, include_merged=True, include_shapes=True)` + - `XarrayMask.from_name(name: str, grid: xr.Dataset | xr.DataArray, mask_dir=...)` +2. **Area calculation** + - `XarrayMask.compute_grid_area(grid: xr.Dataset | xr.DataArray) -> xr.DataArray` +3. **Applying mask** + - `XarrayMask.attach(dataset, include_area=True) -> dict[str, xr.Dataset]` + - Equivalent to current `Cutout.mask(...)` behavior (mask as extra variables). + - `XarrayMask.apply(dataset, mode="where", include_area=False) -> dict[str, xr.Dataset]` + - New convenience method returning mask-applied outputs: + - `mode="where"`: outside mask -> NaN + - `mode="multiply"`: outside mask -> 0 + +This gives both: +- transparent feature-style behavior (`attach`) +- direct filtered outputs (`apply`) + +## Reuse Map (Do Not Reinvent) + +Directly reuse existing code paths: + +- From `src/geodata/mask.py`: + - `Mask.from_name(...)` + - `Mask.load_merged_xr()` / `Mask.load_shape_xr()` +- From `src/geodata/cutout.py`: + - `ds_reformat_index(...)` (move/shared helper) + - `coarsen(...)` (move/shared helper) + - `calc_grid_area(...)` (move/shared helper) +- Keep the same coordinate conventions: + - normalize to `lat`, `lon` + - align mask grid to target dataset grid before applying + +Refactor suggestion: +- Move shared helpers into a new utility module, e.g. `src/geodata/spatial.py` or `src/geodata/mask_xarray.py`, then import from both old and new flows during transition. + +## Migration Phases + +### Phase 0 - Freeze Current Behavior + +- Add tests that lock existing behavior for: + - coarsening/alignment from mask raster to target grid + - area computation + - output structure currently returned by `Cutout.mask(...)` + +This prevents regressions while extracting logic. + +### Phase 1 - Extract Shared Spatial Helpers + +- Move (or duplicate temporarily) these functions out of `cutout.py`: + - `ds_reformat_index` + - `coarsen` + - `calc_grid_area` +- Add unit tests for each helper independent of `Cutout`. + +### Phase 2 - Introduce `XarrayMask` + +- Implement class that: + - loads saved `Mask` by name + - converts mask rasters to xarray + - coarsens/aligned to target grid + - computes area from target grid + - provides `attach()` and `apply()` + +### Phase 3 - Integrate into datasets -> models workflow + +- At model output point (where xarray result exists), call: + - `xmask = XarrayMask.from_name("my_mask", grid=model_ds)` + - `masked = xmask.apply(model_ds, mode="where")` +- Keep `attach()` available for advanced users needing raw mask + area features. + +### Phase 4 - Deprecate Cutout Masking Surface + +- Mark these as deprecated: + - `Cutout.add_mask` + - `Cutout.add_grid_area` + - `Cutout.mask` +- Keep them as wrappers calling new `XarrayMask` for 1-2 releases. + +### Phase 5 - Remove Cutout Dependency + +- Remove or archive old mask-coupled Cutout paths once internal usage is migrated. +- Keep `Cutout` only if still needed for data preparation. + +## Detailed Behavior Decisions + +To avoid ambiguity, define these explicitly: + +- Mask value semantics: + - `mask > 0` means valid/included + - `mask <= 0` means excluded +- Apply scope: + - apply to all data variables by default + - optional include/exclude variable list +- Output keys: + - `"merged_mask"` for merged mask + - shape names for shape masks (same as current behavior) +- Alignment: + - always reformat coords to `lat`/`lon` + - always transpose to `time, lat, lon` when `time` exists +- Area: + - computed from target grid only (not from mask grid) to stay consistent with model outputs + +## Risks and Mitigations + +- Risk: hidden coordinate mismatches (`x/y` vs `lat/lon`, descending latitude). + - Mitigation: centralize coordinate normalization in one helper and test with both styles. +- Risk: users depending on old `Cutout.mask` output shape. + - Mitigation: make `attach()` output identical structure and keep temporary wrappers. +- Risk: performance hit when repeatedly coarsening same mask. + - Mitigation: cache aligned masks keyed by grid signature (lat/lon hashes + mask name). + +## Suggested Minimal First Milestone (1 PR) + +- Add `src/geodata/mask_xarray.py` with: + - `XarrayMask.from_name(...)` + - `compute_grid_area(...)` + - `attach(...)` + - `apply(...)` (`where` + `multiply`) +- Reuse copied helper logic from `cutout.py` initially (refactor later). +- Add tests: + - parity test with `Cutout.mask(...)` behavior for `attach()` + - correctness test for `apply(...)` + - area calculation sanity test + +## Example Future Usage + +```python +import geodata + +# model output +ds_model = model.run(...) # xr.Dataset with dims time/lat/lon (or x/y) + +# load and align mask to ds_model grid +xmask = geodata.XarrayMask.from_name("china", grid=ds_model) + +# 1) feature-style output (raw + mask + area) +attached = xmask.attach(ds_model, include_area=True) + +# 2) direct masked output +masked = xmask.apply(ds_model, mode="where", include_area=True) +china_masked = masked["merged_mask"] +``` + +## Recommended Naming + +- Keep existing `Mask` name for geospatial mask construction object. +- Use a distinct name for xarray adapter to avoid confusion: + - preferred: `XarrayMask` + - alternatives: `MaskApplier`, `MaskDatasetAdapter` + +This separation keeps responsibilities clear: +- `Mask`: build/store masks +- `XarrayMask`: align/apply masks to model outputs diff --git a/docs/source/mask/xarray_mask_workflow.rst b/docs/source/mask/xarray_mask_workflow.rst new file mode 100644 index 00000000..14062f82 --- /dev/null +++ b/docs/source/mask/xarray_mask_workflow.rst @@ -0,0 +1,81 @@ +Xarray masking workflow +========================= + +This page summarizes the **xarray-first masking** work added alongside the +longer-term plan in :doc:`mask_xarray_migration_plan`. The legacy path based on +``Cutout`` (``add_mask``, ``add_grid_area``, ``mask``) is unchanged for now; the +new pieces let you mask **any** model or analysis output +given as an ``xarray.Dataset`` or ``xarray.DataArray``, without threading mask +logic through model classes. + +What was added +-------------- + +**Phase 0 — behavior freeze (tests only)** + +Offline tests lock in legacy masking behavior so refactors do not silently change +results: + +* Coarsening / alignment of saved mask rasters onto a target grid. +* Grid cell area computation consistent with the cutout-style workflow. +* The structure of outputs from ``Cutout.mask(...)`` (keys, variables, dimensions). +* Selected error paths (missing mask, missing area, invalid mask state). + +**Phase 1 — shared spatial helpers** + +The following helpers now live in ``geodata.mask.spatial`` and are re-used from +``cutout`` (and plotting code where relevant): + +* ``ds_reformat_index`` — normalize coordinates toward ``lat`` / ``lon``. +* ``coarsen`` — align a higher-resolution mask grid to a target grid. +* ``calc_grid_area`` / ``calc_shp_area`` — area utilities used by the masking workflow. + +Public names on ``geodata.cutout`` (e.g. ``coarsen``, ``calc_grid_area``) remain +available as aliases for backward compatibility. + +**Phase 2 — ``XarrayMask``** + +``XarrayMask`` (``from geodata import XarrayMask``) provides: + +* ``from_name`` / ``from_mask`` — load a saved ``Mask`` and align + merged and shape masks to a target ``grid`` (your model output or any dataset + with compatible ``x``/``y`` or ``lat``/``lon`` coordinates). +* ``compute_grid_area`` — per-cell area on the target grid (same idea as cutout + grid area). +* ``attach`` — return a dict of datasets like legacy ``Cutout.mask``: original + variables plus ``mask`` and optional ``area``. +* ``apply`` — return masked data (``mode="where"`` for NaN outside mask, + ``mode="multiply"`` for zero outside mask), optionally with ``area``. + +**Integration pattern (no coupling inside models)** + +Masking is intentionally **not** built into wind, pvlib, or other model ``estimate`` +APIs. The intended usage is: + +1. Run the model and obtain ``output_ds`` (or a ``DataArray`` you wrap in a + one-variable dataset). +2. Build ``XarrayMask.from_name("my_mask", grid=output_ds, mask_dir=...)`` if needed. +3. Call ``attach(output_ds)`` or ``apply(output_ds, ...)`` for analysis. + +See the offline tests under ``tests/pr/`` (e.g. ``test_xarray_mask.py``, +``test_wind_xarraymask_integration.py``) for concrete examples. + +Package layout note +------------------- + +The repository currently has both: + +* ``src/geodata/mask.py`` — original ``geodata.mask`` implementation (``Mask``, + raster helpers, etc.). +* ``src/geodata/mask/`` — package namespace that re-exports that API **and** + hosts new modules (``spatial.py``, ``xarray_mask.py``). + +Imports like ``from geodata import Mask`` and ``from geodata import XarrayMask`` +continue to work during this transition. + +See also +-------- + +* :doc:`mask_xarray_migration_plan` — full migration phases and deprecation plan. +* :doc:`mask_on_cutout` — legacy notebook: masks via ``Cutout``. +* :doc:`mask_creation_workflow` — building and saving ``Mask`` objects from rasters. diff --git a/src/geodata/__init__.py b/src/geodata/__init__.py index a9ee1652..190d1669 100644 --- a/src/geodata/__init__.py +++ b/src/geodata/__init__.py @@ -16,12 +16,17 @@ from ._version import __version__ from .cutout import Cutout from .dataset import Dataset -from .mask import Mask +from typing import cast + +from . import mask as _mask_pkg from .plot import * # noqa: F403 from .model import * # noqa: F403 +Mask = cast(type, getattr(_mask_pkg, "Mask")) +XarrayMask = cast(type, getattr(_mask_pkg, "XarrayMask")) + __author__ = "Michael Davidson (UCSD), William Honaker" __copyright__ = "GNU GPL 3 license" -__all__ = ["Cutout", "Dataset", "Mask", "__version__"] +__all__ = ["Cutout", "Dataset", "Mask", "XarrayMask", "__version__"] diff --git a/src/geodata/cutout.py b/src/geodata/cutout.py index 507cf932..732540e3 100644 --- a/src/geodata/cutout.py +++ b/src/geodata/cutout.py @@ -20,13 +20,10 @@ """ import logging -from functools import partial from pathlib import Path -from typing import Literal, Optional, Union +from typing import Optional, Union import numpy as np -import pyproj -import shapely import xarray as xr from shapely.geometry import box from tqdm.auto import tqdm @@ -45,6 +42,7 @@ ) from .datasets._base import BaseDataset from .mask import Mask +from .mask.spatial import calc_grid_area, calc_shp_area, coarsen, ds_reformat_index from .preparation import ( cutout_get_meta, cutout_get_meta_view, @@ -517,145 +515,4 @@ def _convert_cutout( pv = pv -def ds_reformat_index(ds: xr.DataArray) -> xr.DataArray: - """Format the dataArray generated from the convert function. - - Args: - ds (xr.DataArray): dataArray generated from the convert function. - - Returns: - xr.DataArray: DataArray with lat and lon as dimensions. - """ - - if "lat" in ds.dims and "lon" in ds.dims: - return ds.sortby(["lat", "lon"]) - elif "lat" in ds.coords and "lon" in ds.coords: - return ( - ds.reset_coords(["lon", "lat"], drop=True) - .rename({"x": "lon", "y": "lat"}) - .sortby(["lat", "lon"]) - ) - return ds.rename({"x": "lon", "y": "lat"}).sortby(["lat", "lon"]) - - -def _find_intercept(list1, list2, start, threshold=0): - """Find_intercept is a helper function to find the best start point for doing coarsening - in order to make the coordinates of the coarsen as close to the target as possible. - """ - min_res = 0 - init = 0 - for i in range(len(list1) - start): - resid = ((list1[start + i] - list2[0]) % (list2[1] - list2[0])).values.tolist() - if i == 0: - init = resid - if resid <= threshold: - return i - if resid > min_res: - min_res = resid - else: - min_res = resid - break - if min_res == init: - return 0 - else: - return i - - -def coarsen(ori: xr.Dataset, tar: xr.Dataset, func: Literal["sum", "mean"] = "mean"): - """This function will reindex the original xarray dataset according to the coordiantes of the target. - There might be a bias for lattitudes and longitudes. The bias are normally within 0.01 degrees. - In order to not lose too much data, a threshold for bias in degree could be given. - When threshold = 0, it means that the function is going to find the best place with smallest bias. - - Args: - ori (xr.Dataset): The original xarray dataset. - tar (xr.Dataset): The target xarray dataset. - func (Literal['sum', 'mean']): The function to be used for reduction. Defaults to "mean". - - Returns: - xr.Dataset: The reindexed xarray dataset. - - Raises: - ValueError: reduction method can only be 'mean' or 'sum'. - """ - lat_multiple = round( - ((tar.lat[1] - tar.lat[0]) / (ori.lat[1] - ori.lat[0])).values.tolist() - ) - lon_multiple = round( - ((tar.lon[1] - tar.lon[0]) / (ori.lon[1] - ori.lon[0])).values.tolist() - ) - lat_start = _find_intercept(ori.lat, tar.lat, (lat_multiple - 1) // 2) - lon_start = _find_intercept(ori.lon, tar.lon, (lon_multiple - 1) // 2) - - if func == "mean": - _coarsen = ( - ori.isel(lat=slice(lat_start, None), lon=slice(lon_start, None)) - .coarsen( - dim={"lat": lat_multiple, "lon": lon_multiple}, - side={"lat": "left", "lon": "left"}, - boundary="pad", - ) - .mean() - ) - elif func == "sum": - _coarsen = ( - ori.isel(lat=slice(lat_start, None), lon=slice(lon_start, None)) - .coarsen( - dim={"lat": lat_multiple, "lon": lon_multiple}, - side={"lat": "left", "lon": "left"}, - boundary="pad", - ) - .sum() - ) - else: - raise ValueError("func can only be 'mean' or 'sum'") - - return _coarsen.reindex_like(tar, method="nearest") - - -def calc_grid_area(lis_lats_lons): - """Calculate area in km^2 for a grid cell given lats and lon border, with help from: - https://stackoverflow.com/questions/4681737/how-to-calculate-the-area-of-a-polygon-on-the-earths-surface-using-python - - """ - lons, lats = zip(*lis_lats_lons) - ll = list(set(lats))[::-1] - var = [] - for i in range(len(ll)): - var.append("lat_" + str(i + 1)) - st = "" - for v, l in zip(var, ll): # noqa: E741 - st = st + str(v) + "=" + str(l) + " " + "+" - st = ( - st - + "lat_0=" - + str(np.mean(ll)) - + " " - + "+" - + "lon_0" - + "=" - + str(np.mean(lons)) - ) - tx = "+proj=aea +" + st - pa = pyproj.Proj(tx) - - x, y = pa(lons, lats) - cop = {"type": "Polygon", "coordinates": [zip(x, y)]} - - return shapely.geometry.shape(cop).area / 1000000 - - -def calc_shp_area(shp, shp_projection="+proj=latlon"): - """calculate area in km^2 of the shapes for each shp object""" - temp_shape = shapely.ops.transform( - partial( - pyproj.transform, - pyproj.Proj(shp_projection), - pyproj.Proj(proj="aea", lat_1=shp.bounds[1], lat_2=shp.bounds[3]), - ), - shp, - ) - return temp_shape.area / 1000000 - - -__all__ = ["Cutout", "coarsen", "calc_grid_area", "calc_shp_area"] +__all__ = ["Cutout", "coarsen", "calc_grid_area", "calc_shp_area", "ds_reformat_index"] diff --git a/src/geodata/mask/__init__.py b/src/geodata/mask/__init__.py new file mode 100644 index 00000000..3f4102f2 --- /dev/null +++ b/src/geodata/mask/__init__.py @@ -0,0 +1,41 @@ +"""Mask package namespace. + +This package hosts mask-related modules (e.g. spatial helper utilities) while +preserving backward-compatible access to the legacy ``geodata.mask`` module API. +""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from .spatial import calc_grid_area, calc_shp_area, coarsen, ds_reformat_index +from .xarray_mask import XarrayMask + +_LEGACY_MODULE_PATH = Path(__file__).resolve().parent.parent / "mask.py" +_LEGACY_SPEC = importlib.util.spec_from_file_location( + "geodata._legacy_mask_module", _LEGACY_MODULE_PATH +) +if _LEGACY_SPEC is None or _LEGACY_SPEC.loader is None: + raise ImportError(f"Could not load legacy mask module from {_LEGACY_MODULE_PATH}") +_legacy_mask_module = importlib.util.module_from_spec(_LEGACY_SPEC) +_LEGACY_SPEC.loader.exec_module(_legacy_mask_module) + +# Re-export all public names from legacy ``mask.py``. +for _name in dir(_legacy_mask_module): + if _name.startswith("_"): + continue + globals()[_name] = getattr(_legacy_mask_module, _name) + +# Keep explicit access to phase-1 extracted helpers in this namespace. +globals().update( + { + "ds_reformat_index": ds_reformat_index, + "coarsen": coarsen, + "calc_grid_area": calc_grid_area, + "calc_shp_area": calc_shp_area, + "XarrayMask": XarrayMask, + } +) + +__all__ = [name for name in globals() if not name.startswith("_")] diff --git a/src/geodata/mask/spatial.py b/src/geodata/mask/spatial.py new file mode 100644 index 00000000..9c46544a --- /dev/null +++ b/src/geodata/mask/spatial.py @@ -0,0 +1,129 @@ +"""Shared spatial helper utilities for masking workflows.""" + +from functools import partial +from typing import Any, Literal, cast + +import numpy as np +import pyproj +import shapely +import xarray as xr +from shapely import ops + + +def ds_reformat_index(ds: xr.Dataset | xr.DataArray) -> xr.Dataset | xr.DataArray: + """Normalize data coordinates to sorted ``lat``/``lon``.""" + if "lat" in ds.dims and "lon" in ds.dims: + return ds.sortby(["lat", "lon"]) + if "lat" in ds.coords and "lon" in ds.coords: + return ( + ds.reset_coords(["lon", "lat"], drop=True) + .rename({"x": "lon", "y": "lat"}) + .sortby(["lat", "lon"]) + ) + return ds.rename({"x": "lon", "y": "lat"}).sortby(["lat", "lon"]) + + +def _find_intercept(list1, list2, start, threshold=0): + """Find best start offset for coarsening alignment.""" + min_res = 0 + init = 0 + i = 0 + for i in range(len(list1) - start): + resid = ((list1[start + i] - list2[0]) % (list2[1] - list2[0])).values.tolist() + if i == 0: + init = resid + if resid <= threshold: + return i + if resid > min_res: + min_res = resid + else: + min_res = resid + break + if min_res == init: + return 0 + return i + + +def coarsen( + ori: xr.Dataset | xr.DataArray, + tar: xr.Dataset | xr.DataArray, + func: Literal["sum", "mean"] = "mean", +): + """Reindex/coarsen ``ori`` according to target coordinates in ``tar``.""" + lat_multiple = round( + ((tar.lat[1] - tar.lat[0]) / (ori.lat[1] - ori.lat[0])).values.tolist() + ) + lon_multiple = round( + ((tar.lon[1] - tar.lon[0]) / (ori.lon[1] - ori.lon[0])).values.tolist() + ) + lat_start = _find_intercept(ori.lat, tar.lat, (lat_multiple - 1) // 2) + lon_start = _find_intercept(ori.lon, tar.lon, (lon_multiple - 1) // 2) + + if func == "mean": + coarsened = ( + ori.isel(lat=slice(lat_start, None), lon=slice(lon_start, None)) + .coarsen( + dim={"lat": lat_multiple, "lon": lon_multiple}, + side={"lat": "left", "lon": "left"}, + boundary="pad", + ) + ) + reduced = cast(Any, coarsened).mean() + elif func == "sum": + coarsened = ( + ori.isel(lat=slice(lat_start, None), lon=slice(lon_start, None)) + .coarsen( + dim={"lat": lat_multiple, "lon": lon_multiple}, + side={"lat": "left", "lon": "left"}, + boundary="pad", + ) + ) + reduced = cast(Any, coarsened).sum() + else: + raise ValueError("func can only be 'mean' or 'sum'") + + return reduced.reindex_like(tar, method="nearest") + + +def calc_grid_area(lis_lats_lons): + """Calculate area in km^2 for a grid cell defined by corner coordinates.""" + lons, lats = zip(*lis_lats_lons) + ll = list(set(lats))[::-1] + var = [] + for i in range(len(ll)): + var.append("lat_" + str(i + 1)) + st = "" + for v, l in zip(var, ll): # noqa: E741 + st = st + str(v) + "=" + str(l) + " " + "+" + st = ( + st + + "lat_0=" + + str(np.mean(ll)) + + " " + + "+" + + "lon_0" + + "=" + + str(np.mean(lons)) + ) + tx = "+proj=aea +" + st + pa = pyproj.Proj(tx) + + x, y = pa(lons, lats) + cop = {"type": "Polygon", "coordinates": [zip(x, y)]} + return shapely.geometry.shape(cop).area / 1000000 + + +def calc_shp_area(shp, shp_projection="+proj=latlon"): + """Calculate area in km^2 for a shape object.""" + temp_shape = ops.transform( + partial( + pyproj.transform, + pyproj.Proj(shp_projection), + pyproj.Proj(proj="aea", lat_1=shp.bounds[1], lat_2=shp.bounds[3]), + ), + shp, + ) + return temp_shape.area / 1000000 + + +__all__ = ["ds_reformat_index", "coarsen", "calc_grid_area", "calc_shp_area"] diff --git a/src/geodata/mask/xarray_mask.py b/src/geodata/mask/xarray_mask.py new file mode 100644 index 00000000..18e71d3f --- /dev/null +++ b/src/geodata/mask/xarray_mask.py @@ -0,0 +1,173 @@ +"""Xarray-native mask adapter for applying saved Mask objects to datasets.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, cast + +import numpy as np +import xarray as xr + +from .spatial import calc_grid_area, coarsen, ds_reformat_index + + +def _ensure_dataset(data: xr.Dataset | xr.DataArray) -> xr.Dataset: + if isinstance(data, xr.Dataset): + return data + name = data.name or "value" + return data.to_dataset(name=name) + + +def _to_mask_2d(mask: xr.DataArray) -> xr.DataArray: + mask = mask.reset_coords(drop=True) + if "band" in mask.dims: + mask = mask.isel(band=0, drop=True) + return mask.transpose("lat", "lon") + + +@dataclass +class XarrayMask: + """Mask adapter that aligns saved mask rasters to a target xarray grid.""" + + grid: xr.Dataset + merged_mask: xr.DataArray | None = None + shape_masks: dict[str, xr.DataArray] = field(default_factory=dict) + + @classmethod + def from_mask( + cls, + mask, + grid: xr.Dataset | xr.DataArray, + include_merged: bool = True, + include_shapes: bool = True, + ) -> "XarrayMask": + grid_ds = ds_reformat_index(_ensure_dataset(grid)) + grid_ds = cast(xr.Dataset, grid_ds) + merged = None + shapes: dict[str, xr.DataArray] = {} + + if include_merged and mask.merged_mask: + merged = coarsen(mask.load_merged_xr(), grid_ds) + if include_shapes and mask.shape_mask: + shapes = {k: coarsen(v, grid_ds) for k, v in mask.load_shape_xr().items()} + + if merged is None and not shapes: + raise ValueError( + f"No mask found in {mask.name}. Please create a proper mask object first." + ) + + return cls(grid=grid_ds, merged_mask=merged, shape_masks=shapes) + + @classmethod + def from_name( + cls, + name: str, + grid: xr.Dataset | xr.DataArray, + mask_dir: str | None = None, + include_merged: bool = True, + include_shapes: bool = True, + ) -> "XarrayMask": + from geodata import Mask # lazy import to avoid circular imports + + if mask_dir is None: + from geodata import config + + mask = Mask.from_name(name, mask_dir=config.MASK_DIR) + else: + mask = Mask.from_name(name, mask_dir=mask_dir) + return cls.from_mask( + mask, + grid=grid, + include_merged=include_merged, + include_shapes=include_shapes, + ) + + @staticmethod + def compute_grid_area(grid: xr.Dataset | xr.DataArray) -> xr.DataArray: + xr_ds = ds_reformat_index(_ensure_dataset(grid)) + area_arr = np.zeros((xr_ds.lat.shape[0], xr_ds.lon.shape[0])) + lat_diff = np.abs((xr_ds.lat[1].values - xr_ds.lat[0].values)) + for i, lat in enumerate(xr_ds.lat.values): + lat_bottom = lat - lat_diff / 2 + lat_top = lat + lat_diff / 2 + area_arr[i] = np.round( + calc_grid_area( + [ + (xr_ds.lon.values[0], lat_top), + (xr_ds.lon.values[0], lat_bottom), + (xr_ds.lon.values[1], lat_bottom), + (xr_ds.lon.values[1], lat_top), + ] + ), + 2, + ) + return xr.DataArray( + area_arr, + dims=("lat", "lon"), + coords={"lat": xr_ds.lat.values, "lon": xr_ds.lon.values}, + name="area", + ) + + def _target_masks(self) -> dict[str, xr.DataArray]: + res: dict[str, xr.DataArray] = {} + if self.merged_mask is not None: + res["merged_mask"] = _to_mask_2d(self.merged_mask) + for key, value in self.shape_masks.items(): + res[key] = _to_mask_2d(value) + return res + + def attach( + self, dataset: xr.Dataset | xr.DataArray, include_area: bool = True + ) -> dict[str, xr.Dataset]: + ds = ds_reformat_index(_ensure_dataset(dataset)) + if "time" in ds.dims: + ds = ds.transpose("time", "lat", "lon") + + masks = self._target_masks() + if not masks: + raise ValueError("No masks available in XarrayMask.") + + area = self.compute_grid_area(self.grid) if include_area else None + out: dict[str, xr.Dataset] = {} + for key, mask in masks.items(): + cur = ds.assign({"mask": mask}) + if area is not None: + cur = cur.assign({"area": area}) + out[key] = cur + return out + + def apply( + self, + dataset: xr.Dataset | xr.DataArray, + mode: Literal["where", "multiply"] = "where", + include_area: bool = False, + ) -> dict[str, xr.Dataset]: + ds = ds_reformat_index(_ensure_dataset(dataset)) + if "time" in ds.dims: + ds = ds.transpose("time", "lat", "lon") + + masks = self._target_masks() + if not masks: + raise ValueError("No masks available in XarrayMask.") + if mode not in {"where", "multiply"}: + raise ValueError("mode can only be 'where' or 'multiply'") + + area = self.compute_grid_area(self.grid) if include_area else None + out: dict[str, xr.Dataset] = {} + for key, mask in masks.items(): + valid = mask > 0 + cur = ds.copy() + for var in list(cur.data_vars): + da = cur[var] + if "lat" in da.dims and "lon" in da.dims: + if mode == "where": + cur = cur.assign({var: cast(Any, da.where(valid))}) + else: + cur = cur.assign({var: cast(Any, da * valid)}) + if include_area and area is not None: + cur = cur.assign({"area": area}) + out[key] = cast(xr.Dataset, cur) + return out + + +__all__ = ["XarrayMask"] diff --git a/src/geodata/model/wind/_base.py b/src/geodata/model/wind/_base.py index 244f81e3..370d154a 100644 --- a/src/geodata/model/wind/_base.py +++ b/src/geodata/model/wind/_base.py @@ -40,6 +40,7 @@ """ import xarray as xr +from typing import Any, cast from ...resource import get_windturbineconfig from .._base import BaseModel @@ -107,7 +108,7 @@ def _estimate_power( ys: slice | None = None, years: slice | None = None, months: slice | None = None, - ) -> None: + ) -> xr.DataArray: """Estimate wind speed at the given locations and times. Args: @@ -135,7 +136,7 @@ def _estimate_power( turbineconf["V"], turbineconf["POW"], bounds_error=False, - fill_value="extrapolate", + fill_value=cast(Any, "extrapolate"), ) # Calculate the power output @@ -147,4 +148,4 @@ def _estimate_power( output_dtypes=[float], ) - return xr.Dataset({"cf": power / turbineconf["P"]}) + return (power / turbineconf["P"]).rename("cf") diff --git a/src/geodata/plot.py b/src/geodata/plot.py index f6e47872..f12c1fdf 100644 --- a/src/geodata/plot.py +++ b/src/geodata/plot.py @@ -22,7 +22,7 @@ import matplotlib.pyplot as plt import xarray as xr -from .cutout import ds_reformat_index +from .mask.spatial import ds_reformat_index from .mask import show # noqa: F401 plt.rcParams["animation.html"] = "jshtml" diff --git a/tests/pr/mask/test_mask_legacy_error_paths.py b/tests/pr/mask/test_mask_legacy_error_paths.py new file mode 100644 index 00000000..7770be6c --- /dev/null +++ b/tests/pr/mask/test_mask_legacy_error_paths.py @@ -0,0 +1,110 @@ +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pytest +import xarray as xr +from rasterio.transform import from_bounds + +from geodata.cutout import Cutout +from geodata.mask import Mask, save_raster + + +def _build_minimal_cutout() -> Cutout: + cutout = Cutout.__new__(Cutout) + cutout.name = "legacy-error-cutout" + cutout.meta = xr.Dataset( + coords={ + "x": np.array([100.0, 100.25, 100.5]), + "y": np.array([30.5, 30.25, 30.0]), + "year": [2016], + "month": [1], + } + ) + cutout.merged_mask = None + cutout.shape_mask = None + cutout.area = None + cutout.prepared = True + cutout.empty = False + cutout.cutout_dir = Path(".") + return cutout + + +def _sample_dataset() -> xr.Dataset: + t = np.array(["2016-01-01T00:00:00"], dtype="datetime64[ns]") + y = np.array([30.5, 30.25, 30.0]) + x = np.array([100.0, 100.25, 100.5]) + data = np.arange(len(t) * len(y) * len(x), dtype=np.float32).reshape( + len(t), len(y), len(x) + ) + return xr.Dataset( + {"signal": (("time", "y", "x"), data)}, + coords={"time": t, "y": y, "x": x}, + ) + + +def _create_saved_empty_mask(mask_dir: Path, name: str) -> None: + # Create a mask object that is saved but has no merged/shape masks. + mask = Mask(name=name, mask_dir=str(mask_dir)) + mask.save_mask() + + +def _create_unsaved_mask_with_layer(mask_dir: Path, name: str) -> Mask: + west, south, east, north = 100.0, 30.0, 100.75, 30.75 + arr = np.ones((3, 3), dtype=np.uint8) + transform = from_bounds(west, south, east, north, arr.shape[1], arr.shape[0]) + layer_path = mask_dir / f"{name}.tif" + save_raster(arr, transform, str(layer_path)) + mask = Mask(name=name, mask_dir=str(mask_dir)) + mask.add_layer(str(layer_path), layer_name="base") + return mask + + +def test_mask_raises_without_added_masks(): + cutout = _build_minimal_cutout() + ds = _sample_dataset() + + with pytest.raises(ValueError, match="No mask found in cutout"): + cutout.mask(ds) + + +def test_mask_raises_when_true_area_requested_without_area(): + cutout = _build_minimal_cutout() + ds = _sample_dataset() + cutout.merged_mask = xr.DataArray( + np.ones((1, 3, 3), dtype=np.float32), + dims=("band", "lat", "lon"), + coords={ + "band": [1], + "lat": ds["y"].values, + "lon": ds["x"].values, + }, + ) + + with pytest.raises(ValueError, match="No area data found"): + cutout.mask(ds, true_area=True) + + +def test_add_mask_raises_for_saved_mask_without_merged_or_shape(tmp_path, monkeypatch): + cutout = _build_minimal_cutout() + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + name = "empty_saved_mask" + _create_saved_empty_mask(mask_dir, name) + + monkeypatch.setattr("geodata.cutout.config.MASK_DIR", str(mask_dir)) + + with pytest.raises(ValueError, match=f"No mask found in {name}"): + cutout.add_mask(name) + + +def test_mask_load_xarray_raises_when_unsaved(tmp_path): + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + mask = _create_unsaved_mask_with_layer(mask_dir, name="unsaved_mask") + + with pytest.raises(ValueError, match="has not been saved"): + mask.load_merged_xr() + + with pytest.raises(ValueError, match="has not been saved"): + _ = mask.load_shape_xr(names=cast(Any, [])) diff --git a/tests/pr/mask/test_mask_legacy_workflow.py b/tests/pr/mask/test_mask_legacy_workflow.py new file mode 100644 index 00000000..6d75a43e --- /dev/null +++ b/tests/pr/mask/test_mask_legacy_workflow.py @@ -0,0 +1,173 @@ +import uuid +from pathlib import Path +from typing import Any, cast + +import numpy as np +import shapely.geometry +import xarray as xr +from rasterio.transform import from_bounds + +from geodata.cutout import Cutout, calc_grid_area, coarsen, ds_reformat_index +from geodata.datasets import load_dataset +from geodata.mask import Mask, save_raster + + +def _build_cutout(tmp_path: Path) -> Cutout: + dataset_cls = load_dataset("wind_solar_hourly_test") + dataset = dataset_cls(years=slice(2016, 2016), months=slice(1, 1), testing=True) + assert dataset.downloaded, "Fixture NetCDF should be present" + + with xr.open_dataset(dataset.catalog[0].path, engine="h5netcdf") as opened: + if "x" in opened.coords and "y" in opened.coords: + xvals = opened["x"].values + yvals = opened["y"].values + else: + xvals = opened["longitude"].values + yvals = opened["latitude"].values + + # Use a lightweight Cutout instance that still exercises legacy methods + # (add_mask, add_grid_area, mask) without invoking dataset preparation. + cutout = Cutout.__new__(Cutout) + cutout.name = f"legacy-mask-test-{uuid.uuid4().hex[:8]}" + cutout.meta = xr.Dataset( + coords={ + "x": xvals, + "y": yvals, + "year": [2016], + "month": [1], + } + ) + cutout.merged_mask = None + cutout.shape_mask = None + cutout.area = None + cutout.prepared = True + cutout.empty = False + cutout.cutout_dir = tmp_path / "cutouts" + return cutout + + +def _create_and_save_mask(cutout: Cutout, mask_dir: Path, name: str = "legacy_test_mask") -> None: + assert cutout.meta is not None + xr_meta = ds_reformat_index(cast(Any, cutout.meta)) + lon = xr_meta["lon"].values + lat = xr_meta["lat"].values + + lon_step = float(np.abs(lon[1] - lon[0])) + lat_step = float(np.abs(lat[1] - lat[0])) + west = float(lon.min() - lon_step / 2) + east = float(lon.max() + lon_step / 2) + south = float(lat.min() - lat_step / 2) + north = float(lat.max() + lat_step / 2) + + nlon_hi = len(lon) * 2 + nlat_hi = len(lat) * 2 + transform = from_bounds(west, south, east, north, nlon_hi, nlat_hi) + + raster = np.zeros((nlat_hi, nlon_hi), dtype=np.uint8) + # Non-trivial pattern so coarsening does real work. + raster[nlat_hi // 4 : 3 * nlat_hi // 4, nlon_hi // 6 : 5 * nlon_hi // 6] = 1 + + layer_path = mask_dir / "source_layer.tif" + save_raster(raster, transform, str(layer_path)) + + mask = Mask(name=name, mask_dir=str(mask_dir)) + mask.add_layer(str(layer_path), layer_name="source") + mask.merge_layer(show_raster=False) + + centroid_lon = float(np.mean([west, east])) + centroid_lat = float(np.mean([south, north])) + shape = shapely.geometry.box( + west, + south, + centroid_lon, + centroid_lat, + ) + mask.extract_shapes({"region_a": shape}, show_raster=False) + mask.save_mask() + + +def test_legacy_mask_workflow_contract_offline(tmp_path, monkeypatch): + cutout = _build_cutout(tmp_path) + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + mask_name = "legacy_test_mask" + _create_and_save_mask(cutout, mask_dir, name=mask_name) + + monkeypatch.setattr("geodata.cutout.config.MASK_DIR", str(mask_dir)) + cutout.add_mask(mask_name) + cutout.add_grid_area() + + time = np.array(["2016-01-01T00:00:00", "2016-01-01T01:00:00"], dtype="datetime64[ns]") + y = cutout.coords["y"].values + x = cutout.coords["x"].values + payload = np.arange(len(time) * len(y) * len(x), dtype=np.float32).reshape( + len(time), len(y), len(x) + ) + ds = xr.Dataset( + {"signal": (("time", "y", "x"), payload)}, + coords={"time": time, "y": y, "x": x}, + ) + + masked = cutout.mask(ds) + + assert set(masked.keys()) == {"merged_mask", "region_a"} + merged = masked["merged_mask"] + assert isinstance(merged, xr.Dataset) + assert {"signal", "mask", "area"}.issubset(set(merged.data_vars)) + assert tuple(merged["signal"].dims) == ("time", "lat", "lon") + assert tuple(merged["mask"].dims) == ("lat", "lon") + assert tuple(merged["area"].dims) == ("lat", "lon") + + +def test_legacy_add_mask_coarsen_parity_offline(tmp_path, monkeypatch): + cutout = _build_cutout(tmp_path) + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + mask_name = "legacy_test_mask" + _create_and_save_mask(cutout, mask_dir, name=mask_name) + + monkeypatch.setattr("geodata.cutout.config.MASK_DIR", str(mask_dir)) + cutout.add_mask(mask_name, shape_mask=False) + + mask = Mask.from_name(mask_name, mask_dir=str(mask_dir)) + assert cutout.meta is not None + expected = coarsen( + cast(Any, mask.load_merged_xr()), + cast(Any, ds_reformat_index(cast(Any, cutout.meta))), + ) + + assert cutout.merged_mask is not None + np.testing.assert_allclose(cutout.merged_mask.values, expected.values) + assert cutout.merged_mask.shape == expected.shape + + +def test_legacy_add_grid_area_sanity_offline(tmp_path): + cutout = _build_cutout(tmp_path) + cutout.add_grid_area() + + assert cutout.area is not None + area = cutout.area["area"].values + assert np.all(np.isfinite(area)) + assert np.all(area > 0) + + # Area should be constant across longitude for a given latitude row. + row_std = area.std(axis=1) + assert np.allclose(row_std, 0.0, atol=1e-6) + + assert cutout.meta is not None + xr_ds = ds_reformat_index(cast(Any, cutout.meta)) + lat = xr_ds.lat.values + lon = xr_ds.lon.values + lat_diff = float(np.abs(lat[1] - lat[0])) + expected_first_row = np.round( + calc_grid_area( + [ + (lon[0], lat[0] + lat_diff / 2), + (lon[0], lat[0] - lat_diff / 2), + (lon[1], lat[0] - lat_diff / 2), + (lon[1], lat[0] + lat_diff / 2), + ] + ), + 2, + ) + assert np.isclose(area[0, 0], expected_first_row) diff --git a/tests/pr/mask/test_mask_spatial_helpers.py b/tests/pr/mask/test_mask_spatial_helpers.py new file mode 100644 index 00000000..af8ba604 --- /dev/null +++ b/tests/pr/mask/test_mask_spatial_helpers.py @@ -0,0 +1,49 @@ +import numpy as np +import xarray as xr + +from geodata.mask.spatial import calc_grid_area, coarsen, ds_reformat_index + + +def test_ds_reformat_index_renames_and_sorts_xy(): + x = np.array([101.0, 100.5, 100.0]) + y = np.array([30.0, 30.5, 31.0]) + arr = np.arange(9, dtype=np.float32).reshape(3, 3) + da = xr.DataArray(arr, dims=("y", "x"), coords={"x": x, "y": y}, name="signal") + + out = ds_reformat_index(da) + assert out.dims == ("lat", "lon") + assert np.all(np.diff(out["lat"].values) >= 0) + assert np.all(np.diff(out["lon"].values) >= 0) + + +def test_coarsen_mean_on_aligned_grid(): + lat_hi = np.array([0.0, 0.25, 0.5, 0.75]) + lon_hi = np.array([10.0, 10.25, 10.5, 10.75]) + hi = xr.DataArray( + np.arange(16, dtype=np.float32).reshape(4, 4), + dims=("lat", "lon"), + coords={"lat": lat_hi, "lon": lon_hi}, + name="mask", + ) + + lat_lo = np.array([0.125, 0.625]) + lon_lo = np.array([10.125, 10.625]) + lo = xr.Dataset(coords={"lat": lat_lo, "lon": lon_lo}) + + out = coarsen(hi, lo, func="mean") + # Freeze current legacy coarsen behavior. + expected = np.array([[7.5, 9.0], [13.5, 15.0]], dtype=np.float32) + np.testing.assert_allclose(out.values, expected, atol=1e-6) + + +def test_calc_grid_area_positive_and_latitude_sensitive(): + # Avoid perfectly symmetric parallels around 0 that can trip AEA constraints. + cell_low_lat = [(0.0, 1.5), (0.0, 0.5), (1.0, 0.5), (1.0, 1.5)] + cell_high_lat = [(0.0, 60.5), (0.0, 59.5), (1.0, 59.5), (1.0, 60.5)] + + area_low_lat = calc_grid_area(cell_low_lat) + area_high_lat = calc_grid_area(cell_high_lat) + + assert area_low_lat > 0 + assert area_high_lat > 0 + assert area_low_lat > area_high_lat diff --git a/tests/pr/mask/test_xarray_mask.py b/tests/pr/mask/test_xarray_mask.py new file mode 100644 index 00000000..02ab2e33 --- /dev/null +++ b/tests/pr/mask/test_xarray_mask.py @@ -0,0 +1,136 @@ +from pathlib import Path +from typing import Any, cast + +import numpy as np +import shapely.geometry +import xarray as xr +import rasterio as ras +from rasterio.transform import from_bounds + +from geodata import Mask, XarrayMask +from geodata.cutout import Cutout, ds_reformat_index + + +def _build_minimal_cutout() -> Cutout: + cutout = Cutout.__new__(Cutout) + cutout.name = "xarray-mask-test" + cutout.meta = xr.Dataset( + coords={ + "x": np.array([100.0, 100.25, 100.5, 100.75]), + "y": np.array([30.75, 30.5, 30.25, 30.0]), + "year": [2016], + "month": [1], + } + ) + cutout.merged_mask = None + cutout.shape_mask = None + cutout.area = None + cutout.prepared = True + cutout.empty = False + cutout.cutout_dir = Path(".") + return cutout + + +def _create_saved_mask(cutout: Cutout, mask_dir: Path, name: str = "xarray_test_mask") -> None: + assert cutout.meta is not None + xr_meta = ds_reformat_index(cast(Any, cutout.meta)) + lon = xr_meta["lon"].values + lat = xr_meta["lat"].values + + lon_step = float(np.abs(lon[1] - lon[0])) + lat_step = float(np.abs(lat[1] - lat[0])) + west = float(lon.min() - lon_step / 2) + east = float(lon.max() + lon_step / 2) + south = float(lat.min() - lat_step / 2) + north = float(lat.max() + lat_step / 2) + + nlon_hi = len(lon) * 2 + nlat_hi = len(lat) * 2 + transform = from_bounds(west, south, east, north, nlon_hi, nlat_hi) + + arr = np.zeros((nlat_hi, nlon_hi), dtype=np.uint8) + arr[nlat_hi // 4 : 3 * nlat_hi // 4, nlon_hi // 4 : 3 * nlon_hi // 4] = 1 + layer_path = mask_dir / "source.tif" + with ras.open( + str(layer_path), + "w", + driver="GTiff", + height=arr.shape[0], + width=arr.shape[1], + count=1, + dtype=arr.dtype, + compress="lzw", + crs="+proj=latlong", + transform=transform, + ) as dst: + dst.write(arr, 1) + + mask = Mask(name=name, mask_dir=str(mask_dir)) + mask.add_layer(str(layer_path), layer_name="source") + mask.merge_layer(show_raster=False) + shape = shapely.geometry.box(west, south, (west + east) / 2, (south + north) / 2) + mask.extract_shapes({"region_a": shape}, show_raster=False) + mask.save_mask() + + +def _sample_dataset_from_cutout(cutout: Cutout) -> xr.Dataset: + assert cutout.meta is not None + y = cutout.meta["y"].values + x = cutout.meta["x"].values + t = np.array(["2016-01-01T00:00:00", "2016-01-01T01:00:00"], dtype="datetime64[ns]") + vals = np.arange(len(t) * len(y) * len(x), dtype=np.float32).reshape( + len(t), len(y), len(x) + ) + return xr.Dataset({"signal": (("time", "y", "x"), vals)}, coords={"time": t, "y": y, "x": x}) + + +def test_xarraymask_attach_matches_legacy_contract(tmp_path, monkeypatch): + cutout = _build_minimal_cutout() + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + mask_name = "xarray_test_mask" + _create_saved_mask(cutout, mask_dir, name=mask_name) + + monkeypatch.setattr("geodata.cutout.config.MASK_DIR", str(mask_dir)) + cutout.add_mask(mask_name) + cutout.add_grid_area() + + ds = _sample_dataset_from_cutout(cutout) + legacy = cutout.mask(ds) + + assert cutout.meta is not None + xmask = XarrayMask.from_name(mask_name, grid=cutout.meta, mask_dir=str(mask_dir)) + attached = xmask.attach(ds, include_area=True) + + assert set(attached.keys()) == set(legacy.keys()) + for key in attached: + xr.testing.assert_allclose(attached[key]["mask"], legacy[key]["mask"]) + xr.testing.assert_allclose(attached[key]["area"], legacy[key]["area"]) + xr.testing.assert_allclose(attached[key]["signal"], legacy[key]["signal"]) + + +def test_xarraymask_apply_where_and_multiply(tmp_path): + cutout = _build_minimal_cutout() + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + mask_name = "xarray_test_mask" + _create_saved_mask(cutout, mask_dir, name=mask_name) + + ds = _sample_dataset_from_cutout(cutout) + assert cutout.meta is not None + xmask = XarrayMask.from_name(mask_name, grid=cutout.meta, mask_dir=str(mask_dir)) + + attached = xmask.attach(ds, include_area=False) + merged_mask = attached["merged_mask"]["mask"] + + where_out = xmask.apply(ds, mode="where", include_area=True)["merged_mask"] + multiply_out = xmask.apply(ds, mode="multiply", include_area=False)["merged_mask"] + + valid = merged_mask > 0 + expected_where = attached["merged_mask"]["signal"].where(valid) + expected_multiply = attached["merged_mask"]["signal"] * valid + + xr.testing.assert_allclose(where_out["signal"], expected_where) + xr.testing.assert_allclose(multiply_out["signal"], expected_multiply) + assert "area" in where_out + assert "area" not in multiply_out diff --git a/tests/pr/test_wind_xarraymask_integration.py b/tests/pr/test_wind_xarraymask_integration.py new file mode 100644 index 00000000..d509e369 --- /dev/null +++ b/tests/pr/test_wind_xarraymask_integration.py @@ -0,0 +1,101 @@ +from pathlib import Path + +import numpy as np +import rasterio as ras +import xarray as xr +from dask.distributed import Client +from rasterio.transform import from_bounds + +from geodata import XarrayMask +from geodata.datasets import load_dataset +from geodata.model.wind import WindInterpolationModel + + +def _create_saved_mask_from_output_grid( + output: xr.DataArray, + mask_dir: Path, + name: str = "wind_xmask", +) -> None: + x = output["x"].values + y = output["y"].values + + lon = np.sort(np.asarray(x, dtype=float)) + lat = np.sort(np.asarray(y, dtype=float)) + lon_step = float(np.abs(lon[1] - lon[0])) + lat_step = float(np.abs(lat[1] - lat[0])) + west = float(lon.min() - lon_step / 2) + east = float(lon.max() + lon_step / 2) + south = float(lat.min() - lat_step / 2) + north = float(lat.max() + lat_step / 2) + + nlon_hi = len(lon) * 2 + nlat_hi = len(lat) * 2 + transform = from_bounds(west, south, east, north, nlon_hi, nlat_hi) + + arr = np.zeros((nlat_hi, nlon_hi), dtype=np.uint8) + arr[nlat_hi // 4 : 3 * nlat_hi // 4, nlon_hi // 4 : 3 * nlon_hi // 4] = 1 + + source_tif = mask_dir / "source.tif" + with ras.open( + str(source_tif), + "w", + driver="GTiff", + height=arr.shape[0], + width=arr.shape[1], + count=1, + dtype=arr.dtype, + compress="lzw", + crs="+proj=latlong", + transform=transform, + ) as dst: + dst.write(arr, 1) + + from geodata import Mask + + mask = Mask(name=name, mask_dir=str(mask_dir)) + mask.add_layer(str(source_tif), layer_name="source") + mask.merge_layer(show_raster=False) + mask.save_mask() + + +def test_wind_estimate_with_xarray_mask_offline(tmp_path): + years = slice(2016, 2016) + months = slice(1, 1) + + with Client(processes=True, threads_per_worker=1): + ds_cls = load_dataset("wind_3d_hourly_test") + ds = ds_cls(years=years, months=months) + assert ds.downloaded, "Wind fixture NetCDF should be present" + + model = WindInterpolationModel(ds) + model.prepare(force=True) + + base = model.estimate(years=years, months=months, height=12) + assert isinstance(base, xr.DataArray) + + mask_dir = tmp_path / "masks" + mask_dir.mkdir(parents=True, exist_ok=True) + mask_name = "wind_xmask" + _create_saved_mask_from_output_grid(base, mask_dir, name=mask_name) + + base_ds = base.to_dataset(name=base.name or "value") + xmask = XarrayMask.from_name(mask_name, grid=base_ds, mask_dir=str(mask_dir)) + masked = xmask.apply( + base_ds, + mode="where", + include_area=True, + ) + + assert isinstance(masked, dict) + assert set(masked.keys()) == {"merged_mask"} + + merged = masked["merged_mask"] + assert "area" in merged + value_vars = [v for v in merged.data_vars if v not in {"area"}] + assert len(value_vars) == 1 + var = value_vars[0] + + attached = xmask.attach(base, include_area=False)["merged_mask"] + valid = attached["mask"] > 0 + expected = attached[var].where(valid) + xr.testing.assert_allclose(merged[var], expected)