diff --git a/bluemath_tk/core/metrics.py b/bluemath_tk/core/metrics.py index 84eb62a..b6fa7e7 100644 --- a/bluemath_tk/core/metrics.py +++ b/bluemath_tk/core/metrics.py @@ -28,6 +28,7 @@ from collections.abc import Callable, Iterable import numpy as np +import pandas as pd ArrayLike = np.ndarray | Iterable[float] @@ -58,7 +59,9 @@ def paired_finite( - reference: ArrayLike, estimate: ArrayLike + reference: ArrayLike, + estimate: ArrayLike, + bounds: tuple[float, float] | None = None, ) -> tuple[np.ndarray, np.ndarray]: """ Flatten two paired arrays and keep only the pairs finite on both sides. @@ -69,6 +72,9 @@ def paired_finite( Reference values. estimate : array-like Estimated values, paired element-wise with *reference*. + bounds : tuple of float, optional + ``(low, high)``: also drop pairs with either side outside it, e.g. + corrupted-but-finite model output. Default is None (finite only). Returns ------- @@ -89,6 +95,9 @@ def paired_finite( f"({ref.size} != {est.size})" ) mask = np.isfinite(ref) & np.isfinite(est) + if bounds is not None: + low, high = bounds + mask &= (ref >= low) & (ref <= high) & (est >= low) & (est <= high) return ref[mask], est[mask] @@ -639,6 +648,7 @@ def compute_metrics( *, circular: bool = False, rrmse_normalization: str = "mean", + bounds: tuple[float, float] | None = None, ) -> dict[str, float]: """ Compute several metrics at once on the finite pairs of two arrays. @@ -658,6 +668,8 @@ def compute_metrics( ``NaN`` so tables keep the same columns. Default is False. rrmse_normalization : {"mean", "rms"}, optional Passed to :func:`rrmse`. Default is "mean". + bounds : tuple of float, optional + Passed to :func:`paired_finite`. Default is None. Returns ------- @@ -682,7 +694,7 @@ def compute_metrics( f"Unknown metric(s) {unknown}; available: {list(AVAILABLE_METRICS)}" ) - ref, est = paired_finite(reference, estimate) + ref, est = paired_finite(reference, estimate, bounds) out: dict[str, float] = {"n": int(ref.size)} for name in names: if ref.size == 0: @@ -696,3 +708,49 @@ def compute_metrics( out[name] = _LINEAR[name](ref, est) return out + + +def pool_metrics( + table: pd.DataFrame, + by: Iterable[str], + metrics: Iterable[str] = ("bias", "rmse", "rrmse", "r2"), +) -> pd.DataFrame: + """ + Pool per-group metrics into one row per *by* group, weighted by ``n``. + + ``bias`` (n-weighted mean) and ``rmse`` (root of the n-weighted mean + square) are the exact pooled values. Ratio and correlation metrics + (``rrmse``, ``r2``, ...) need the raw pairs to be pooled exactly, so they + are n-weighted means, an approximation. + + Parameters + ---------- + table : pd.DataFrame + One row per group, with an ``n`` column and the *metrics* columns. + by : iterable of str + Columns that define the pooled rows. + metrics : iterable of str, optional + Metric columns to pool. + + Returns + ------- + pd.DataFrame + The *by* columns, the pooled *metrics* and the summed ``n``. + """ + + by, metrics = list(by), list(metrics) + + def pool(group: pd.DataFrame) -> pd.Series: + w = group["n"].to_numpy(dtype=float) + if w.sum() <= 0: + return pd.Series({m: np.nan for m in metrics} | {"n": 0}) + out = { + m: float(np.sqrt(np.average(group[m] ** 2, weights=w))) + if m == "rmse" + else float(np.average(group[m], weights=w)) + for m in metrics + } + return pd.Series(out | {"n": int(w.sum())}) + + out = table.groupby(by, sort=False)[["n", *metrics]].apply(pool).reset_index() + return out.astype({"n": int}) diff --git a/bluemath_tk/core/models.py b/bluemath_tk/core/models.py index 43f4c9b..4bb7b0a 100644 --- a/bluemath_tk/core/models.py +++ b/bluemath_tk/core/models.py @@ -20,7 +20,7 @@ from sklearn.preprocessing import StandardScaler from .constants import EARTH_RADIUS -from .logging import get_file_logger +from .logging import _coerce_level, _console_handlers, get_file_logger from .operations import ( denormalize, destandarize, @@ -31,6 +31,36 @@ ) +#: Console level of every model logger (PCA, MDA, GP, ...); see :func:`set_models_console_level`. +MODELS_CONSOLE_LEVEL: Union[int, str] = "WARNING" + + +def set_models_console_level(level: Union[int, str]) -> None: + """ + Console level of the model loggers, existing and future (files keep their level). + + Models log routine messages on every fit; ``"ERROR"`` keeps a notebook + that fits many models readable while ``logs/*.log`` keep the detail. + + Parameters + ---------- + level : int or str + Logging level, e.g. ``"ERROR"``. + """ + + global MODELS_CONSOLE_LEVEL + MODELS_CONSOLE_LEVEL = level + names = set() + stack = [BlueMathModel] + while stack: + cls = stack.pop() + names.add(cls.__name__) + stack.extend(cls.__subclasses__()) + for name in names & set(logging.Logger.manager.loggerDict): + for handler in _console_handlers(logging.getLogger(name)): + handler.setLevel(_coerce_level(level)) + + class BlueMathModel(ABC): """ Abstract base class for handling default functionalities across the project. @@ -148,8 +178,13 @@ def logger(self) -> logging.Logger: - Thread-safe logger creation """ - if self._logger is None: - self._logger = get_file_logger(name=self.__class__.__name__) + # an unpickled model refers to its logger by name: in a new process + # that logger has no handlers yet + if self._logger is None or not self._logger.handlers: + self._logger = get_file_logger( + name=self._logger.name if self._logger else self.__class__.__name__, + console_level=MODELS_CONSOLE_LEVEL, + ) return self._logger @logger.setter @@ -200,7 +235,9 @@ def set_logger_name( >>> model.set_logger_name("my_model", level="DEBUG", console=False) """ - self.logger = get_file_logger(name=name, level=level, console=console) + self.logger = get_file_logger( + name=name, level=level, console=console, console_level=MODELS_CONSOLE_LEVEL + ) def save_model(self, model_path: str, exclude_attributes: List[str] = None) -> None: """ diff --git a/bluemath_tk/core/plotting/labels.py b/bluemath_tk/core/plotting/labels.py new file mode 100644 index 0000000..068bdb1 --- /dev/null +++ b/bluemath_tk/core/plotting/labels.py @@ -0,0 +1,89 @@ +""" +Point labels that do not overlap each other. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +from matplotlib.axes import Axes + +#: Candidate offsets (points, ha, va): 8 compass directions per ring, near to far. +_OFFSETS: list[tuple[tuple[float, float], str, str]] = [ + ((radius * ux / ((ux**2 + uy**2) ** 0.5), radius * uy / ((ux**2 + uy**2) ** 0.5)), ha, va) + for radius in (6, 11, 16, 22, 28, 34, 42, 55) + for ux, uy, ha, va in ( + (1, 1, "left", "bottom"), + (1, -1, "left", "top"), + (-1, 1, "right", "bottom"), + (-1, -1, "right", "top"), + (0, 1, "center", "bottom"), + (0, -1, "center", "top"), + (1, 0, "left", "center"), + (-1, 0, "right", "center"), + ) +] + + +def _overlap(a: Sequence[float], b: Sequence[float], margin: float) -> bool: + return not ( + a[2] + margin < b[0] or b[2] + margin < a[0] or a[3] + margin < b[1] or b[3] + margin < a[1] + ) + + +def place_labels( + ax: Axes, + x: Sequence[float], + y: Sequence[float], + texts: Sequence[str], + *, + leader_colors: Sequence[str] | None = None, + **text_kwargs, +) -> None: + """ + Annotate points, each at the nearest offset that clears the labels already placed. + + The axes limits must be final (the label sizes are measured on the + rendered figure). A label pushed far from its point gets a leader line. + + Parameters + ---------- + ax : Axes + Axes with its final limits. + x, y : sequence of float + Points in data coordinates. + texts : sequence of str + One label per point. + leader_colors : sequence of str, optional + Leader-line colour per point. Default is grey. + **text_kwargs + Passed to ``ax.annotate`` (``fontsize``, ``path_effects``, ...). + """ + + ax.figure.canvas.draw() + renderer = ax.figure.canvas.get_renderer() + inv = ax.transData.inverted() + pad = abs(inv.transform((3.5 * ax.figure.dpi / 72.0, 0))[0] - inv.transform((0, 0))[0]) + placed: list[tuple[float, float, float, float]] = [] + for i, (px, py, text) in enumerate(zip(x, y, texts)): + for offset, ha, va in _OFFSETS: + probe = ax.annotate(text, (px, py), xytext=offset, textcoords="offset points", + ha=ha, va=va, **text_kwargs) + bbox = probe.get_window_extent(renderer) + probe.remove() + (x0, y0), (x1, y1) = inv.transform((bbox.x0, bbox.y0)), inv.transform((bbox.x1, bbox.y1)) + box = (min(x0, x1), min(y0, y1), max(x0, x1), max(y0, y1)) + if not any(_overlap(box, other, pad) for other in placed): + break + placed.append(box) + leader = (offset[0] ** 2 + offset[1] ** 2) ** 0.5 > 9 + color = leader_colors[i] if leader_colors is not None else "0.4" + ax.annotate( + text, (px, py), xytext=offset, textcoords="offset points", ha=ha, va=va, + arrowprops=dict(arrowstyle="-", color=color, lw=0.7, alpha=0.75, shrinkA=0, shrinkB=3) + if leader else None, + **text_kwargs, + ) + + +__all__ = ["place_labels"] diff --git a/bluemath_tk/core/plotting/scatter.py b/bluemath_tk/core/plotting/scatter.py index e09dad6..7c5beaa 100644 --- a/bluemath_tk/core/plotting/scatter.py +++ b/bluemath_tk/core/plotting/scatter.py @@ -126,6 +126,7 @@ def validation_scatter( seed: int = 0, fmt: str = ".3f", rrmse_normalization: str = "mean", + bounds: tuple[float, float] | None = None, ) -> dict[str, float]: """ Density scatter of an estimate against its reference, with metrics. @@ -177,6 +178,9 @@ def validation_scatter( Seed for the drawing subsample. Default is 0. fmt : str, optional Number format for the annotation. Default is ".3f". + bounds : tuple of float, optional + Drop pairs outside ``(low, high)`` before plotting and scoring (see + :func:`bluemath_tk.core.metrics.paired_finite`). rrmse_normalization : {"mean", "rms"}, optional Passed to :func:`bluemath_tk.core.metrics.rrmse`. Default is "mean". @@ -186,7 +190,7 @@ def validation_scatter( The metrics computed on all finite pairs (always includes ``n``). """ - ref, est = paired_finite(reference, estimate) + ref, est = paired_finite(reference, estimate, bounds) if circular: ref, est = np.mod(ref, 360.0), np.mod(est, 360.0) diff --git a/bluemath_tk/interpolation/gps.py b/bluemath_tk/interpolation/gps.py index 8dba72e..9fdc764 100644 --- a/bluemath_tk/interpolation/gps.py +++ b/bluemath_tk/interpolation/gps.py @@ -683,6 +683,29 @@ def fit( self.is_fitted = True self.logger.info("GP models fitted successfully") + def to(self, device: str | torch.device) -> "ExactGPInterpolation": + """ + Move the fitted models and likelihoods to *device* (e.g. ``"cpu"``). + + Parameters + ---------- + device : str or torch.device + Target device for prediction. + + Returns + ------- + ExactGPInterpolation + This instance. + """ + + self.device = torch.device(device) + for model in self._models.values(): + model.to(self.device) + for likelihood in self._likelihoods.values(): + likelihood.to(self.device) + + return self + def predict( self, dataset: pd.DataFrame, diff --git a/bluemath_tk/topo_bathy/mesh.py b/bluemath_tk/topo_bathy/mesh.py new file mode 100644 index 0000000..ceff94b --- /dev/null +++ b/bluemath_tk/topo_bathy/mesh.py @@ -0,0 +1,276 @@ +""" +Unstructured model meshes: one reader, depth at points, and a raster of the mesh. + +A wave model runs on its own mesh, which need not cover the same area as the +bathymetry raster used to draw contours. :func:`read_mesh` reads the usual +formats into one :class:`Mesh` (elevation, positive up, as GEBCO): + +- UGRID netCDF (``mesh2d_*``): triangles, quads or a mix (e.g. a D-Flow FM + quadtree), with fill values for the missing nodes of smaller faces; +- ADCIRC ``fort.14`` / ``.grd`` (depth positive down in the file), used by SWAN. + +:func:`depth_at_points` interpolates the mesh depth linearly inside its +triangles (NaN outside the mesh), :func:`mesh_to_raster` samples the mesh on +a regular grid (NaN outside), and :func:`mask_raster_to_mesh` blanks a raster +outside the mesh, so contours traced on it stay inside the model domain. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import xarray as xr + + +@dataclass +class Mesh: + """ + An unstructured mesh in one convention. + + Attributes + ---------- + x, y : np.ndarray + Node coordinates (lon/lat degrees or projected metres). + faces : np.ndarray + ``(n_faces, max_nodes)`` 0-based node indices, ``-1`` where a face has + fewer nodes (triangles in a quad mesh). + z : np.ndarray + Node elevation (m, positive up; negative below the reference level). + geographic : bool + True for lon/lat coordinates. + open_boundary : np.ndarray or None + ``(n, 2)`` polyline of the open (sea) boundary when the format has one + (ADCIRC), else None. + """ + + x: np.ndarray + y: np.ndarray + faces: np.ndarray + z: np.ndarray + geographic: bool = True + open_boundary: np.ndarray | None = None + + @property + def depth(self) -> np.ndarray: + """Node depth (m, positive down).""" + + return -self.z + + @property + def triangles(self) -> np.ndarray: + """``(n, 3)`` triangles: faces with more than three nodes split as fans.""" + + faces = np.asarray(self.faces, dtype=int) + tris = [] + for k in range(2, faces.shape[1]): + sel = faces[:, k] >= 0 + tris.append(faces[sel][:, [0, k - 1, k]]) + + return np.vstack(tris) + + def triangulation(self): + """Build a :class:`matplotlib.tri.Triangulation` of :attr:`triangles`.""" + + import matplotlib.tri as mtri + + return mtri.Triangulation(self.x, self.y, self.triangles) + + +def _looks_geographic(x: np.ndarray, y: np.ndarray) -> bool: + return bool(np.nanmax(np.abs(x)) <= 360 and np.nanmax(np.abs(y)) <= 90) + + +def _read_ugrid(path: Path) -> Mesh: + with xr.open_dataset(path) as ds: + faces_var = ds["mesh2d_face_nodes"] + start = int(faces_var.attrs.get("start_index", 0)) + values = np.asarray(faces_var.values, dtype=float) + missing = ~np.isfinite(values) | (values < start) + faces = np.where(missing, -1, values - start).astype(int) + x = ds["mesh2d_node_x"].values.astype(float) + y = ds["mesh2d_node_y"].values.astype(float) + z = ds["mesh2d_node_z"].values.astype(float) + units = str(ds["mesh2d_node_x"].attrs.get("units", "")) + geographic = "degree" in units or _looks_geographic(x, y) + + return Mesh(x=x, y=y, faces=faces, z=z, geographic=geographic) + + +def _read_adcirc(path: Path) -> Mesh: + from ..wrappers.swan.swan_utils import read_adcirc_grid, swan_boundary_side + + grid = read_adcirc_grid(str(path)) + v = grid["vertices"] + open_boundary = swan_boundary_side(grid)[0] if grid["open_boundaries"] else None + + return Mesh( + x=v[:, 0], + y=v[:, 1], + faces=np.asarray(grid["triangles"], dtype=int), + z=-v[:, 2], + geographic=_looks_geographic(v[:, 0], v[:, 1]), + open_boundary=open_boundary, + ) + + +def read_mesh(path: str | Path) -> Mesh: + """ + Read a model mesh: UGRID netCDF (``.nc``) or ADCIRC (``fort.14``, ``.grd``). + + Parameters + ---------- + path : str or Path + Mesh file. + + Returns + ------- + Mesh + Nodes, faces (0-based) and elevation (positive up). + + Raises + ------ + ValueError + If the format is not recognised. + """ + + path = Path(path) + if path.suffix == ".nc": + return _read_ugrid(path) + if path.suffix in (".grd", ".14") or path.name == "fort.14": + return _read_adcirc(path) + + raise ValueError(f"Unknown mesh format: {path.name} (expected .nc, fort.14, .grd)") + + +def depth_at_points(mesh: Mesh, x: np.ndarray, y: np.ndarray) -> np.ndarray: + """ + Mesh depth at points, linear inside the triangles. + + Parameters + ---------- + mesh : Mesh + Model mesh. + x, y : np.ndarray + Point coordinates, in the mesh coordinates. + + Returns + ------- + np.ndarray + Depth (m, positive down) at each point; NaN outside the mesh. + """ + + import matplotlib.tri as mtri + + interp = mtri.LinearTriInterpolator(mesh.triangulation(), mesh.depth) + x, y = np.asarray(x, dtype=float), np.asarray(y, dtype=float) + + return np.asarray(interp(x, y).filled(np.nan), dtype=float) + + +def _median_edge(mesh: Mesh) -> float: + tri = mesh.triangles + xy = np.c_[mesh.x, mesh.y] + + return float(np.median(np.hypot(*(xy[tri[:, 1]] - xy[tri[:, 0]]).T))) + + +def _raster( + values: np.ndarray, x: np.ndarray, y: np.ndarray, geographic: bool +) -> xr.DataArray: + xname, yname = ("lon", "lat") if geographic else ("x", "y") + da = xr.DataArray( + values, + dims=(yname, xname), + coords={yname: y, xname: x}, + name="elevation", + attrs={"units": "m", "positive": "up"}, + ) + if geographic: + try: + da = da.rio.set_spatial_dims(x_dim=xname, y_dim=yname).rio.write_crs( + "EPSG:4326" + ) + except AttributeError: # rioxarray not installed + pass + + return da + + +def mesh_to_raster( + mesh: Mesh, + resolution: float | None = None, + bounds: tuple[float, float, float, float] | None = None, +) -> xr.DataArray: + """ + Sample a mesh's elevation on a regular grid. + + Parameters + ---------- + mesh : Mesh + Model mesh. + resolution : float, optional + Cell size, in the mesh coordinates. Default: the median edge length + of the mesh. + bounds : tuple, optional + ``(xmin, ymin, xmax, ymax)``. Default: the mesh extent. + + Returns + ------- + xr.DataArray + Elevation (m, positive up) on ``(lat, lon)`` (geographic, with an + EPSG:4326 CRS when rioxarray is installed) or ``(y, x)``, north up; + NaN outside the mesh. ``.rio.to_raster(path)`` writes a GeoTIFF. + """ + + import matplotlib.tri as mtri + + if resolution is None: + resolution = _median_edge(mesh) + xmin, ymin, xmax, ymax = bounds or ( + mesh.x.min(), + mesh.y.min(), + mesh.x.max(), + mesh.y.max(), + ) + x = np.arange(xmin + resolution / 2, xmax, resolution) + y = np.arange(ymax - resolution / 2, ymin, -resolution) + interp = mtri.LinearTriInterpolator(mesh.triangulation(), mesh.z) + xx, yy = np.meshgrid(x, y) + values = np.asarray(interp(xx, yy).filled(np.nan), dtype=float) + + return _raster(values, x, y, mesh.geographic) + + +def mask_raster_to_mesh(da: xr.DataArray, mesh: Mesh) -> xr.DataArray: + """ + Set a raster to NaN outside a mesh. + + Contours traced on the result (and smoothing with + :func:`bluemath_tk.waves.hywaves.output_points.smooth_raster`, which + respects NaNs) then stay inside the model domain. + + Parameters + ---------- + da : xr.DataArray + 2-D raster on ``(y, x)``, ``(lat, lon)`` or ``(latitude, longitude)``, + in the mesh coordinates. + mesh : Mesh + Model mesh. + + Returns + ------- + xr.DataArray + A copy, NaN where the cell centre is outside the mesh. + """ + + from ..waves.hywaves.output_points import xy_coord_names + + xname, yname = xy_coord_names(da) + xx, yy = np.meshgrid(da[xname].values, da[yname].values) + inside = mesh.triangulation().get_trifinder()(xx, yy) >= 0 + if da.dims.index(xname) == 0: + inside = inside.T + + return da.where(inside) diff --git a/bluemath_tk/waves/hywaves/__init__.py b/bluemath_tk/waves/hywaves/__init__.py index 5adfeea..189c1d3 100644 --- a/bluemath_tk/waves/hywaves/__init__.py +++ b/bluemath_tk/waves/hywaves/__init__.py @@ -12,12 +12,13 @@ contributions. """ -from . import metamodel, offshore, output_points, reconstruction, sectors, wind_correction +from . import metamodel, offshore, output_points, plots, reconstruction, sectors, wind_correction __all__ = [ "metamodel", "offshore", "output_points", + "plots", "reconstruction", "sectors", "wind_correction", diff --git a/bluemath_tk/waves/hywaves/metamodel.py b/bluemath_tk/waves/hywaves/metamodel.py index 039edde..e84bc8d 100644 --- a/bluemath_tk/waves/hywaves/metamodel.py +++ b/bluemath_tk/waves/hywaves/metamodel.py @@ -5,8 +5,8 @@ (sampled Hs, Tp, Dir, Spr and WL), on ``(case_num, sites)``. The metamodel maps offshore forcing to nearshore fields, ``hs`` included, in absolute terms: -1. :func:`mask_dry_sites` -> :func:`detect_removed_sites` / :func:`drop_sites` - -> :func:`fill_dry_sites` - site filter and dry-case handling. +1. :func:`mask_zero_energy` -> :func:`detect_removed_sites` / :func:`drop_sites` + -> :func:`fill_zero_energy` - site filter and zero-energy cases. 2. :func:`fit_mda` - MDA centroids of the forcing are the training cases (the rest is held out). 3. :func:`fit_pcas` - one PCA per predicted variable over the site fields @@ -57,7 +57,7 @@ MIN_HS = 0.0 #: Default PCA NaN handling: fraction of valid cases a site needs to be kept, -#: and the fill value for the remaining NaNs. After :func:`fill_dry_sites` the +#: and the fill value for the remaining NaNs. After :func:`fill_zero_energy` the #: training fields have no NaN, so this is only a safety net. DEFAULT_PCA_PREPROCESSING: dict[str, dict[str, float]] = { "nan_threshold_to_drop": {"hs": 0.99, "dir_u": 0.99, "dir_v": 0.99, "spr": 0.99}, @@ -86,11 +86,11 @@ def detect_removed_sites( e.g. ``{"hs_ratio_max": 3}`` drops sites where ``hs`` exceeds three times the offshore Hs of a case (numerical anomalies). - ``{var}_nan_max``: bound on the fraction of cases where ``var`` is NaN, - e.g. ``{"hs_nan_max": 0.01}`` drops sites outside the model grid - (always NaN) and sites that are dry in more than 1 % of the cases. - Apply :func:`mask_dry_sites` first so that dry means NaN for every model. + e.g. ``{"hs_nan_max": 0.01}`` drops sites where the model failed in + more than 1 % of the cases (sites are inside the mesh and always wet, + so NaN ``hs`` is never a dry point; ``hs = 0`` is not NaN). - NaNs are skipped by the bound rules, so a dry case never trips them. + NaNs are skipped by the bound rules. Parameters ---------- @@ -152,75 +152,76 @@ def _site_fields(cases: xr.Dataset) -> list[str]: ] -def mask_dry_sites(cases: xr.Dataset, var: str = "hs") -> xr.Dataset: +def mask_zero_energy(cases: xr.Dataset, var: str = "hs") -> xr.Dataset: """ - Set every output field to NaN where a site is dry or outside the grid. + Blank the fields of site-cases where no wave energy arrives. - Models flag such points differently: SWAN writes its exception value (NaN - after reading) in every quantity, SnapWave writes ``hs = 0`` with - placeholder ``tp``/``dir`` and NaN ``spr``. Here a site-case is dry when - *var* is NaN or ``<= 0``, and then all its fields become NaN, so both - models look the same to :func:`detect_removed_sites` and - :func:`fill_dry_sites`. + Output sites are inside the mesh and always wet (see + :func:`bluemath_tk.waves.hywaves.output_points.site_status`), so + ``hs = 0`` is a real result: the goal sends no energy there in that case. + Models then write NaN (SWAN) or a placeholder (SnapWave) in the other + fields; here they all become NaN, ``hs`` stays 0 and NaN ``hs`` (a model + failure) stays NaN for :func:`detect_removed_sites`. Parameters ---------- cases : xr.Dataset Model cases on ``(case_num, sites)``. var : str, optional - Wave height variable that defines dry. Default is "hs". + Wave height variable. Default is "hs". Returns ------- xr.Dataset - A copy with dry site-cases set to NaN. + A copy, ``hs = 0`` (``hs <= 0`` clipped to 0) and the other fields NaN + where no energy arrives. """ - dry = cases[var].isnull() | (cases[var] <= 0) + zero = cases[var] <= 0 out = cases.copy() + out[var] = cases[var].where(~zero, 0.0) for name in _site_fields(cases): - out[name] = cases[name].where(~dry) + if name != var: + out[name] = cases[name].where(~zero) return out -def fill_dry_sites(cases: xr.Dataset, var: str = "hs") -> xr.Dataset: +def fill_zero_energy(cases: xr.Dataset, var: str = "hs") -> xr.Dataset: """ - Fill the dry site-cases left after the site filter, for PCA training. + Fill the fields of zero-energy site-cases for PCA training. - *var* becomes 0 (no water, no waves). Every other field takes the case - forcing (``{field}_forcing``, e.g. ``tp_forcing``): no value is physical - for a dry point, the forcing is a smooth function of the GP inputs (a - constant would add a jump the GP has to learn), and it does not matter in - the reconstruction, where a contribution with ``hs = 0`` adds no energy. + Where ``hs = 0`` every other field takes the case forcing + (``{field}_forcing``, e.g. ``tp_forcing``): no value is physical without + waves, the forcing is a smooth function of the GP inputs (a constant would + add a jump the GP has to learn), and it does not matter in the + reconstruction, where a contribution with ``hs = 0`` adds no energy. Fields without a forcing counterpart keep their NaNs (the PCA fills them). Parameters ---------- cases : xr.Dataset - Output of :func:`mask_dry_sites` after :func:`drop_sites`. + Output of :func:`mask_zero_energy` (after :func:`drop_sites`). var : str, optional Wave height variable. Default is "hs". Returns ------- xr.Dataset - A copy with the dry site-cases filled. + A copy with the zero-energy site-cases filled. """ - dry = cases[var].isnull() - if not bool(dry.any()): + zero = cases[var] == 0 + if not bool(zero.any()): return cases out = cases.copy() - out[var] = cases[var].fillna(0.0) for name in _site_fields(cases): if name != var and f"{name}_forcing" in cases: - out[name] = cases[name].where(~dry, cases[f"{name}_forcing"]) + out[name] = cases[name].where(~zero, cases[f"{name}_forcing"]) logger.info( - "Filled %d dry site-cases at %d sites (%s = 0, other fields = forcing)", - int(dry.sum()), - int(dry.any("case_num").sum()), - var, + "Filled %d zero-energy site-cases at %d sites (other fields = forcing)", + int(zero.sum()), + int(zero.any("case_num").sum()), ) return out @@ -539,7 +540,9 @@ def predict_fields( pcas : mapping Fitted PCAs by target name. forcing : pd.DataFrame - Forcing rows (columns as at fit time); the index labels the rows. + Forcing rows (columns as at fit time); the index labels the rows and + its name (``case_num`` if unnamed) is the row dimension of the output, + e.g. ``time`` for a hindcast. vars_to_predict : mapping Same specification as at fit time. sites : list of str @@ -548,11 +551,11 @@ def predict_fields( Returns ------- xr.Dataset - Predicted variables on ``(case_num, sites)``. + Predicted variables on ``(row, sites)``. """ - if forcing.index.name != "case_num": - forcing = forcing.rename_axis("case_num") + row_dim = forcing.index.name or "case_num" + forcing = forcing.rename_axis("case_num") rows = forcing.index.values blocks = split_pcs(gp.predict(dataset=forcing).values, pcas) @@ -591,7 +594,8 @@ def predict_fields( arr = arr.clip(min=MIN_HS) out[name] = arr.sel(sites=sites) - return xr.Dataset(out) + ds = xr.Dataset(out) + return ds if row_dim == "case_num" else ds.rename(case_num=row_dim) # --------------------------------------------------------------------------- diff --git a/bluemath_tk/waves/hywaves/output_points.py b/bluemath_tk/waves/hywaves/output_points.py index aecf122..8b6a507 100644 --- a/bluemath_tk/waves/hywaves/output_points.py +++ b/bluemath_tk/waves/hywaves/output_points.py @@ -2,8 +2,15 @@ HyWaves output points: where nearshore waves are computed and reconstructed. Typical sources are points along depth contours of a bathymetry raster -(:func:`get_isobath_points`) and unstructured-mesh nodes inside a polygon -(:func:`get_mesh_nodes_in_polygon`), plus any observation sites (buoys). +(:func:`get_isobath_points`) and model-mesh nodes inside a polygon (thinned +evenly with :func:`spatial_subsample`), plus any observation sites (buoys). +Every site is then checked against the model mesh (:func:`site_status`: inside +it and never dry; depths from :func:`bluemath_tk.topo_bathy.mesh.depth_at_points`). +Sites are decided once, before any run: a site is kept only if its mesh depth +plus the lowest water level of the cases stays above the models' minimum +depth. Every kept site is then wet in every case, so ``hs = 0`` at a site is a +result (the goal sends no energy there), not a dry or failed point; see +:func:`bluemath_tk.waves.hywaves.metamodel.mask_zero_energy`. Raw contours of a high-resolution bathymetry follow every small sandbank and pit, which puts points in odd places. :func:`get_isobath_points` can smooth @@ -15,8 +22,6 @@ from __future__ import annotations -from typing import Any - import matplotlib.pyplot as plt import numpy as np import shapely @@ -298,42 +303,118 @@ def get_isobath_points( return np.vstack(segments) if segments else np.empty((0, 2), dtype=float) -def get_mesh_nodes_in_polygon( - polygon: Polygon | MultiPolygon, - mesh: xr.Dataset, - wet_only: bool = False, -) -> list[dict[str, Any]]: +def spatial_subsample( + x: np.ndarray, + y: np.ndarray, + fraction: float, + seed: int = 0, + geographic: bool = True, +) -> np.ndarray: """ - Unstructured-mesh nodes inside a polygon. + Pick a random subset of points that still covers their whole area. + + Stratified sampling: the area is cut into equal square cells, sized so that + about ``fraction * len(x)`` cells hold points, and one point is drawn at + random from each of those cells. Dense clusters (e.g. a finer mesh near + the coast) are thinned more than sparse areas, so the subset spreads + evenly over the area instead of following the point density. Parameters ---------- - polygon : Polygon or MultiPolygon - Area, in the mesh coordinates. - mesh : xr.Dataset - UGRID mesh with ``mesh2d_node_x``, ``mesh2d_node_y``, ``mesh2d_node_z``. - wet_only : bool, optional - Keep only nodes with negative bed level. Default is False. + x, y : np.ndarray + Point coordinates (lon/lat degrees or projected metres). + fraction : float + Share of the points to keep, in ``(0, 1]``. The result has about + ``fraction * len(x)`` points (one per occupied cell). + seed : int, optional + Random seed. Default is 0. + geographic : bool, optional + Lon/lat input (cells are square in metres). Default is True. Returns ------- - list[dict] - One ``{"lon", "lat", "depth", "node_index"}`` dict per node. + np.ndarray + Sorted indices of the kept points. + + Raises + ------ + ValueError + If *fraction* is not in ``(0, 1]``. """ - lon = mesh["mesh2d_node_x"].values - lat = mesh["mesh2d_node_y"].values - z = mesh["mesh2d_node_z"].values - mask = shapely.contains_xy(polygon, lon, lat) - if wet_only: - mask &= z < 0.0 - - return [ - { - "lon": float(lon[i]), - "lat": float(lat[i]), - "depth": float(z[i]), - "node_index": int(i), - } - for i in np.flatnonzero(mask) - ] + if not 0 < fraction <= 1: + raise ValueError(f"fraction must be in (0, 1], got {fraction}") + x, y = np.asarray(x, dtype=float), np.asarray(y, dtype=float) + n = x.size + target = max(1, int(round(fraction * n))) + if target >= n: + return np.arange(n) + if geographic: + lat0 = np.radians(np.mean(y)) + x = np.radians(x) * EARTH_RADIUS_M * np.cos(lat0) + y = np.radians(y) * EARTH_RADIUS_M + + def cells(size: float) -> np.ndarray: + ij = np.floor(np.c_[x - x.min(), y - y.min()] / size).astype(np.int64) + return np.unique(ij, axis=0, return_inverse=True)[1].ravel() + + lo, hi = 1e-9, float(max(np.ptp(x), np.ptp(y))) + 1.0 + for _ in range(60): # occupied cells decrease with the cell size + mid = np.sqrt(lo * hi) + if cells(mid).max() + 1 > target: + lo = mid + else: + hi = mid + cell = cells( + lo + if abs(cells(lo).max() + 1 - target) < abs(cells(hi).max() + 1 - target) + else hi + ) + + rng = np.random.default_rng(seed) + order = rng.permutation(n) + _, first = np.unique(cell[order], return_index=True) + + return np.sort(order[first]) + + +#: :func:`site_status` labels. +SITE_OK = "ok" +SITE_OUTSIDE = "outside_mesh" +SITE_CAN_DRY = "can_dry" + + +def site_status( + depth: np.ndarray, wl_min: float = 0.0, min_depth: float = 0.1 +) -> np.ndarray: + """ + Whether each output site can be used: inside the mesh and always wet. + + A site is kept only if the model mesh covers it and its water depth stays + above *min_depth* at the lowest water level the cases can have, so every + run computes it and ``hs = 0`` there always means "no energy arrives". + + Parameters + ---------- + depth : np.ndarray + Mesh depth at the sites (m, positive down; NaN outside the mesh, see + :func:`bluemath_tk.topo_bathy.mesh.depth_at_points`). + wl_min : float, optional + Lowest water level of the cases (m, e.g. the lower ``wl`` bound of the + metamodel). Default is 0. + min_depth : float, optional + Minimum depth the model computes (SnapWave ``hmin``, SWAN ``DEPMIN``). + Default is 0.1. + + Returns + ------- + np.ndarray + :data:`SITE_OK`, :data:`SITE_OUTSIDE` or :data:`SITE_CAN_DRY` per site. + """ + + depth = np.asarray(depth, dtype=float) + status = np.full(depth.shape, SITE_OK, dtype=object) + status[depth + float(wl_min) < float(min_depth)] = SITE_CAN_DRY + status[~np.isfinite(depth)] = SITE_OUTSIDE + + return status diff --git a/bluemath_tk/waves/hywaves/plots.py b/bluemath_tk/waves/hywaves/plots.py new file mode 100644 index 0000000..a995c18 --- /dev/null +++ b/bluemath_tk/waves/hywaves/plots.py @@ -0,0 +1,535 @@ +""" +Project: BlueMath_tk +Sub-Module: waves.hywaves.plots +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) + +Figures of the HyWaves method: + +- goals: colours (:func:`goal_colors`), sectors with the sampled or observed + wave directions on a map (:func:`plot_goal_sectors`) and per goal + (:func:`plot_direction_ranges`); +- the met-ocean forcing of dynamic runs (:func:`plot_dynamic_forcing`); +- the wind correction's K (:func:`plot_wind_k_comparison`); +- the MDA-size sweep of the holdout metrics (:func:`plot_mda_sweep`). + +Goals are mappings ``{goal_id: {"lon", "lat", "angles": [left, right]}}``, +directions nautical, coming-from, clockwise. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import matplotlib.colors as mcolors +import matplotlib.dates as mdates +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from matplotlib.axes import Axes +from matplotlib.figure import Figure +from matplotlib.gridspec import GridSpec +from matplotlib.patches import Wedge + +from ...core.metrics import METRIC_LABELS +from ...core.operations import nautical_sector_to_wedge_angles +from ...core.plotting.labels import place_labels + +Goals = Mapping[str, Mapping[str, Any]] + + +def goal_colors(goal_ids: Sequence[Any]) -> dict[Any, tuple]: + """ + One colour per goal from viridis, by rank (never repeats, unlike tab10). + + Goals are numbered along the domain boundary, so the colour is also a + position cue. + """ + + positions = np.linspace(0.08, 0.92, len(goal_ids)) if len(goal_ids) > 1 else [0.5] + cmap = plt.get_cmap("viridis") + return {gid: cmap(pos) for gid, pos in zip(goal_ids, positions)} + + +def plot_goal_sectors( + ax: Axes, + goals: Goals, + colors: Mapping[int, Any], + *, + directions: pd.DataFrame | None = None, + radius: float = 0.35, + arrow_length: float = 0.15, + max_arrows: int = 25, +) -> None: + """ + Goal sectors as wedges, goal stars with non-overlapping labels, wave arrows. + + Parameters + ---------- + ax : Axes + Map axes with its final limits (labels are placed on the rendered figure). + goals : mapping + ``{goal_id: {"lon", "lat", "angles"}}``. + colors : mapping + Colour by integer goal id (see :func:`goal_colors`). + directions : pd.DataFrame, optional + ``goal`` and ``dir`` columns: wave-from directions drawn as arrows + into each goal (a random subsample of *max_arrows*). + radius : float, optional + Wedge radius (degrees). + arrow_length : float, optional + Arrow length (degrees). + max_arrows : int, optional + Arrows per goal. + """ + + keys = sorted(goals, key=int) + crowded = len(keys) > 15 + for key in keys: + goal = goals[key] + color = colors[int(key)] + theta1, theta2 = nautical_sector_to_wedge_angles(*goal["angles"]) + ax.add_patch(Wedge((goal["lon"], goal["lat"]), radius, theta1, theta2, facecolor=color, + edgecolor=color, alpha=0.15, linewidth=0.8, zorder=4)) + ax.plot(goal["lon"], goal["lat"], marker="*", color=color, markersize=9 if crowded else 12, zorder=6) + if directions is None: + continue + sub = directions[directions["goal"] == int(key)] + if len(sub) > max_arrows: + sub = sub.sample(max_arrows, random_state=0) + theta = np.deg2rad(sub["dir"].to_numpy(float)) + for dx, dy in zip(np.sin(theta) * arrow_length, np.cos(theta) * arrow_length): + ax.annotate( + "", xy=(goal["lon"], goal["lat"]), xytext=(goal["lon"] + dx, goal["lat"] + dy), + arrowprops=dict(arrowstyle="-|>", color=color, lw=0.9, alpha=0.55, shrinkA=0, + shrinkB=2, mutation_scale=8), + zorder=5, + ) + place_labels( + ax, + [goals[k]["lon"] for k in keys], + [goals[k]["lat"] for k in keys], + [f"G{k}" for k in keys], + leader_colors=[colors[int(k)] for k in keys], + fontsize=7.5 if crowded else 9, + weight="bold", + zorder=7, + path_effects=[pe.withStroke(linewidth=2.2, foreground="white")], + ) + + +def _sector_spans(left: float, right: float) -> list[tuple[float, float]]: + """One or two ``(start, end)`` spans on [0, 360) of a nautical sector.""" + left_n, right_n = float(np.mod(left, 360.0)), float(np.mod(right, 360.0)) + if left < 0 or left_n > right_n: + return [(left_n, 360.0), (0.0, right_n)] + return [(left_n, right_n)] + + +def _circular_spans(dirs: np.ndarray) -> list[tuple[float, float]]: + """One or two spans on [0, 360) tightly bounding *dirs* (jumps the largest gap).""" + vals = np.sort(np.mod(np.asarray(dirs, dtype=float), 360.0)) + if vals.size <= 1: + v = float(vals[0]) if vals.size else 0.0 + return [(v, v)] + gaps = np.append(np.diff(vals), vals[0] + 360.0 - vals[-1]) + i = int(np.argmax(gaps)) + if i == len(vals) - 1: + return [(float(vals[0]), float(vals[-1]))] + start, end = float(vals[i + 1]), float(vals[i]) + return [(start, 360.0), (0.0, end)] + + +def plot_direction_ranges( + ax: Axes, + cases: pd.DataFrame, + goals: Goals, + colors: Mapping[int, Any], + *, + bracket: str = "observed", + max_marks: int = 200, + title: str = "Wave directions per goal", +) -> bool: + """ + One row per goal: its sector (light bar) and the directions of *cases*. + + Up to *max_marks* cases per goal every direction is a tick with a solid + range bar; above that the row is a density strip (alpha = density). + + Parameters + ---------- + ax : Axes + Target axes. + cases : pd.DataFrame + ``goal``, ``dir`` and ``tp`` columns (one row per case or time step). + goals : mapping + ``{goal_id: {"angles"}}``. + colors : mapping + Colour by integer goal id. + bracket : {"observed", "sector"}, optional + What the crisp outline bounds: the observed interval (designed cases) + or the configured sector (hindcast, whose directions may spill out). + max_marks : int, optional + Cases per goal above which rows become density strips. + title : str, optional + Axes title. + + Returns + ------- + bool + Whether the density-strip style was used (for a legend text). + """ + + goal_ids = sorted(cases["goal"].unique()) + dense = bool((cases.groupby("goal").size() > max_marks).any()) + for y, gid in enumerate(goal_ids): + color = colors[gid] + sub = cases[cases["goal"] == gid] + dirs = sub["dir"].to_numpy(float) + sector = _sector_spans(*goals[str(gid)]["angles"]) + for x0, x1 in sector: + ax.barh(y, x1 - x0, left=x0, height=0.55, color=color, alpha=0.18, edgecolor=color, linewidth=0.6) + spans = sector if bracket == "sector" else _circular_spans(dirs) + if dense: + counts, _ = np.histogram(np.mod(dirs, 360.0), bins=120, range=(0, 360)) + rgba = np.zeros((1, 120, 4)) + rgba[0, :, :3] = mcolors.to_rgb(color) + rgba[0, :, 3] = counts / max(counts.max(), 1) * 0.85 + ax.imshow(rgba, extent=(0, 360, y - 0.31, y + 0.31), aspect="auto", + interpolation="nearest", zorder=2) + edge = tuple(c * 0.55 for c in mcolors.to_rgb(color)) + for x0, x1 in spans: + ax.barh(y, x1 - x0, left=x0, height=0.62, fill=False, edgecolor=edge, linewidth=1.4, zorder=3) + else: + for x0, x1 in spans: + ax.barh(y, x1 - x0, left=x0, height=0.28, color=color, alpha=0.85, edgecolor="white", linewidth=0.5) + ax.scatter(dirs, np.full(len(dirs), y), marker="|", s=90, color=color, linewidths=1.2, zorder=4) + ax.text(362, y, f"tp {sub['tp'].min():.1f}-{sub['tp'].max():.1f}s", va="center", fontsize=8, color="#333333") + ax.set_yticks(range(len(goal_ids)), [f"G{g}" for g in goal_ids]) + ax.set_xlim(0, 395) + ax.set_xticks(np.arange(0, 361, 45)) + ax.set_xlabel("wave-from direction (° nautical)") + ax.set_title(title) + ax.grid(axis="x", alpha=0.25) + ax.invert_yaxis() + return dense + +#: Metrics drawn by :func:`plot_mda_sweep`, one row each. +SWEEP_METRICS: tuple[str, ...] = ("bias", "rmse", "rrmse", "r2") + + +def plot_mda_sweep( + table: pd.DataFrame, + series: Sequence[str], + variables: Sequence[str], + *, + metrics: Sequence[str] = SWEEP_METRICS, + units: Mapping[str, str] | None = None, + colors: Mapping[str, str] | None = None, + title: str | None = None, + figsize: tuple[float, float] | None = None, +) -> Figure: + """ + Holdout metrics against the MDA size: one panel per metric and variable. + + Parameters + ---------- + table : pd.DataFrame + One row per series, ``mda_candidate`` and ``variable``, with the + *metrics* columns (e.g. pooled with + :func:`bluemath_tk.core.metrics.pool_metrics`). + series : sequence of str + Columns whose values make one line; the legend joins them. + variables : sequence of str + Variables, one column of panels each. + metrics : sequence of str, optional + Metrics, one row of panels each (those without values are skipped). + units : mapping, optional + Units per variable for the axis labels. + colors : mapping, optional + Line colour by the last *series* value (e.g. a group name); others + follow the default colour cycle. + title : str, optional + Figure title. + figsize : tuple, optional + Default scales with the panel grid. + + Returns + ------- + Figure + The panel grid, with one shared legend. + """ + + units, colors = units or {}, colors or {} + metrics = [m for m in metrics if m in table and table[m].notna().any()] + nrows, ncols = len(metrics), len(variables) + fig, axes = plt.subplots( + nrows, + ncols, + figsize=figsize or (4.2 * ncols, 3.2 * nrows), + sharex=True, + squeeze=False, + ) + for j, var in enumerate(variables): + sub = table[table["variable"] == var] + for key, line in sub.groupby(list(series), sort=False): + key = key if isinstance(key, tuple) else (key,) + line = line.sort_values("mda_candidate") + for i, metric in enumerate(metrics): + axes[i, j].plot( + line["mda_candidate"], + line[metric], + "o-", + lw=1.5, + ms=4, + color=colors.get(str(key[-1])), + label=" / ".join(str(k) for k in key), + ) + for i, metric in enumerate(metrics): + ax = axes[i, j] + unit = "" if metric in ("rrmse", "r2", "r", "si") else f" [{units.get(var, '')}]" + ax.set_ylabel(f"{METRIC_LABELS.get(metric, metric)}{unit}") + if metric == "r2": + ax.set_ylim(0, 1.05) + ax.grid(True, alpha=0.3) + axes[0, j].set_title(var) + axes[-1, j].set_xlabel("MDA train cases") + + handles, labels = axes[0, 0].get_legend_handles_labels() + if handles: + fig.legend( + handles, + labels, + loc="upper center", + ncol=min(4, len(labels)), + bbox_to_anchor=(0.5, 1.04), + fontsize=8, + ) + if title: + fig.suptitle(title, fontsize=12, y=1.08 if handles else 1.0) + fig.tight_layout() + return fig + + + +def plot_dynamic_forcing( + aoi: Any, + goals: Goals, + offshore: xr.Dataset, + *, + wind: xr.Dataset | None = None, + water_level: xr.DataArray | None = None, + angle_range: tuple[float, float] | None = None, + wind_name: str | None = None, + water_level_name: str | None = None, + title: str = "Dynamic met-ocean forcing", +) -> Figure: + """ + Wind and water level given to dynamic runs, next to the offshore waves. + + Top: mean wind over the period (shading: mean speed; arrows: mean + vector, short where the direction keeps changing) and the wind at the + storm hour (highest goal-mean Hs) with the goals' wave directions. + Bottom: wind rose of the domain-mean wind; wind-wave angle at the goals + for hours with wind sea, with *angle_range* shaded; water-level envelope + over the goals. + + Parameters + ---------- + aoi : geopandas.GeoDataFrame + Area of interest (outline and extent). + goals : mapping + ``{goal_id: {"lon", "lat"}}``. + offshore : xr.Dataset + ``hs``, ``tp``, ``dir``, ``spr`` on ``(time, goal)``. + wind : xr.Dataset, optional + ``u10``/``v10`` on ``(time, latitude, longitude)``. + water_level : xr.DataArray, optional + Water level on ``(time, goal)``. + angle_range : tuple of float, optional + Wind-wave angles (deg) a wind correction was trained on. + wind_name, water_level_name : str, optional + Source names for the titles. + title : str, optional + Figure title. + + Returns + ------- + Figure + """ + + from ...core.plotting.wind import plot_wind_arrows, plot_wind_rose + from ..wind import wind_sea_weight, wind_speed_direction + + offshore = offshore.transpose("time", "goal", ...) + times = pd.DatetimeIndex(offshore["time"].values) + goal_ids = [int(g) for g in offshore["goal"].values] + goal_xy = np.array([[goals[str(g)]["lon"], goals[str(g)]["lat"]] for g in goal_ids]) + storm = times[int(np.nanargmax(offshore["hs"].mean("goal").values))] + + fig = plt.figure(figsize=(16, 10.5), constrained_layout=True) + gs = GridSpec(2, 6, figure=fig, height_ratios=[1.25, 1.0]) + ax_mean, ax_storm = fig.add_subplot(gs[0, :3]), fig.add_subplot(gs[0, 3:]) + ax_rose = fig.add_subplot(gs[1, :2], projection="polar") + ax_angle, ax_wl = fig.add_subplot(gs[1, 2:4]), fig.add_subplot(gs[1, 4:]) + fig.suptitle(title, fontsize=14, fontweight="bold") + + aoi = aoi.to_crs(4326) + x0, y0, x1, y1 = aoi.total_bounds + + def frame(ax: Axes, panel_title: str) -> None: + aoi.boundary.plot(ax=ax, color="black", lw=1.0, zorder=4) + ax.scatter(*goal_xy.T, s=25, c="white", edgecolors="black", zorder=7) + for g, (x, y) in zip(goal_ids, goal_xy): + ax.annotate(f"G{g}", (x, y), xytext=(3, 3), textcoords="offset points", fontsize=7, zorder=8) + ax.set(xlim=(x0 - 0.15, x1 + 0.15), ylim=(y0 - 0.15, y1 + 0.15)) + ax.set_aspect("equal", adjustable="box") + ax.set_title(panel_title, fontsize=11) + ax.tick_params(labelsize=8) + + def empty(ax: Axes, text: str) -> None: + ax.text(0.5, 0.5, text, ha="center", va="center", transform=ax.transAxes) + ax.set_axis_off() + + if wind is not None: + wind = wind.sel(time=slice(times.min(), times.max())) + lon2, lat2 = np.meshgrid(wind["longitude"].values, wind["latitude"].values) + speed = np.hypot(wind["u10"], wind["v10"]) + vmax = max(float(np.nanmax(speed.max("time").values)), 1.0) + mean_speed = speed.mean("time").values + s_mean, d_mean = wind_speed_direction(wind["u10"].mean("time").values, wind["v10"].mean("time").values) + mesh = ax_mean.pcolormesh(lon2, lat2, mean_speed, cmap="YlGnBu", vmin=0, vmax=vmax * 0.6, shading="auto") + plt.colorbar(mesh, ax=ax_mean, fraction=0.03, label="mean u10 (m/s)") + plot_wind_arrows(ax_mean, lon2, lat2, s_mean, d_mean, max_arrows=400, key_speed=5.0, scale=80.0, width=0.003) + frame(ax_mean, f"Mean wind {times.min():%Y-%m} → {times.max():%Y-%m} " + f"(arrows: mean vector; steadiness {np.nanmean(s_mean) / np.nanmean(mean_speed):.2f})") + + at = wind.sel(time=storm, method="nearest") + s_storm, d_storm = wind_speed_direction(at["u10"].values, at["v10"].values) + mesh = ax_storm.pcolormesh(lon2, lat2, s_storm, cmap="YlOrRd", vmin=0, vmax=vmax, shading="auto") + plt.colorbar(mesh, ax=ax_storm, fraction=0.03, label="u10 (m/s)") + plot_wind_arrows(ax_storm, lon2, lat2, s_storm, d_storm, max_arrows=400, scale=280.0, width=0.003) + off = offshore.sel(time=storm) + for (x, y), d in zip(goal_xy, off["dir"].values): + th = np.deg2rad(d) + ax_storm.annotate("", xy=(x, y), xytext=(x + 0.25 * np.sin(th), y + 0.25 * np.cos(th)), + arrowprops=dict(arrowstyle="-|>", color="#08306b", lw=2, mutation_scale=14), zorder=8) + frame(ax_storm, f"Storm hour {storm:%Y-%m-%d %H:%M} (max goal-mean Hs " + f"{float(off['hs'].mean()):.1f} m; blue = wave dir)") + + dm_speed, dm_dir = wind_speed_direction( + wind["u10"].mean(["latitude", "longitude"]).values, wind["v10"].mean(["latitude", "longitude"]).values + ) + plot_wind_rose(ax_rose, dm_speed, dm_dir) + ax_rose.set_title(f"Domain-mean wind ({wind_name or 'u10'})", fontsize=11, pad=12) + ax_rose.legend(fontsize=7, loc="upper left", bbox_to_anchor=(1.02, 1.0), frameon=False) + + at_goals = wind.interp( + latitude=xr.DataArray(goal_xy[:, 1], dims="goal"), + longitude=xr.DataArray(goal_xy[:, 0], dims="goal"), + ).reindex(time=times, method="nearest", tolerance=pd.Timedelta("1h")) + w_speed, w_dir = wind_speed_direction(at_goals["u10"].values, at_goals["v10"].values) + weight = wind_sea_weight(offshore["hs"].values, offshore["tp"].values, offshore["spr"].values) + angle = np.mod(w_dir - offshore["dir"].values + 180.0, 360.0) - 180.0 + sel = (weight > 0) & np.isfinite(angle) & (w_speed > 1) + ax_angle.hist(angle[sel], bins=np.arange(-180, 181, 10), color="#4c72b0", edgecolor="white") + panel = "Wind − wave direction, wind-sea hours" + if angle_range is not None: + ax_angle.axvspan(*angle_range, color="#55a868", alpha=0.18, label="wind-correction training range") + outside = np.mean((angle[sel] < angle_range[0]) | (angle[sel] > angle_range[1])) * 100 if sel.any() else 0 + ax_angle.legend(fontsize=8, loc="upper left") + panel += f" ({outside:.0f}% outside)" + ax_angle.set(title=panel, xlabel="angle (deg; + = wind clockwise of waves)", ylabel="goal-hours", xlim=(-180, 180)) + else: + for ax in (ax_mean, ax_storm, ax_angle, ax_rose): + empty(ax, "no wind forcing") + + if water_level is not None: + wl = water_level.sel(goal=goal_ids).reindex(time=times).load() + ax_wl.fill_between(times, wl.min("goal"), wl.max("goal"), color="#9ecae1", label="goal range", lw=0) + ax_wl.plot(times, wl.mean("goal"), color="#08519c", lw=0.6, label="goal mean") + missing = wl.isnull().all("goal").values + edges = np.flatnonzero(np.diff(np.r_[0, missing.astype(int), 0])) + for k, (a, b) in enumerate(zip(edges[::2], edges[1::2])): + ax_wl.axvspan(times[a], times[b - 1], color="#d62728", alpha=0.12, label="no data (wl = 0)" if k == 0 else None) + ax_wl.set(title=f"Water level ({water_level_name or 'per goal'})", ylabel="m") + ax_wl.legend(fontsize=8, loc="upper right") + locator = mdates.AutoDateLocator(maxticks=6) + ax_wl.xaxis.set_major_locator(locator) + ax_wl.xaxis.set_major_formatter(mdates.ConciseDateFormatter(locator)) + else: + empty(ax_wl, "no water-level forcing") + return fig + + +def plot_wind_k_comparison(df: pd.DataFrame, title: str = "How K is estimated") -> Figure: + """ + The wind correction's ``K``: fitted model and growth formulas against the runs. + + Left: each estimate against the runs' ``K`` (reference, x-axis), with + RMSE and bias. Right: median ``K`` against how developed the wind sea + already is (``g Hs_ws / U^2``; 0.24 = fully developed), in bins of at + least 30 site-cases. + + Parameters + ---------- + df : pd.DataFrame + ``k_runs``, ``k_fitted``, ``k_formula_deep``, ``k_formula_depth`` and + ``hs_ws_u2`` per site-case. + title : str, optional + Figure title. + + Returns + ------- + Figure + """ + + estimates = { + "k_fitted": ("fitted model", "#1f77b4"), + "k_formula_deep": ("formula, deep water", "#ff7f0e"), + "k_formula_depth": ("formula, depth-limited", "#2ca02c"), + } + fig, (ax_sc, ax_bin) = plt.subplots(1, 2, figsize=(14, 6), constrained_layout=True) + fig.suptitle(title, fontsize=14, fontweight="bold") + ref = df["k_runs"].to_numpy() + top = float(np.nanpercentile(ref, 99.5)) + for col, (label, color) in estimates.items(): + est = df[col].to_numpy() + ok = np.isfinite(ref) & np.isfinite(est) + err = est[ok] - ref[ok] + ax_sc.scatter(ref[ok], est[ok], s=4, alpha=0.25, color=color, rasterized=True, + label=f"{label}: RMSE {np.sqrt(np.mean(err**2)):.3f}, bias {np.mean(err):+.3f}") + ax_sc.plot([1, top], [1, top], "k--", lw=1) + ax_sc.set(xlim=(0.9, top), ylim=(0.9, top), xlabel="K from the runs (with / without wind)", + ylabel="K estimated", title=f"Held-out cases · {len(df):,} site-cases") + ax_sc.legend(fontsize=9, markerscale=4, loc="upper left") + + bins = np.geomspace(max(float(df["hs_ws_u2"].min()), 1e-3), float(df["hs_ws_u2"].max()), 16) + mid = np.sqrt(bins[:-1] * bins[1:]) + grouped = df.groupby(pd.cut(df["hs_ws_u2"], bins), observed=False) + enough = (grouped.size() >= 30).to_numpy() + medians = grouped[["k_runs", *estimates]].median()[enough] + ax_bin.plot(mid[enough], medians["k_runs"], "k-o", ms=4, label="runs (reference)") + for col, (label, color) in estimates.items(): + ax_bin.plot(mid[enough], medians[col], "-o", ms=3, color=color, label=label) + ax_bin.axvline(0.24, color="0.5", ls=":", lw=1) + ax_bin.text(0.24, ax_bin.get_ylim()[1], " fully developed", va="top", fontsize=8, color="0.4") + ax_bin.set(xscale="log", xlabel="g Hs_ws / U² (how developed the wind sea already is)", + ylabel="median K", title="K vs development of the wind sea") + ax_bin.legend(fontsize=9) + return fig + + +__all__ = [ + "SWEEP_METRICS", + "goal_colors", + "plot_direction_ranges", + "plot_dynamic_forcing", + "plot_goal_sectors", + "plot_mda_sweep", + "plot_wind_k_comparison", +] diff --git a/bluemath_tk/waves/hywaves/reconstruction.py b/bluemath_tk/waves/hywaves/reconstruction.py index d8c433d..ecd44f4 100644 --- a/bluemath_tk/waves/hywaves/reconstruction.py +++ b/bluemath_tk/waves/hywaves/reconstruction.py @@ -10,6 +10,16 @@ contribution, the local wind sea, is added per site batch before the sum (see :mod:`.wind_correction`). +The sum reproduces a model run forced by every goal at once only if each +metamodel case imposes a fraction ``w(goal, point)`` of its goal's *energy* +along the open boundary and the fractions of all goals add up to 1. Both +models do: SnapWave mixes the spectra of the two nearest goals at every +boundary node with inverse-distance weights; SWAN (``BOUNDSPEC ... VARIABLE +PAR``) interpolates the spectra linearly along the boundary, so the energy is +linear and Hs falls as ``sqrt(w)`` between goals, and the boundary ends take +their nearest goal in full (see +:func:`bluemath_tk.wrappers.swan.swan_utils.boundary_par_table`). + The spectral construction allocates the full ``(time, sites, goal x partition, freq, dir)`` array, so long periods or many sites are processed in site batches sized to a memory budget (:func:`site_batch_size`), optionally across @@ -235,7 +245,7 @@ def _sum_site_batches( for k, i in enumerate(range(0, len(site_ids), batch), start=1): t_batch = time.monotonic() sites = site_ids[i : i + batch] - merged = xr.merge([c.sel(sites=sites) for c in chunks]) + merged = xr.merge([c.sel(sites=sites) for c in chunks], join="outer", compat="no_conflicts") wind_k = None if wind_correction is not None: merged, wind_k = _add_wind_sea(merged, wind_correction) diff --git a/bluemath_tk/waves/hywaves/wind_correction.py b/bluemath_tk/waves/hywaves/wind_correction.py index 5ce2bb5..59b5dc3 100644 --- a/bluemath_tk/waves/hywaves/wind_correction.py +++ b/bluemath_tk/waves/hywaves/wind_correction.py @@ -22,6 +22,27 @@ The fetch at each site and direction comes from :func:`bluemath_tk.waves.fetch.fetch_table`. + +Why not run the metamodel cases with wind: each case is one goal's +contribution and the reconstruction sums them, so every case would grow its +own copy of the local wind sea (counted once per goal); and growth slows as +the sea develops, which depends on the total energy, not on one goal's share. + +``K`` is learned from the model, not from a growth formula (which saturates +at ``K = 1`` for storm seas where the models still add 15-35 %). Its inputs +are dimensionless with the wind speed (lengths ``U^2/g``), so ~100 training +cases cover every site and wind: + +- ``g Hs_ws / U^2``: how developed the wind sea is (0.24 = fully developed); + the strongest effect, young seas grow most; +- ``g F / U^2``: fetch left upwind inside the domain; +- ``g d / U^2`` and ``g d_mean / U^2``: depth at the site and along the + fetch (shallow water limits growth); +- ``cos(wind - wave)``: following winds grow most. + +Limits: training winds are uniform; wind-wave angles beyond the trained range +(e.g. +-60 deg) keep the edge value (tree models do not extrapolate); in +SnapWave (one period per node) wind energy over swell takes the swell period. """ from __future__ import annotations diff --git a/tests/core/test_core_metrics.py b/tests/core/test_core_metrics.py index 4aa069c..bdf7bec 100644 --- a/tests/core/test_core_metrics.py +++ b/tests/core/test_core_metrics.py @@ -212,5 +212,28 @@ def test_unknown_metric(self): compute_metrics([1.0], [1.0], ["nse"]) +class TestBoundsAndPooling(unittest.TestCase): + def test_bounds_drop_pairs_outside(self): + out = compute_metrics([1.0, 2.0, 3.0], [1.0, 2.0, 9e13], ["rmse"], bounds=(0.0, 15.0)) + self.assertEqual(out["n"], 2) + self.assertEqual(out["rmse"], 0.0) + + def test_pool_metrics(self): + import pandas as pd + + from bluemath_tk.core.metrics import pool_metrics + + table = pd.DataFrame( + {"set": ["a", "a", "b"], "n": [1, 3, 2], "bias": [0.0, 1.0, 2.0], + "rmse": [1.0, 3.0, 2.0], "rrmse": [0.1, 0.3, 0.2], "r2": [0.5, 0.9, 0.8]} + ) + out = pool_metrics(table, ["set"]).set_index("set") + self.assertAlmostEqual(out.loc["a", "bias"], 0.75) + self.assertAlmostEqual(out.loc["a", "rmse"], np.sqrt((1 + 27) / 4)) + self.assertAlmostEqual(out.loc["a", "rrmse"], 0.25) + self.assertEqual(out.loc["a", "n"], 4) + self.assertEqual(out["n"].dtype.kind, "i") + + if __name__ == "__main__": unittest.main() diff --git a/tests/core/test_models_logging.py b/tests/core/test_models_logging.py new file mode 100644 index 0000000..d51e752 --- /dev/null +++ b/tests/core/test_models_logging.py @@ -0,0 +1,49 @@ +import contextlib +import io +import logging +import unittest + +from bluemath_tk.core import models +from bluemath_tk.core.models import BlueMathModel, set_models_console_level + + +class _DummyModel(BlueMathModel): + def __init__(self, logger: logging.Logger | None = None) -> None: + super().__init__() + if logger is not None: + self._logger = logger + + +class TestModelsConsoleLevel(unittest.TestCase): + def setUp(self): + self.addCleanup(set_models_console_level, models.MODELS_CONSOLE_LEVEL) + set_models_console_level("ERROR") + + def test_new_model_is_quiet_on_console(self): + with contextlib.redirect_stderr(io.StringIO()) as buf: + _DummyModel().logger.warning("routine message") + self.assertEqual(buf.getvalue(), "") + + def test_unpickled_logger_is_repaired(self): + # a model loaded from disk refers to its logger by name; in a new + # process that logger has no handlers and propagates to the root + bare = logging.getLogger("_BareUnpickledModel") + bare.handlers = [] + bare.propagate = True + model = _DummyModel(bare) + with contextlib.redirect_stderr(io.StringIO()) as buf: + model.logger.warning("must not reach the console") + self.assertEqual(buf.getvalue(), "") + self.assertFalse(model.logger.propagate) + self.assertEqual(model.logger.name, "_BareUnpickledModel") + + def test_existing_model_loggers_follow_the_level(self): + logger = _DummyModel().logger + set_models_console_level("CRITICAL") + consoles = [h for h in logger.handlers if not isinstance(h, logging.FileHandler)] + self.assertTrue(consoles) + self.assertTrue(all(h.level == logging.CRITICAL for h in consoles)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/topo_bathy/test_mesh.py b/tests/topo_bathy/test_mesh.py new file mode 100644 index 0000000..97f08da --- /dev/null +++ b/tests/topo_bathy/test_mesh.py @@ -0,0 +1,215 @@ +"""Tests for bluemath_tk.topo_bathy.mesh and the output-site gate.""" + +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import xarray as xr + +from bluemath_tk.topo_bathy.mesh import ( + Mesh, + depth_at_points, + mask_raster_to_mesh, + mesh_to_raster, + read_mesh, +) +from bluemath_tk.waves.hywaves.output_points import ( + SITE_CAN_DRY, + SITE_OK, + SITE_OUTSIDE, + site_status, + spatial_subsample, +) + + +def _sloping_grid(n: int = 11) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Nodes of [0, 1]^2 (degrees), index grid, elevation -20 x (deeper east).""" + + x, y = np.meshgrid(np.linspace(0, 1, n), np.linspace(0, 1, n)) + idx = np.arange(n * n).reshape(n, n) + return x.ravel(), y.ravel(), idx, -20.0 * x.ravel() + + +def _quad_mesh() -> Mesh: + """Quads on the west half, triangles on the east half, padded with -1.""" + + x, y, idx, z = _sloping_grid() + a, b = idx[:-1, :-1], idx[:-1, 1:] + c, d = idx[1:, :-1], idx[1:, 1:] + west = (slice(None), slice(0, 5)) + east = (slice(None), slice(5, None)) + quads = np.c_[a[west].ravel(), b[west].ravel(), d[west].ravel(), c[west].ravel()] + tris = np.r_[ + np.c_[a[east].ravel(), b[east].ravel(), d[east].ravel()], + np.c_[a[east].ravel(), d[east].ravel(), c[east].ravel()], + ] + faces = np.r_[quads, np.c_[tris, np.full(len(tris), -1)]] + return Mesh(x=x, y=y, faces=faces, z=z) + + +class TestMesh(unittest.TestCase): + """Depth at points and rasters from a mixed quad/triangle mesh.""" + + def setUp(self): + """Build the test data.""" + + self.mesh = _quad_mesh() + + def test_triangles(self): + """Each quad becomes two triangles; padded triangles stay one.""" + + self.assertEqual(self.mesh.triangles.shape, (50 * 2 + 100, 3)) + self.assertGreaterEqual(self.mesh.triangles.min(), 0) + + def test_depth_at_points(self): + """Linear depth inside the mesh (positive down), NaN outside.""" + + depth = depth_at_points(self.mesh, [0.25, 0.75, 1.5], [0.5, 0.5, 0.5]) + np.testing.assert_allclose(depth[:2], [5.0, 15.0]) + self.assertTrue(np.isnan(depth[2])) + + def test_mesh_to_raster(self): + """Raster of the mesh elevation, north up, NaN outside the mesh.""" + + da = mesh_to_raster(self.mesh, resolution=0.05, bounds=(-0.5, 0.0, 1.0, 1.0)) + self.assertEqual(da.dims, ("lat", "lon")) + self.assertTrue(bool((np.diff(da.lat) < 0).all())) + row = da.isel(lat=10) + self.assertTrue(bool(row.sel(lon=slice(None, 0)).isnull().all())) + inside = row.sel(lon=slice(0, None)) + np.testing.assert_allclose(inside, -20.0 * inside.lon, atol=1e-9) + + def test_mask_raster_to_mesh(self): + """A larger raster keeps its values only inside the mesh.""" + + lon = np.linspace(-0.5, 1.5, 41) + lat = np.linspace(1.0, 0.0, 21) + da = xr.DataArray( + np.full((21, 41), -5.0), + coords={"lat": lat, "lon": lon}, + dims=("lat", "lon"), + ) + masked = mask_raster_to_mesh(da, self.mesh) + inside = (da.lon >= 0) & (da.lon <= 1) + self.assertTrue(bool(masked.where(~inside).isnull().all())) + self.assertTrue( + bool((masked.isel(lat=slice(1, -1)).where(inside, -5.0) == -5.0).all()) + ) + + +class TestReadMesh(unittest.TestCase): + """UGRID (with fill values) and ADCIRC give the same mesh.""" + + def setUp(self): + """Build the test data.""" + + self.tmp = tempfile.TemporaryDirectory() + self.dir = Path(self.tmp.name) + self.mesh = _quad_mesh() + + def tearDown(self): + """Remove the temporary folder.""" + + self.tmp.cleanup() + + def test_ugrid_quadtree(self): + """1-based face nodes with fill values, elevation positive up.""" + + faces = np.where(self.mesh.faces >= 0, self.mesh.faces + 1, -999) + ds = xr.Dataset( + { + "mesh2d_node_x": ("n", self.mesh.x, {"units": "degrees_east"}), + "mesh2d_node_y": ("n", self.mesh.y), + "mesh2d_node_z": ("n", self.mesh.z), + "mesh2d_face_nodes": ( + ("f", "m"), + faces, + {"start_index": 1, "_FillValue": -999}, + ), + } + ) + path = self.dir / "mesh.nc" + ds.to_netcdf(path) + mesh = read_mesh(path) + np.testing.assert_array_equal(mesh.faces, self.mesh.faces) + self.assertTrue(mesh.geographic) + np.testing.assert_allclose(depth_at_points(mesh, [0.5], [0.5]), [10.0]) + + def test_adcirc(self): + """ADCIRC depth is positive down in the file.""" + + tri = self.mesh.triangles + lines = ["EPSG:4326", f"{len(tri)} {len(self.mesh.x)}"] + lines += [ + f"{i + 1} {x} {y} {-z}" + for i, (x, y, z) in enumerate(zip(self.mesh.x, self.mesh.y, self.mesh.z)) + ] + lines += [ + f"{i + 1} 3 {a + 1} {b + 1} {c + 1}" for i, (a, b, c) in enumerate(tri) + ] + lines += ["0", "0", "0", "0"] + path = self.dir / "fort.14" + path.write_text("\n".join(lines) + "\n") + mesh = read_mesh(path) + self.assertIsNone(mesh.open_boundary) + np.testing.assert_allclose(mesh.z, self.mesh.z) + np.testing.assert_allclose(depth_at_points(mesh, [0.5], [0.5]), [10.0]) + + def test_unknown_format(self): + """Other extensions are rejected.""" + + with self.assertRaises(ValueError): + read_mesh(self.dir / "mesh.msh") + + +class TestSiteStatus(unittest.TestCase): + """Sites must be inside the mesh and wet at the lowest water level.""" + + def test_status(self): + """Deep, dry-at-low-water, just-wet and outside sites.""" + + depth = np.array([10.0, 3.0, 3.7, np.nan]) + status = site_status(depth, wl_min=-3.5, min_depth=0.1) + self.assertEqual( + status.tolist(), [SITE_OK, SITE_CAN_DRY, SITE_OK, SITE_OUTSIDE] + ) + self.assertEqual(site_status([0.05])[0], SITE_CAN_DRY) + + +class TestSpatialSubsample(unittest.TestCase): + """A fraction of the points, spread over the whole area.""" + + def setUp(self): + """Build the test data.""" + + rng = np.random.default_rng(1) + # dense cluster in one corner of a sparse square (metres) + dense = rng.uniform(0, 1000, (4000, 2)) + sparse = rng.uniform(0, 10000, (1000, 2)) + self.xy = np.r_[dense, sparse] + + def test_size_and_coverage(self): + """About the requested size, and spread over the area.""" + + idx = spatial_subsample(*self.xy.T, fraction=0.05, geographic=False) + self.assertLess(abs(len(idx) - 250), 25) + kept = self.xy[idx] + in_dense = (kept < 1000).all(axis=1).mean() + self.assertLess(in_dense, 0.2) # 80 % of the input, a few % of the area + quadrants = np.unique((kept // 5000).astype(int), axis=0) + self.assertEqual(len(quadrants), 4) + + def test_reproducible(self): + """Same seed, same subset; fraction bounds are checked.""" + + a = spatial_subsample(*self.xy.T, fraction=0.1, seed=3, geographic=False) + b = spatial_subsample(*self.xy.T, fraction=0.1, seed=3, geographic=False) + np.testing.assert_array_equal(a, b) + self.assertEqual(len(spatial_subsample(*self.xy.T, fraction=1.0)), len(self.xy)) + with self.assertRaises(ValueError): + spatial_subsample(*self.xy.T, fraction=0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/waves/test_hywaves.py b/tests/waves/test_hywaves.py index 262584a..7c1edf4 100644 --- a/tests/waves/test_hywaves.py +++ b/tests/waves/test_hywaves.py @@ -180,23 +180,6 @@ def test_smooth_raster_keeps_nans(self): smooth[:, 100], da[:, 100], atol=0.05 ) # linear field - def test_mesh_nodes_in_polygon(self): - """Only (wet) nodes inside the polygon are returned.""" - - mesh = xr.Dataset( - { - "mesh2d_node_x": ("n", [0.5, 0.6, 2.0]), - "mesh2d_node_y": ("n", [0.5, 0.6, 2.0]), - "mesh2d_node_z": ("n", [-3.0, 1.0, -5.0]), - } - ) - nodes = output_points.get_mesh_nodes_in_polygon(box(0, 0, 1, 1), mesh) - self.assertEqual([n["node_index"] for n in nodes], [0, 1]) - wet = output_points.get_mesh_nodes_in_polygon( - box(0, 0, 1, 1), mesh, wet_only=True - ) - self.assertEqual([n["node_index"] for n in wet], [0]) - def _synthetic_cases(n_cases: int = 60, n_sites: int = 5, seed: int = 0) -> xr.Dataset: """SnapWave-like cases: hs/dir/spr fields that depend smoothly on the forcing. @@ -215,7 +198,7 @@ def _synthetic_cases(n_cases: int = 60, n_sites: int = 5, seed: int = 0) -> xr.D hs = gain[None, :] * hs_off[:, None] * (0.5 + tp[:, None] / 32) hs = np.clip(hs + 0.05 * wl[:, None], 0, None) hs[:, -1] = np.minimum(hs[:, -1], 1.5) - hs[:, 0] = np.nan # a dry site, dropped by the PCA NaN threshold + hs[:, 0] = np.nan # a site the model never computes, dropped by the PCA sites = [f"s{i}" for i in range(n_sites)] return xr.Dataset( { @@ -267,25 +250,25 @@ def test_site_filter(self): bad = metamodel.detect_removed_sites(cases, {"hs_ratio_max": 3.0}) self.assertEqual(bad, ["s2"]) - def test_dry_sites(self): - """Dry site-cases are masked, filtered by fraction and filled.""" + def test_zero_energy(self): + """Zero hs is kept as a result; only NaN hs counts as a model failure.""" cases = self.cases.copy(deep=True) - cases["hs"][:3, 1] = 0.0 # SnapWave-style dry: hs = 0, placeholder dir + cases["hs"][:3, 1] = 0.0 # no energy: SnapWave placeholder dir cases["dir"][:3, 1] = 270.0 - cases["spr"][:3, 1] = np.nan - cases["hs"][:30, 2] = np.nan # SWAN-style dry in half of the cases + cases["spr"][:3, 1] = np.nan # SWAN-style NaN + cases["hs"][:30, 2] = np.nan # model failure in half of the cases - masked = metamodel.mask_dry_sites(cases) - for var in ("hs", "dir", "spr"): + masked = metamodel.mask_zero_energy(cases) + np.testing.assert_array_equal(masked.hs.sel(sites="s1")[:3], 0.0) + for var in ("dir", "spr"): self.assertTrue(bool(masked[var][:3, 1].isnull().all())) self.assertFalse(bool(masked.hs_forcing.isnull().any())) - bad = metamodel.detect_removed_sites(masked, {"hs_nan_max": 0.1}) - self.assertEqual(bad, ["s0", "s2"]) # outside the grid, dry 50 % + bad = metamodel.detect_removed_sites(masked, {"hs_nan_max": 0.01}) + self.assertEqual(bad, ["s0", "s2"]) # s1 (zero energy) is kept - filled = metamodel.fill_dry_sites(metamodel.drop_sites(masked, bad)) - self.assertFalse(bool(filled.hs.isnull().any())) + filled = metamodel.fill_zero_energy(metamodel.drop_sites(masked, bad)) np.testing.assert_array_equal(filled.hs.sel(sites="s1")[:3], 0.0) np.testing.assert_allclose( filled.dir.sel(sites="s1")[:3], cases.dir_forcing[:3] @@ -320,6 +303,14 @@ def test_fit_and_predict(self): true_hs = self.cases.hs.isel(case_num=test_idx).sel(sites=sites) self.assertLess(float(np.abs(pred.hs - true_hs.values).max()), 0.1) + hindcast = test_forcing.set_index( + pd.date_range("2020-01-01", periods=len(test_idx), freq="h").rename("time") + ) + by_time = metamodel.predict_fields(gp, pcas, hindcast, VARS, sites) + self.assertEqual(by_time.hs.dims, ("time", "sites")) + np.testing.assert_allclose(by_time.hs.values, pred.hs.values) + self.assertIs(gp.to("cpu"), gp) + def test_goal_forcing_and_valid_steps(self): """Partition forcing is renamed; empty partitions and off-sector steps drop.""" diff --git a/tests/waves/test_hywaves_plots.py b/tests/waves/test_hywaves_plots.py new file mode 100644 index 0000000..02260ad --- /dev/null +++ b/tests/waves/test_hywaves_plots.py @@ -0,0 +1,88 @@ +import unittest + +import matplotlib + +matplotlib.use("Agg") +import geopandas as gpd +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from shapely.geometry import box + +from bluemath_tk.waves.hywaves import plots + +GOALS = { + "1": {"lon": 0.0, "lat": 0.0, "angles": [-45.0, 45.0]}, + "2": {"lon": 0.5, "lat": 0.0, "angles": [180.0, 300.0]}, +} + + +class TestHyWavesPlots(unittest.TestCase): + def tearDown(self): + plt.close("all") + + def test_goal_colors_never_repeat(self): + colors = plots.goal_colors(list(range(30))) + self.assertEqual(len(set(colors.values())), 30) + + def test_goal_sectors_and_direction_ranges(self): + rng = np.random.default_rng(0) + cases = pd.DataFrame( + { + "goal": np.repeat([1, 2], 50), + "dir": np.r_[rng.uniform(-45, 45, 50) % 360, rng.uniform(180, 300, 50)], + "tp": rng.uniform(5, 12, 100), + } + ) + colors = plots.goal_colors([1, 2]) + fig, (ax_map, ax_bars) = plt.subplots(2, 1) + ax_map.set(xlim=(-1, 1.5), ylim=(-1, 1)) + plots.plot_goal_sectors(ax_map, GOALS, colors, directions=cases, max_arrows=5) + self.assertEqual(len(ax_map.patches), 2) + self.assertEqual({t.get_text() for t in ax_map.texts if t.get_text()}, {"G1", "G2"}) + self.assertFalse(plots.plot_direction_ranges(ax_bars, cases, GOALS, colors)) + self.assertTrue(plots.plot_direction_ranges(ax_bars, cases, GOALS, colors, max_marks=10)) + + def test_circular_spans_jump_the_gap(self): + self.assertEqual(plots._circular_spans(np.array([350.0, 355.0, 5.0, 10.0])), [(350.0, 360.0), (0.0, 10.0)]) + self.assertEqual(plots._sector_spans(-45.0, 45.0), [(315.0, 360.0), (0.0, 45.0)]) + + def test_mda_sweep(self): + table = pd.DataFrame( + { + "metamodel": ["buoys"] * 4, + "mda_candidate": [50, 100, 50, 100], + "variable": ["hs", "hs", "dir", "dir"], + "bias": [0.1, 0.05, 1.0, 0.5], + "rmse": [0.3, 0.2, 10.0, 8.0], + "rrmse": [0.2, 0.1, np.nan, np.nan], + "r2": [0.9, 0.95, 0.8, 0.85], + } + ) + fig = plots.plot_mda_sweep(table, ["metamodel"], ["hs", "dir"], units={"hs": "m", "dir": "deg"}) + self.assertEqual(len(fig.axes), 8) + + def test_dynamic_forcing_without_wind_any_dim_order(self): + times = pd.date_range("2020-01-01", periods=24, freq="h") + offshore = xr.Dataset( + {v: (("goal", "time"), np.full((2, 24), x)) for v, x in (("hs", 2.0), ("tp", 8.0), ("dir", 270.0), ("spr", 30.0))}, + coords={"goal": [1, 2], "time": times}, + ) + level = xr.DataArray(np.zeros((24, 2)), dims=("time", "goal"), coords={"time": times, "goal": [1, 2]}) + aoi = gpd.GeoDataFrame(geometry=[box(-1, -1, 1.5, 1)], crs="EPSG:4326") + fig = plots.plot_dynamic_forcing(aoi, GOALS, offshore, water_level=level) + self.assertEqual(len(fig.axes), 5) + + def test_wind_k_comparison(self): + rng = np.random.default_rng(1) + k = 1 + rng.uniform(0, 0.5, 200) + df = pd.DataFrame( + {"k_runs": k, "k_fitted": k + 0.01, "k_formula_deep": k * 1.1, + "k_formula_depth": k * 1.05, "hs_ws_u2": rng.uniform(0.01, 0.3, 200)} + ) + self.assertEqual(len(plots.plot_wind_k_comparison(df).axes), 2) + + +if __name__ == "__main__": + unittest.main()