feat: implement xarray-based masking workflow and migration plan

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

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.

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__"]