From e280155884b4214f7a45ec4da0264196c5962bdb Mon Sep 17 00:00:00 2001 From: tausiaj Date: Wed, 4 Nov 2026 17:53:35 +0100 Subject: [PATCH 1/2] [JTH] several features for wind and swan --- bluemath_tk/topo_bathy/mesh.py | 130 ++++++- bluemath_tk/waves/calibration.py | 329 ++++++++++++++++++ bluemath_tk/waves/hywaves/offshore.py | 100 +++++- bluemath_tk/waves/hywaves/plots.py | 294 ++++++++++++++++ bluemath_tk/waves/hywaves/store.py | 2 +- bluemath_tk/waves/hywaves/wind_correction.py | 112 ++++-- bluemath_tk/wrappers/snapwave/__init__.py | 10 +- .../wrappers/snapwave/snapwave_utils.py | 120 +++---- .../wrappers/snapwave/snapwave_wrapper.py | 46 ++- bluemath_tk/wrappers/swan/swan_plots.py | 2 +- bluemath_tk/wrappers/swan/swan_utils.py | 1 + bluemath_tk/wrappers/swan/swan_wrapper.py | 31 +- tests/data/swan/boundary/INPUT | 2 + tests/topo_bathy/test_mesh.py | 34 ++ tests/waves/test_calibration.py | 74 +++- tests/waves/test_hywaves_offshore.py | 45 +++ tests/waves/test_hywaves_plots.py | 47 +++ tests/waves/test_wind_correction.py | 14 + tests/wrappers/test_snapwave_wrapper.py | 44 ++- tests/wrappers/test_swan_boundary_wrapper.py | 13 +- 20 files changed, 1320 insertions(+), 130 deletions(-) diff --git a/bluemath_tk/topo_bathy/mesh.py b/bluemath_tk/topo_bathy/mesh.py index ceff94b..5cd7b03 100644 --- a/bluemath_tk/topo_bathy/mesh.py +++ b/bluemath_tk/topo_bathy/mesh.py @@ -10,7 +10,8 @@ - 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 +triangles (NaN outside the mesh), :func:`open_boundary_nodes` finds the open +(sea) part of the mesh outline from points on it, :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. """ @@ -144,6 +145,133 @@ def read_mesh(path: str | Path) -> Mesh: raise ValueError(f"Unknown mesh format: {path.name} (expected .nc, fort.14, .grd)") +def outline_rings(mesh: Mesh) -> list[np.ndarray]: + """ + Closed outlines of the mesh: node indices in order along each ring. + + Parameters + ---------- + mesh : Mesh + Model mesh. + + Returns + ------- + list of np.ndarray + One array of node indices per ring (outer boundary and islands). + """ + + tri = mesh.triangles + edges = np.sort(np.vstack([tri[:, [0, 1]], tri[:, [1, 2]], tri[:, [2, 0]]]), axis=1) + unique, counts = np.unique(edges, axis=0, return_counts=True) + neighbours: dict[int, list[int]] = {} + for a, b in unique[counts == 1]: + neighbours.setdefault(int(a), []).append(int(b)) + neighbours.setdefault(int(b), []).append(int(a)) + + rings, seen = [], set() + for start in neighbours: + if start in seen: + continue + ring, previous, current = [start], None, start + seen.add(start) + while True: + following = next(n for n in neighbours[current] if n != previous) + if following == start or following in seen: + break + ring.append(following) + seen.add(following) + previous, current = current, following + rings.append(np.asarray(ring)) + + return rings + + +def _turn_angle(xy: np.ndarray, before: int, node: int, after: int) -> float: + """Absolute change of direction (deg) at *node* walking before -> node -> after.""" + + v1, v2 = xy[node] - xy[before], xy[after] - xy[node] + cross = v1[0] * v2[1] - v1[1] * v2[0] + + return abs(float(np.degrees(np.arctan2(cross, v1 @ v2)))) + + +def open_boundary_nodes( + mesh: Mesh, points: np.ndarray, max_turn: float = 45.0 +) -> np.ndarray: + """ + Open (sea) boundary of a mesh: the outline through *points*, out to the coast. + + The outline arc from the first to the last point (passing every point, in + order) is extended at both ends along the outline over two corners + sharper than *max_turn*: past the first (the end of the offshore boundary, + where the lateral boundary starts) and up to the second (where the + lateral boundary meets the coast). So the first and last points must sit + on the offshore boundary, at most at its corners (e.g. HyWaves goals). + + Parameters + ---------- + mesh : Mesh + Model mesh. + points : np.ndarray + ``(n, 2)`` points on the open boundary, in order along it. + max_turn : float, optional + Change of direction (deg) that ends the boundary. Default is 45. + + Returns + ------- + np.ndarray + Node indices of the open boundary, in order along it. + + Raises + ------ + ValueError + If the points are not in order along one outline. + """ + + points = np.asarray(points, dtype=float).reshape(-1, 2) + scale = np.cos(np.radians(np.nanmean(mesh.y))) if mesh.geographic else 1.0 + nodes_xy = np.column_stack([mesh.x * scale, mesh.y]) + points_xy = points * [scale, 1.0] + + rings = outline_rings(mesh) + ring_of = np.concatenate([np.full(len(r), k) for k, r in enumerate(rings)]) + place = np.concatenate([np.arange(len(r)) for r in rings]) + outline_xy = nodes_xy[np.concatenate(rings)] + nearest = [int(np.argmin(((outline_xy - p) ** 2).sum(1))) for p in points_xy] + if len(set(ring_of[nearest])) != 1: + raise ValueError("the points are on different mesh outlines") + ring = rings[ring_of[nearest[0]]] + xy = nodes_xy[ring] + n = len(ring) + position = [int(place[i]) for i in nearest] + + step = 1 + if [(p - position[0]) % n for p in position] != sorted((p - position[0]) % n for p in position): + step = -1 + if [(position[0] - p) % n for p in position] != sorted((position[0] - p) % n for p in position): + raise ValueError("the points are not in order along the mesh outline") + arc = [position[0]] + while arc[-1] != position[-1]: + arc.append((arc[-1] + step) % n) + + def extend(end: int, inner: int, direction: int) -> list[int]: + added, previous, current, corners = [], inner, end, 0 + while len(arc) + len(added) < n: + following = (current + direction) % n + if previous != current and _turn_angle(xy, previous, current, following) > max_turn: + corners += 1 + if corners == 2: + break + added.append(following) + previous, current = current, following + return added + + head = extend(arc[0], arc[1] if len(arc) > 1 else arc[0], -step) + tail = extend(arc[-1], arc[-2] if len(arc) > 1 else arc[-1], step) + + return ring[np.asarray(head[::-1] + arc + tail)] + + def depth_at_points(mesh: Mesh, x: np.ndarray, y: np.ndarray) -> np.ndarray: """ Mesh depth at points, linear inside the triangles. diff --git a/bluemath_tk/waves/calibration.py b/bluemath_tk/waves/calibration.py index 79659e5..d6ad1ca 100644 --- a/bluemath_tk/waves/calibration.py +++ b/bluemath_tk/waves/calibration.py @@ -26,6 +26,13 @@ The pairs (e.g. from :func:`bluemath_tk.waves.altimetry.collocate`) are split into a training part (``train_fraction``, 75 % by default) used for the fit and a test part used only to validate it (:meth:`CalVal.metrics`). + +A constant factor per sector corrects the level of a hindcast, not a storm tail +that is too short. :class:`QuantileMapping` is the second step: it maps the Hs +quantiles of the directionally calibrated hindcast onto those of the reference, +over the pairs of a whole area (a tail needs many storms: altimeter passes +around one point hold too few), and keeps the correction constant beyond +``tail_quantile``, where the quantiles rest on a handful of pairs. """ from __future__ import annotations @@ -44,6 +51,14 @@ #: Metrics reported by :meth:`CalVal.metrics` (reference first: bias = hindcast - reference). CALIBRATION_METRICS: tuple[str, ...] = ("bias", "rmse", "si", "r2", "slope") +#: Default labels and colours of the :class:`QuantileMapping` stages. +MAPPING_STAGES: dict[str, tuple[str, str]] = { + "reference": ("reference", "black"), + "raw": ("raw", "0.6"), + "hindcast": ("hindcast", "tab:blue"), + "mapped": ("mapped", "tab:orange"), +} + def sector_index(direction: np.ndarray, sector_size: float) -> np.ndarray: """ @@ -449,3 +464,317 @@ def plot(self, title: str | None = None) -> Figure: fig.tight_layout() return fig + + +class QuantileMapping(BlueMathModel): + """ + Monotonic map of hindcast Hs onto the Hs distribution of a reference. + + Each pair of samples gives the quantiles of both at the same probabilities: + below the first one the correction is a factor, between them linear, and + above ``tail_quantile`` a constant offset. + + Parameters + ---------- + tail_quantile : float, optional + Last quantile mapped. Default is 0.999. + n_tail : int, optional + Quantiles between 0.99 and ``tail_quantile``. Default is 12. + min_tail_samples : int, optional + Pairs needed above ``tail_quantile``. Default is 5. + hs_prefix : str, optional + Partition Hs prefix. Default "phs". + + Attributes + ---------- + table : pd.DataFrame + ``probability``, ``hindcast`` and ``reference`` quantiles. + + Notes + ----- + The map aligns the two distributions, not each pair: the mapped Hs of a + pair can move away from its reference even when the quantiles agree + (:meth:`plot`, last panel). The samples given to :meth:`fit` are kept for + :meth:`quantiles` and :meth:`plot` but not saved with the model. + """ + + def __init__( + self, + tail_quantile: float = 0.999, + n_tail: int = 12, + min_tail_samples: int = 5, + hs_prefix: str = "phs", + ) -> None: + super().__init__() + self.set_logger_name(name="QuantileMapping", level="INFO", console=True) + if not 0.99 <= tail_quantile < 1: + raise ValueError(f"tail_quantile must be in [0.99, 1), got {tail_quantile}") + self.tail_quantile = float(tail_quantile) + self.n_tail = int(n_tail) + self.min_tail_samples = int(min_tail_samples) + self.hs_prefix = hs_prefix + self.table: pd.DataFrame | None = None + self._samples: pd.DataFrame | None = None + self._exclude_attributes += ["_samples"] + + def fit(self, hindcast: np.ndarray, reference: np.ndarray) -> QuantileMapping: + """ + Fit the map on paired Hs samples. + + Parameters + ---------- + hindcast, reference : np.ndarray + Hs of the hindcast (already directionally calibrated) and of the + reference, same shape; pairs with a NaN are dropped. + + Returns + ------- + QuantileMapping + The fitted instance. + """ + + hindcast, reference = np.asarray(hindcast, float), np.asarray(reference, float) + self._samples = pd.DataFrame({"hindcast": hindcast, "reference": reference}) + finite = np.isfinite(hindcast) & np.isfinite(reference) + hindcast, reference = hindcast[finite], reference[finite] + tail_samples = hindcast.size * (1.0 - self.tail_quantile) + if tail_samples < self.min_tail_samples: + raise ValueError( + f"{hindcast.size} pairs leave {tail_samples:.1f} above quantile " + f"{self.tail_quantile}; {self.min_tail_samples} are needed" + ) + probability = np.unique( + np.r_[np.linspace(0.01, 0.99, 99), np.linspace(0.99, self.tail_quantile, self.n_tail)] + ) + self.table = pd.DataFrame( + { + "probability": probability, + "hindcast": np.maximum.accumulate(np.quantile(hindcast, probability)), + "reference": np.maximum.accumulate(np.quantile(reference, probability)), + } + ) + last = self.table.iloc[-1] + self.logger.info( + "Fitted on %d pairs: %+.2f m above %.2f m", + hindcast.size, + last["reference"] - last["hindcast"], + last["hindcast"], + ) + + return self + + def map(self, hs: np.ndarray) -> np.ndarray: + """Mapped Hs (NaN stays NaN).""" + + if self.table is None: + raise ValueError("Run fit() first") + hs = np.asarray(hs, float) + low, high = self.table["hindcast"].to_numpy(), self.table["reference"].to_numpy() + mapped = np.interp(hs, low, high) + mapped = np.where(hs < low[0], hs * high[0] / max(low[0], 1e-6), mapped) + mapped = np.where(hs > low[-1], hs + high[-1] - low[-1], mapped) + + return np.where(np.isnan(hs), np.nan, mapped) + + def correct(self, data: pd.DataFrame | xr.Dataset) -> pd.DataFrame | xr.Dataset: + """ + Map the bulk ``hs`` and scale the partition Hs by the same ratio. + + Parameters + ---------- + data : pd.DataFrame or xr.Dataset + Hindcast with ``hs`` and any partition Hs. + + Returns + ------- + pd.DataFrame or xr.Dataset + A copy with corrected Hs; other variables unchanged. + """ + + hs = np.asarray(data["hs"], float) + ratio = np.divide(self.map(hs), hs, out=np.ones_like(hs), where=hs > 0) + factor = ratio if isinstance(data, pd.DataFrame) else data["hs"].copy(data=ratio) + out = data.copy() + names = data.columns if isinstance(data, pd.DataFrame) else data.data_vars + for name in names: + if name == "hs" or str(name).startswith(self.hs_prefix): + out[name] = data[name] * factor + + return out + + # ------------------------------------------------------------------ + # Diagnostics + # ------------------------------------------------------------------ + + def evaluated_samples(self, raw: np.ndarray | None = None) -> pd.DataFrame: + """ + The fitted pairs with their mapped Hs. + + Parameters + ---------- + raw : np.ndarray, optional + Hs before any earlier correction, aligned with the samples given + to :meth:`fit` (e.g. before the directional step). + + Returns + ------- + pd.DataFrame + ``reference``, ``raw`` (if given), ``hindcast`` and ``mapped``, + pairs with a NaN dropped. + """ + + if self._samples is None: + raise ValueError("Run fit() first (the samples are not stored after loading)") + samples = self._samples.copy() + if raw is not None: + raw = np.asarray(raw, float) + if raw.shape != (len(samples),): + raise ValueError(f"raw has shape {raw.shape}; the fit had {len(samples)} pairs") + samples["raw"] = raw + samples["mapped"] = self.map(samples["hindcast"].to_numpy()) + columns = [c for c in ("reference", "raw", "hindcast", "mapped") if c in samples] + + return samples[columns].dropna().reset_index(drop=True) + + def quantiles( + self, + probabilities: tuple[float, ...] = (0.5, 0.9, 0.99, 0.999), + raw: np.ndarray | None = None, + ) -> pd.DataFrame: + """ + Hs quantiles of the reference and of each stage, on the fitted pairs. + + Parameters + ---------- + probabilities : tuple of float, optional + Non-exceedance probabilities. + raw : np.ndarray, optional + See :meth:`evaluated_samples`. + + Returns + ------- + pd.DataFrame + Index ``probability``; one column per stage and ``n_above``, the + pairs above the quantile (what the value rests on). + """ + + samples = self.evaluated_samples(raw) + table = samples.quantile(list(probabilities)) + table.index.name = "probability" + table["n_above"] = np.round(len(samples) * (1.0 - table.index.to_numpy())).astype(int) + + return table + + def plot( + self, + raw: np.ndarray | None = None, + labels: dict[str, str] | None = None, + title: str | None = None, + ) -> Figure: + """ + What the map does: the correction, the quantiles, the tail and the pairs. + + Panels: (a) correction (mapped - hindcast) against hindcast Hs, with the + fitted quantiles, the three regions (factor below the first quantile, + interpolated, constant offset beyond ``tail_quantile``), how many pairs + each holds and their histogram; (b) Q-Q against the reference; (c) + exceedance probability; (d) mapped Hs of each pair against its + reference. + + Parameters + ---------- + raw : np.ndarray, optional + See :meth:`evaluated_samples`; adds the raw stage to (b) and (c). + labels : dict, optional + Stage labels (keys ``reference``, ``raw``, ``hindcast``, + ``mapped``), e.g. ``{"reference": "satellite"}``. + title : str, optional + Figure title. + + Returns + ------- + Figure + The 2 x 2 panels. + """ + + samples = self.evaluated_samples(raw) + names = {key: (labels or {}).get(key, label) for key, (label, _) in MAPPING_STAGES.items()} + colors = {key: color for key, (_, color) in MAPPING_STAGES.items()} + stages = [c for c in ("raw", "hindcast", "mapped") if c in samples] + low, high = self.table["hindcast"].to_numpy(), self.table["reference"].to_numpy() + hindcast = samples["hindcast"].to_numpy() + n = len(samples) + fig, axes = plt.subplots(2, 2, figsize=(13, 10.5)) + + ax = axes[0, 0] + x_max = max(hindcast.max(), low[-1]) * 1.1 + x = np.linspace(0.0, x_max, 500) + ax_hist = ax.twinx() + ax_hist.hist(hindcast, bins=60, range=(0, x_max), color="0.85", log=True) + ax_hist.set_ylabel("pairs per bin (log)", color="0.5") + ax.set_zorder(ax_hist.get_zorder() + 1) + ax.patch.set_visible(False) + regions = ( + (0.0, low[0], "factor", "0.6"), + (low[0], low[-1], "interpolated", "tab:blue"), + (low[-1], x_max, f"offset {high[-1] - low[-1]:+.2f} m", "tab:red"), + ) + for k, (start, end, label, color) in enumerate(regions): + ax.axvspan(start, end, color=color, alpha=0.07) + inside = int(((hindcast >= start) & (hindcast < end)).sum()) + ax.text(0.01 * x_max if k == 0 else (start + end) / 2, 0.97, f"{label}\n{inside} pairs", + transform=ax.get_xaxis_transform(), ha="left" if k == 0 else "center", va="top", + fontsize=8, color=color) + ax.axhline(0.0, color="k", lw=0.8) + ax.plot(x, self.map(x) - x, color=colors["mapped"], lw=2, label="mapped - hindcast") + ax.plot(low, high - low, "o", ms=3, color="k", label="fitted quantiles") + bottom, top = ax.get_ylim() + ax.set_ylim(bottom, top + 0.3 * (top - bottom)) + ax.set(xlim=(0, x_max), xlabel=f"{names['hindcast']} Hs (m)", ylabel="correction (m)", + title=f"(a) the map · tail at p{100 * self.tail_quantile:g} = {low[-1]:.2f} m") + ax.legend(loc="lower right", fontsize=8) + + ax = axes[0, 1] + probability = np.unique(np.r_[np.linspace(0.01, 0.99, 99), 1.0 - np.logspace(-2, np.log10(1.0 / n), 40)]) + reference_q = np.quantile(samples["reference"], probability) + for stage in stages: + ax.plot(reference_q, np.quantile(samples[stage], probability), ".-", ms=3, lw=1, + color=colors[stage], label=names[stage]) + limit = max(reference_q.max(), samples[stages].quantile(probability[-1]).max()) * 1.05 + ax.plot([0, limit], [0, limit], "k--", lw=0.8) + tail = np.quantile(samples["reference"], self.tail_quantile) + ax.axvline(tail, color="tab:red", ls=":", lw=1) + ax.text(tail, 0.02, f" p{100 * self.tail_quantile:g}", transform=ax.get_xaxis_transform(), + color="tab:red", fontsize=8) + ax.set(xlim=(0, limit), ylim=(0, limit), aspect="equal", + xlabel=f"{names['reference']} Hs quantile (m)", ylabel="Hs quantile (m)", + title="(b) Q-Q: on the 1:1 line the distributions agree") + ax.legend(fontsize=8) + + ax = axes[1, 0] + exceedance = 1.0 - np.arange(1, n + 1) / (n + 1) + for stage in ("reference", *stages): + ax.semilogy(np.sort(samples[stage].to_numpy()), exceedance, lw=1.2 + (stage == "reference"), + color=colors[stage], label=names[stage]) + ax.axhline(1.0 - self.tail_quantile, color="tab:red", ls=":", lw=1) + ax.text(0.01, 1.0 - self.tail_quantile, f"p{100 * self.tail_quantile:g}", color="tab:red", + fontsize=8, va="bottom", transform=ax.get_yaxis_transform()) + ax.set(xlabel="Hs (m)", ylabel="exceedance probability", + title=f"(c) tail · {n} pairs, {int((hindcast > low[-1]).sum())} beyond the last quantile") + ax.grid(alpha=0.3, which="both") + ax.legend(fontsize=8) + + validation_scatter( + axes[1, 1], + reference=samples["reference"].to_numpy(), + estimate=samples["mapped"].to_numpy(), + reference_label=f"{names['reference']} Hs", + estimate_label=f"{names['mapped']} Hs", + title="(d) pairs: the map aligns distributions, not pairs", + units="m", + ) + if title: + fig.suptitle(title) + fig.tight_layout() + + return fig diff --git a/bluemath_tk/waves/hywaves/offshore.py b/bluemath_tk/waves/hywaves/offshore.py index 54bb241..6dcd62c 100644 --- a/bluemath_tk/waves/hywaves/offshore.py +++ b/bluemath_tk/waves/hywaves/offshore.py @@ -14,17 +14,23 @@ (:data:`EDH_ERA5_WAVES_URL`, hourly, 0.5 degree, 1940 onwards) at the goal nodes. ERA5 has no partition peak periods: ``ptp{i}`` are the partition mean periods (``mpww``, ``mwp1..3``, mean period Tm-1,0), shorter than the peak -period (about Tp / 1.1 for a JONSWAP spectrum); the attributes say so. The -three swells share the total-swell directional width (``dwps``). +period (about Tp / 1.1 for a JONSWAP spectrum); the attributes say so. +:func:`peak_period_samples` measures Tp / Tm-1,0 on ERA5 itself (hours one +partition dominates, against the bulk peak period ``pp1d``) and +:func:`partition_peak_periods` converts ``ptp{i}`` with it. The three swells +share the total-swell directional width (``dwps``). """ from __future__ import annotations import os +from collections.abc import Mapping import dask import numpy as np +import pandas as pd import xarray as xr +from scipy.integrate import trapezoid #: Earth Data Hub (DestinE) ERA5 single-level ocean-wave variables. EDH_ERA5_WAVES_URL = ( @@ -175,3 +181,93 @@ def era5_at_goals( } return out + + +def jonswap_tm10_ratio(gamma: float) -> float: + """ + Tm-1,0 / Tp of a JONSWAP spectrum (``gamma = 1``: Pierson-Moskowitz). + + Parameters + ---------- + gamma : float + Peak enhancement factor. + + Returns + ------- + float + ``m_-1 / m_0`` in units of the peak frequency. + """ + + f = np.linspace(0.3, 20.0, 40000) # f / fp + sigma = np.where(f <= 1.0, 0.07, 0.09) + s = f**-5 * np.exp(-1.25 * f**-4) * gamma ** np.exp(-((f - 1.0) ** 2) / (2 * sigma**2)) + return float(trapezoid(s / f, f) / trapezoid(s, f)) + + +def peak_period_samples(ds: xr.Dataset, min_fraction: float = 0.9) -> pd.DataFrame: + """ + Tp / Tm-1,0 at the hours one partition holds most of the energy. + + There the bulk peak period ``tp`` is the peak of that partition, so + ``tp / ptp{i}`` measures the ratio of its peak to its mean period. + + Parameters + ---------- + ds : xr.Dataset + ``hs``, ``tp`` and ``phs{i}``, ``ptp{i}`` (as :func:`era5_at_goals`). + min_fraction : float, optional + Least share of the energy ``(phs{i} / hs)²`` of the partition. + Default is 0.9. + + Returns + ------- + pd.DataFrame + ``partition`` (``"sea"`` for 0, ``"swell"`` for 1, 2, ...), ``ratio`` + (Tp / Tm-1,0), one row per selected hour and goal. + """ + + parts = sorted(int(name[3:]) for name in ds.data_vars if name.startswith("phs")) + rows = [] + for i in parts: + share = (ds[f"phs{i}"] / ds["hs"]) ** 2 + ratio = (ds["tp"] / ds[f"ptp{i}"]).where(share >= min_fraction).values.ravel() + ratio = ratio[np.isfinite(ratio)] + rows.append(pd.DataFrame({"partition": "sea" if i == 0 else "swell", "ratio": ratio})) + return pd.concat(rows, ignore_index=True) + + +def partition_peak_periods(ds: xr.Dataset, ratios: Mapping[str, float]) -> xr.Dataset: + """ + Partition peak periods from mean periods: ``ptp{i}`` x the Tp / Tm-1,0 ratio. + + Parameters + ---------- + ds : xr.Dataset + ``ptp{i}`` holding Tm-1,0 (as :func:`era5_at_goals`). + ratios : Mapping[str, float] + Tp / Tm-1,0 for ``"sea"`` (partition 0) and ``"swell"`` (1, 2, ...), + e.g. the medians of :func:`peak_period_samples`. + + Returns + ------- + xr.Dataset + Copy with ``ptp{i}`` peak periods; the ratio used is in each + variable's ``tp_tm10_ratio`` attribute. + """ + + out = ds.copy() + for name in [v for v in ds.data_vars if v.startswith("ptp")]: + kind = "sea" if name == "ptp0" else "swell" + out[name] = ds[name] * ratios[kind] + label = "wind sea" if kind == "sea" else f"swell {name[3:]}" + out[name].attrs = { + **ds[name].attrs, + "long_name": f"Peak wave period (Tm-1,0 x {ratios[kind]:.3f}), {label}", + "tp_tm10_ratio": float(ratios[kind]), + } + out.attrs["note"] = ( + "ptp* are peak periods from the partition mean periods (Tm-1,0) " + f"x Tp / Tm-1,0 (sea {ratios['sea']:.3f}, swell {ratios['swell']:.3f}); " + "the swells share the total-swell directional width (dwps)" + ) + return out diff --git a/bluemath_tk/waves/hywaves/plots.py b/bluemath_tk/waves/hywaves/plots.py index a995c18..6ec4884 100644 --- a/bluemath_tk/waves/hywaves/plots.py +++ b/bluemath_tk/waves/hywaves/plots.py @@ -524,12 +524,306 @@ def plot_wind_k_comparison(df: pd.DataFrame, title: str = "How K is estimated") return fig +def plot_peak_period_ratio( + samples: pd.DataFrame, theory: Mapping[str, float], title: str = "Tp / Tm-1,0" +) -> Figure: + """ + Measured Tp / Tm-1,0 per partition kind, against spectral shapes. + + Parameters + ---------- + samples : pd.DataFrame + ``partition`` and ``ratio`` (e.g. ``offshore.peak_period_samples``). + theory : Mapping[str, float] + Label -> ratio of reference spectra (e.g. JONSWAP for a few gamma), + drawn as vertical lines. + title : str, optional + Figure title. + + Returns + ------- + Figure + """ + + kinds = list(dict.fromkeys(samples["partition"])) + fig, axes = plt.subplots(1, len(kinds), figsize=(6 * len(kinds), 4), sharey=True, + constrained_layout=True, squeeze=False) + fig.suptitle(title, fontsize=14, fontweight="bold") + bins = np.linspace(0.9, 1.5, 61) + colors = plt.get_cmap("Greys")(np.linspace(0.4, 0.9, len(theory))) + for ax, kind in zip(axes[0], kinds): + ratio = samples.loc[samples["partition"] == kind, "ratio"] + ax.hist(ratio, bins=bins, density=True, color="#1f77b4", alpha=0.6) + ax.axvline(ratio.median(), color="#d62728", lw=2, label=f"median {ratio.median():.3f}") + for (label, value), color in zip(theory.items(), colors): + ax.axvline(value, color=color, ls="--", lw=1.2, label=f"{label} {value:.3f}") + ax.set(title=f"{kind} · {len(ratio):,} hours", xlabel="Tp / Tm-1,0", ylabel="density") + ax.legend(fontsize=9) + return fig + + +def _direction_arrow(ax: Axes, direction: float, xy: tuple[float, float], color: str, label: str) -> None: + """Arrow towards where waves / wind go, from a nautical coming-from *direction* (axes fraction).""" + rad = np.deg2rad(direction) + dx, dy = -0.1 * np.sin(rad), -0.1 * np.cos(rad) + start = (xy[0] - dx / 2, xy[1] - dy / 2) + ax.annotate("", xy=(start[0] + dx, start[1] + dy), xytext=start, xycoords="axes fraction", + arrowprops={"arrowstyle": "-|>", "color": color, "lw": 2}) + ax.text(xy[0], xy[1] - 0.1, label, transform=ax.transAxes, ha="center", va="top", fontsize=8, color=color) + + +def plot_wind_case_comparison(ds: xr.Dataset, size: float = 4) -> Figure: + """ + Training cases with and without wind, and the wind correction, at every site. + + One row per case: what the wind adds in the runs (with - without wind), + what the correction adds (corrected - without wind), on the same scale; + their difference (corrected - with wind) with its bias and RMSE; and the + run with wind (x, reference) against the run without wind and the + corrected one. Arrows: waves (black) and wind (red), towards where they + go. + + Parameters + ---------- + ds : xr.Dataset + ``hs0``, ``hs1``, ``hs_corrected`` on ``(case, sites)``, with site + ``longitude`` / ``latitude`` and case ``hs``, ``tp``, ``u10``, + ``dir``, ``u10dir``, ``wind_wave_angle``, ``ws_weight``, + ``hs_ws_u2`` coordinates. + size : float, optional + Panel size (inches). Default is 4. + + Returns + ------- + Figure + """ + + n = ds.sizes["case"] + fig, axes = plt.subplots(n, 4, figsize=(4 * size, n * size), squeeze=False, constrained_layout=True) + lon, lat = ds["longitude"].values, ds["latitude"].values + aspect = 1.0 / np.cos(np.deg2rad(np.nanmean(lat))) + for row, case in zip(axes, ds["case"].values): + c = ds.sel(case=case) + added = {"wind added, runs": (c["hs1"] - c["hs0"]).values, + "wind added, correction": (c["hs_corrected"] - c["hs0"]).values} + lim = max(float(np.nanpercentile(np.abs(np.r_[tuple(added.values())]), 99)), 0.05) + for ax, (label, value) in zip(row, added.items()): + sc = ax.scatter(lon, lat, c=value, s=2, cmap="RdBu_r", vmin=-lim, vmax=lim, rasterized=True) + ax.set_title(f"{label} (median {np.nanmedian(value):+.2f} m)", fontsize=10) + fig.colorbar(sc, ax=row[:2], label="Hs difference (m)", shrink=0.8) + err = (c["hs_corrected"] - c["hs1"]).values + sc = row[2].scatter(lon, lat, c=err, s=2, cmap="RdBu_r", vmin=-lim, vmax=lim, rasterized=True) + fig.colorbar(sc, ax=row[2], label="corrected - with wind (m)", shrink=0.8) + row[2].set_title(f"bias {np.nanmean(err):+.2f} m, RMSE {np.sqrt(np.nanmean(err**2)):.2f} m", fontsize=10) + for ax in row[:3]: + ax.set_aspect(aspect) + ax.tick_params(labelsize=7) + _direction_arrow(row[0], float(c["dir"]), (0.12, 0.88), "black", "waves") + _direction_arrow(row[0], float(c["u10dir"]), (0.30, 0.88), "#d62728", "wind") + row[0].set_ylabel( + f"case {case}: Hs {float(c['hs']):.1f} m, Tp {float(c['tp']):.1f} s\n" + f"U {float(c['u10']):.1f} m/s, angle {float(c['wind_wave_angle']):+.0f}°, " + f"w {float(c['ws_weight']):.2f}, gHs/U² {float(c['hs_ws_u2']):.2f}", + fontsize=9, + ) + ax = row[3] + ref = c["hs1"].values + ax.scatter(ref, c["hs0"].values, s=2, color="0.6", alpha=0.5, label="without wind", rasterized=True) + ax.scatter(ref, c["hs_corrected"].values, s=2, color="#1f77b4", alpha=0.5, label="corrected", rasterized=True) + top = float(np.nanmax(np.r_[ref, c["hs_corrected"].values])) + ax.plot([0, top], [0, top], "k--", lw=1) + ax.set(xlim=(0, top), ylim=(0, top), xlabel="Hs with wind (m)", ylabel="Hs (m)", title="all sites") + ax.legend(fontsize=8, markerscale=4, loc="upper left") + return fig + + +def plot_k_response( + ds: xr.Dataset, ranges: Mapping[str, tuple[float, float]] | None = None, title: str = "K response" +) -> Figure: + """ + ``K`` of the wind correction against each input, one line per site. + + Parameters + ---------- + ds : xr.Dataset + ``k_u10``, ``k_angle``, ``k_hs_ws_u2`` on ``(sites, )``, site + ``depth`` and ``fetch`` (m) coordinates and the base state + (``wind_dir``, ``hs_ws``, ``u10``, ``angle``) in the attributes. + ranges : Mapping[str, tuple of float], optional + ``u10``, ``angle``, ``hs_ws_u2`` -> range of the training cases; + outside it is shaded (the model keeps its edge value there). + title : str, optional + Figure title. + + Returns + ------- + Figure + """ + + panels = { + "u10": ("wind speed U10 (m/s)", f"Hs_ws {ds.attrs['hs_ws']:g} m, angle {ds.attrs['angle']:g}°"), + "angle": ("wind - wave angle (°)", f"Hs_ws {ds.attrs['hs_ws']:g} m, U {ds.attrs['u10']:g} m/s"), + "hs_ws_u2": ("g Hs_ws / U² (how developed the wind sea is)", + f"U {ds.attrs['u10']:g} m/s, angle {ds.attrs['angle']:g}°"), + } + fig, axes = plt.subplots(1, 3, figsize=(17, 4.8), sharey=True, constrained_layout=True) + fig.suptitle(f"{title} · wind from {ds.attrs['wind_dir']:g}°", fontsize=14, fontweight="bold") + for ax, (name, (xlabel, base)) in zip(axes, panels.items()): + x = ds[name].values + for site in ds["sites"].values: + s = ds.sel(sites=site) + ax.plot(x, s[f"k_{name}"], lw=2, + label=f"{site} ({float(s['depth']):.0f} m, fetch {float(s['fetch']) / 1e3:.0f} km)") + if ranges is not None and name in ranges: + low, high = ranges[name] + ax.axvspan(x.min(), low, color="0.85", zorder=0) + ax.axvspan(high, x.max(), color="0.85", zorder=0, label="outside training") + if name == "hs_ws_u2": + ax.set_xscale("log") + ax.axvline(0.24, color="0.4", ls=":", lw=1) + ax.text(0.24, 1.0, " fully developed", transform=ax.get_xaxis_transform(), va="top", fontsize=8, color="0.4") + ax.axhline(1, color="k", lw=0.8) + ax.set(xlabel=xlabel, xlim=(x.min(), x.max()), title=base) + axes[0].set_ylabel("K = Hs with / without wind") + axes[0].legend(fontsize=8, loc="upper left") + return fig + + +def plot_wind_sea_event( + ws: xr.Dataset, + hs: Mapping[str, xr.DataArray], + ranges: Mapping[str, tuple[float, float]] | None = None, + title: str = "", +) -> Figure: + """ + The wind correction at one site through an event. + + Rows: Hs of every source; the wind sea the correction sees (``hs_ws``) + and the wind (``u10``); how developed the wind sea is + (``g Hs_ws / U^2``); the wind-wave angle; ``K``. Shaded: outside the + training cases (*ranges*). + + Parameters + ---------- + ws : xr.Dataset + ``hs_ws``, ``u10``, ``hs_ws_u2``, ``angle``, ``k`` on ``time`` (one + site of ``wind_sea_inputs``). + hs : Mapping[str, xr.DataArray] + Label -> Hs on ``time`` (e.g. buoy, dynamic run, metamodel without and + with the correction). + ranges : Mapping[str, tuple of float], optional + ``hs_ws_u2``, ``angle`` -> training range. + title : str, optional + Figure title. + + Returns + ------- + Figure + """ + + fig, axes = plt.subplots(5, 1, figsize=(13, 13), sharex=True, constrained_layout=True, + gridspec_kw={"height_ratios": [2, 1, 1, 1, 1]}) + fig.suptitle(title, fontsize=14, fontweight="bold") + times = ws["time"].values + ax = axes[0] + for label, series in hs.items(): + marker = "o" if label.lower().startswith("buoy") else None + ax.plot(series["time"], series, lw=1.5, marker=marker, ms=2.5, ls="-" if marker is None else "none", label=label) + ax.set(ylabel="Hs (m)") + ax.legend(fontsize=9, ncol=2) + ax = axes[1] + ax.plot(times, ws["hs_ws"], color="#2ca02c", lw=1.5, label="wind sea Hs_ws (m)") + ax.set_ylabel("Hs_ws (m)", color="#2ca02c") + twin = ax.twinx() + twin.plot(times, ws["u10"], color="#d62728", lw=1.2, label="U10") + twin.set_ylabel("U10 (m/s)", color="#d62728") + for ax, name, label, log in ((axes[2], "hs_ws_u2", "g Hs_ws / U²", True), (axes[3], "angle", "wind - wave (°)", False)): + values = ws[name].values if name != "angle" else (ws[name].values + 180.0) % 360.0 - 180.0 + ax.plot(times, values, color="k", lw=1.2) + if ranges is not None and name in ranges: + low, high = ranges[name] + ax.axhspan(low, high, color="#1f77b4", alpha=0.12, label="training range") + ax.legend(fontsize=8, loc="upper left") + if log: + ax.set_yscale("log") + ax.axhline(0.24, color="0.4", ls=":", lw=1) + ax.set_ylabel(label) + axes[4].plot(times, ws["k"], color="#9467bd", lw=1.5) + axes[4].axhline(1, color="k", lw=0.8) + axes[4].set_ylabel("K") + axes[4].xaxis.set_major_formatter(mdates.DateFormatter("%d %b %H:%M")) + return fig + + +def plot_event_maps( + reference: xr.DataArray, + estimates: Mapping[str, xr.DataArray], + k: xr.DataArray | None = None, + reference_label: str = "reference", + title: str = "", + size: float = 4.5, +) -> Figure: + """ + Hs at every site at one time: a reference, each estimate minus it, and ``K``. + + Parameters + ---------- + reference : xr.DataArray + Hs on ``sites`` with ``site_longitude`` / ``site_latitude``. + estimates : Mapping[str, xr.DataArray] + Label -> Hs on the same sites. + k : xr.DataArray, optional + ``K`` of the wind correction on the same sites. + reference_label : str, optional + Name of the reference. + title : str, optional + Figure title. + size : float, optional + Panel size (inches). + + Returns + ------- + Figure + """ + + n = 1 + len(estimates) + int(k is not None) + fig, axes = plt.subplots(1, n, figsize=(n * size, size * 1.15), constrained_layout=True, squeeze=False) + axes = axes[0] + fig.suptitle(title, fontsize=14, fontweight="bold") + lon, lat = reference["site_longitude"].values, reference["site_latitude"].values + sc = axes[0].scatter(lon, lat, c=reference.values, s=3, cmap="viridis", vmin=0, rasterized=True) + fig.colorbar(sc, ax=axes[0], label="Hs (m)", shrink=0.8) + axes[0].set_title(f"{reference_label} (median {float(reference.median()):.2f} m)", fontsize=10) + diffs = {label: (est.sel(sites=reference["sites"]) - reference).values for label, est in estimates.items()} + lim = max(float(np.nanpercentile(np.abs(np.concatenate(list(diffs.values()))), 99)), 0.05) if diffs else 1.0 + for ax, (label, diff) in zip(axes[1:], diffs.items()): + sc = ax.scatter(lon, lat, c=diff, s=3, cmap="RdBu_r", vmin=-lim, vmax=lim, rasterized=True) + ax.set_title(f"{label} - {reference_label}\nbias {np.nanmean(diff):+.2f}, RMSE {np.sqrt(np.nanmean(diff**2)):.2f} m", + fontsize=10) + if diffs: + fig.colorbar(sc, ax=axes[1 : 1 + len(diffs)], label="Hs difference (m)", shrink=0.8) + if k is not None: + sc = axes[-1].scatter(lon, lat, c=k.sel(sites=reference["sites"]).values, s=3, cmap="magma_r", vmin=1, rasterized=True) + fig.colorbar(sc, ax=axes[-1], label="K", shrink=0.8) + axes[-1].set_title("K of the wind correction", fontsize=10) + aspect = 1.0 / np.cos(np.deg2rad(np.nanmean(lat))) + for ax in axes: + ax.set_aspect(aspect) + ax.tick_params(labelsize=7) + return fig + + __all__ = [ "SWEEP_METRICS", "goal_colors", "plot_direction_ranges", "plot_dynamic_forcing", + "plot_event_maps", "plot_goal_sectors", + "plot_k_response", "plot_mda_sweep", + "plot_peak_period_ratio", + "plot_wind_case_comparison", "plot_wind_k_comparison", + "plot_wind_sea_event", ] diff --git a/bluemath_tk/waves/hywaves/store.py b/bluemath_tk/waves/hywaves/store.py index ac98aa9..9092131 100644 --- a/bluemath_tk/waves/hywaves/store.py +++ b/bluemath_tk/waves/hywaves/store.py @@ -230,7 +230,7 @@ def read(self, time: Any, sites: Iterable[str]) -> xr.Dataset: valid = times != _NAT data = {} for name, extra in self.variables.items(): - values = np.stack([self.group[name].oindex[m, :, site_idx] for m in month_idx]) + values = self.group[name].oindex[month_idx, :, site_idx] # one call: concurrent chunk reads data[name] = (("time", "sites", *extra), values[valid]) coords = {"time": times[valid].view("datetime64[ns]"), "sites": sites} coords.update({d: self.coords[d] for extra in self.variables.values() for d in extra}) diff --git a/bluemath_tk/waves/hywaves/wind_correction.py b/bluemath_tk/waves/hywaves/wind_correction.py index 59b5dc3..9f52833 100644 --- a/bluemath_tk/waves/hywaves/wind_correction.py +++ b/bluemath_tk/waves/hywaves/wind_correction.py @@ -297,6 +297,73 @@ def _vector_mean(direction: xr.DataArray, weight: xr.DataArray, dims: list[str]) return np.mod(np.rad2deg(np.arctan2(s, c)), 360.0) +def wind_sea_inputs( + cube: xr.Dataset, + correction: WindCorrection, + fetch: xr.Dataset, + wind: xr.Dataset | None = None, +) -> xr.Dataset | None: + """ + The wind sea at each site and step of one batch, the inputs of ``K`` and ``K``. + + Parameters + ---------- + cube, correction, fetch, wind + As for :func:`wind_sea_contribution`. + + Returns + ------- + xr.Dataset or None + On ``(time, sites)``: ``e_ws`` (wind-sea energy ``sum w hs^2``, m²), + ``hs_ws``, ``tp_ws`` and ``ws_dir`` (its Hs, energy-weighted Tp and + direction), ``u10`` and ``u10dir`` (wind), ``angle`` (wind - wave), + ``fetch``, ``mean_depth`` (along the wind) and ``depth`` (site), + ``hs_ws_u2`` (``g Hs_ws / U^2``) and ``k``; None if nothing in the + batch is wind sea. + """ + + parts = ["goal", "partition"] + weight = cube["ws_weight"].fillna(0.0) + energy = (weight * cube["hs"].fillna(0.0) ** 2).transpose("time", "sites", ...) + e_ws = energy.sum(parts) + if not bool((e_ws > 0).any()): + return None + + with np.errstate(invalid="ignore", divide="ignore"): + ws_dir = _vector_mean(cube["dir"].fillna(0.0), energy, parts) + tp_ws = (energy * cube["tp"].fillna(0.0)).sum(parts) / e_ws + if wind is None: + u10 = (energy * cube["u10"].fillna(0.0)).sum(parts) / e_ws + u10dir = _vector_mean(cube["u10dir"].fillna(0.0), energy * cube["u10"].fillna(0.0), parts) + else: + u10 = wind["u10"].sel(time=cube["time"], sites=cube["sites"]) + u10dir = wind["u10dir"].sel(time=cube["time"], sites=cube["sites"]) + u10 = u10.fillna(0.0).transpose("time", "sites") + u10dir = u10dir.transpose("time", "sites") + hs_ws = np.sqrt(e_ws) + + fetch = fetch.sel(sites=cube["sites"].values) + name = "fetch_eff" if "fetch_eff" in fetch else "fetch" + fetch_u = _fetch_along(fetch, name, u10dir) + depth_u = _fetch_along(fetch, "mean_depth", u10dir) + site_depth = fetch["site_depth"].broadcast_like(hs_ws) + angle = u10dir - ws_dir + k = correction.predict( + hs_ws.values, u10.values, fetch_u.values, site_depth.values, depth_u.values, angle.values + ) + with np.errstate(invalid="ignore", divide="ignore"): + hs_ws_u2 = GRAVITY * hs_ws / u10**2 + + return xr.Dataset( + { + "e_ws": e_ws, "hs_ws": hs_ws, "tp_ws": tp_ws, "ws_dir": ws_dir, + "u10": u10, "u10dir": u10dir, "angle": angle, "fetch": fetch_u, + "mean_depth": depth_u, "depth": site_depth.reset_coords(drop=True), "hs_ws_u2": hs_ws_u2, + "k": xr.DataArray(k, dims=hs_ws.dims, coords=hs_ws.coords), + } + ) + + def wind_sea_contribution( cube: xr.Dataset, correction: WindCorrection, @@ -340,47 +407,20 @@ def wind_sea_contribution( if nothing in the batch is wind sea. """ - parts = ["goal", "partition"] - weight = cube["ws_weight"].fillna(0.0) - energy = (weight * cube["hs"].fillna(0.0) ** 2).transpose("time", "sites", ...) - e_ws = energy.sum(parts) - if not bool((e_ws > 0).any()): + ws = wind_sea_inputs(cube, correction, fetch, wind) + if ws is None: return None - - with np.errstate(invalid="ignore", divide="ignore"): - ws_dir = _vector_mean(cube["dir"].fillna(0.0), energy, parts) - if wind is None: - u10 = (energy * cube["u10"].fillna(0.0)).sum(parts) / e_ws - u10dir = _vector_mean(cube["u10dir"].fillna(0.0), energy * cube["u10"].fillna(0.0), parts) - else: - u10 = wind["u10"].sel(time=cube["time"], sites=cube["sites"]) - u10dir = wind["u10dir"].sel(time=cube["time"], sites=cube["sites"]) - u10 = u10.fillna(0.0).transpose("time", "sites") - u10dir = u10dir.transpose("time", "sites") - hs_ws = np.sqrt(e_ws) - - fetch = fetch.sel(sites=cube["sites"].values) - name = "fetch_eff" if "fetch_eff" in fetch else "fetch" - fetch_u = _fetch_along(fetch, name, u10dir) - depth_u = _fetch_along(fetch, "mean_depth", u10dir) - site_depth = fetch["site_depth"].broadcast_like(hs_ws) - angle = u10dir - ws_dir - k = correction.predict( - hs_ws.values, u10.values, fetch_u.values, site_depth.values, depth_u.values, angle.values - ) - k = xr.DataArray(k, dims=hs_ws.dims, coords=hs_ws.coords) - hs_add = np.sqrt(np.clip(k**2 - 1.0, 0.0, None) * e_ws) + hs_add = np.sqrt(np.clip(ws["k"] ** 2 - 1.0, 0.0, None) * ws["e_ws"]) if tp == "growth": from ..wind import equivalent_fetch, growth_tp with np.errstate(invalid="ignore", divide="ignore"): - grown = equivalent_fetch(hs_ws.values, u10.values) + fetch_u.values - tp_add = growth_tp(u10.values, np.where(np.isfinite(grown), grown, 1e12), depth_u.values) - tp_add = xr.DataArray(tp_add, dims=hs_ws.dims, coords=hs_ws.coords) + grown = equivalent_fetch(ws["hs_ws"].values, ws["u10"].values) + ws["fetch"].values + tp_add = growth_tp(ws["u10"].values, np.where(np.isfinite(grown), grown, 1e12), ws["mean_depth"].values) + tp_add = xr.DataArray(tp_add, dims=hs_add.dims, coords=hs_add.coords) elif tp == "wind_sea": - with np.errstate(invalid="ignore", divide="ignore"): - tp_add = (energy * cube["tp"].fillna(0.0)).sum(parts) / e_ws + tp_add = ws["tp_ws"] else: raise ValueError(f"tp must be 'wind_sea' or 'growth', got {tp!r}") @@ -389,9 +429,9 @@ def wind_sea_contribution( { "hs": hs_add.where(keep), "tp": tp_add.where(keep), - "dir": u10dir.where(keep), + "dir": ws["u10dir"].where(keep), "spr": xr.full_like(hs_add, float(spr)).where(keep), - "k": k, + "k": ws["k"], } ) diff --git a/bluemath_tk/wrappers/snapwave/__init__.py b/bluemath_tk/wrappers/snapwave/__init__.py index fe6db22..ad64491 100644 --- a/bluemath_tk/wrappers/snapwave/__init__.py +++ b/bluemath_tk/wrappers/snapwave/__init__.py @@ -9,12 +9,13 @@ from .snapwave_utils import ( BOUNDARY_VARS, FORCING_FILES, + MAX_ERROR, POINT_VARS, boundary_time_series, - buffer_polygon, - build_enclosure_polygon, + open_boundary_enclosure, parse_snapwave_inp, read_boundary_nodes, + read_convergence, read_enclosure_polygon, read_forcing_table, read_his_file, @@ -31,15 +32,16 @@ __all__ = [ "BOUNDARY_VARS", "FORCING_FILES", + "MAX_ERROR", "POINT_VARS", "SnapWaveDynamicModelWrapper", "SnapWaveMetaModelWrapper", "SnapWaveModelWrapper", "boundary_time_series", - "buffer_polygon", - "build_enclosure_polygon", + "open_boundary_enclosure", "parse_snapwave_inp", "read_boundary_nodes", + "read_convergence", "read_enclosure_polygon", "read_forcing_table", "read_his_file", diff --git a/bluemath_tk/wrappers/snapwave/snapwave_utils.py b/bluemath_tk/wrappers/snapwave/snapwave_utils.py index 72c4dfd..460f765 100644 --- a/bluemath_tk/wrappers/snapwave/snapwave_utils.py +++ b/bluemath_tk/wrappers/snapwave/snapwave_utils.py @@ -12,12 +12,13 @@ from __future__ import annotations +import re from pathlib import Path import numpy as np import pandas as pd import xarray as xr -from shapely.geometry import MultiPolygon, Polygon +from shapely.geometry import LineString, MultiPolygon, Polygon #: SnapWave history (``his_file``) variables -> short names. POINT_VARS: dict[str, str] = { @@ -45,83 +46,32 @@ # --------------------------------------------------------------------------- -def _largest_polygon(geometry: Polygon | MultiPolygon) -> Polygon: - """Return *geometry* itself, or the largest part of a MultiPolygon.""" - - if isinstance(geometry, MultiPolygon): - geometry = max(geometry.geoms, key=lambda g: g.area) - if not isinstance(geometry, Polygon): - raise TypeError(f"Expected Polygon/MultiPolygon, got {geometry.geom_type}") - - return geometry - - -def buffer_polygon( - polygon: Polygon | MultiPolygon, - buffer_m: float, - geo_crs: str = "EPSG:4326", -) -> Polygon | MultiPolygon: - """ - Buffer a lon/lat polygon by *buffer_m* metres in a local UTM projection. - - Parameters - ---------- - polygon : Polygon or MultiPolygon - Input geometry in *geo_crs*. - buffer_m : float - Outward buffer distance in metres. ``0`` returns *polygon* unchanged. - geo_crs : str, optional - CRS of input and output coordinates. Default is "EPSG:4326". - - Returns - ------- - Polygon or MultiPolygon - Buffered geometry in *geo_crs*. +def open_boundary_enclosure(boundary: np.ndarray) -> Polygon: """ + SnapWave enclosure of an open boundary: a thin polygon around its nodes. - if buffer_m <= 0: - return polygon - - import geopandas as gpd - - gdf = gpd.GeoDataFrame(geometry=[polygon], crs=geo_crs) - work_crs = gdf.estimate_utm_crs() - if work_crs is None: - raise ValueError("Could not estimate a UTM CRS for the polygon.") - - return gdf.to_crs(work_crs).buffer(buffer_m).to_crs(geo_crs).iloc[0] - - -def build_enclosure_polygon( - polygon: Polygon | MultiPolygon, - buffer_m: float = 0.0, - geo_crs: str = "EPSG:4326", -) -> Polygon: - """ - Build a simple SnapWave enclosure around an area of interest. - - Exterior ring only (holes dropped), buffered outwards, then reduced to its - minimum rotated rectangle (4 vertices). + SnapWave forces every mesh edge node inside the enclosure, so it must hold + only the open (sea) boundary; a coast node inside it would radiate the + boundary spectra into the domain. The polygon is the boundary polyline + buffered by a quarter of its shortest segment, with flat ends. Parameters ---------- - polygon : Polygon or MultiPolygon - Area of interest; for a MultiPolygon the largest part is used. - buffer_m : float, optional - Outward buffer in metres. Default is 0. - geo_crs : str, optional - CRS of input and output coordinates. Default is "EPSG:4326". + boundary : np.ndarray + ``(n, 2)`` open-boundary nodes in order (e.g. + :func:`~bluemath_tk.topo_bathy.mesh.open_boundary_nodes`). Returns ------- Polygon - The enclosure quadrilateral. + The enclosure. """ - exterior = Polygon(_largest_polygon(polygon).exterior.coords) - buffered = _largest_polygon(buffer_polygon(exterior, buffer_m, geo_crs=geo_crs)) + boundary = np.asarray(boundary, dtype=float).reshape(-1, 2) + segments = np.hypot(*np.diff(boundary, axis=0).T) + width = 0.25 * float(segments[segments > 0].min()) - return buffered.minimum_rotated_rectangle + return LineString(boundary).buffer(width, cap_style="flat", join_style="mitre") def write_polygon_vertices_to_txt( @@ -392,6 +342,44 @@ def read_his_file( # --------------------------------------------------------------------------- +#: Final relative error above which a SnapWave run did not converge. +MAX_ERROR = 0.01 + + +def read_convergence(case_dir: str | Path, log_file: str = "wrapper_out.log") -> float: + """ + Relative error of SnapWave's last iteration, from the run log. + + SnapWave stops at ``niter`` sweeps whatever the error; an error near 1 + means the solution never settled (e.g. wind mode under swell and light + wind), while converged runs end below ~1e-3. + + Parameters + ---------- + case_dir : str or Path + Case folder. + log_file : str, optional + Standard output of the run. Default is "wrapper_out.log". + + Returns + ------- + float + Error of the last iteration. + + Raises + ------ + ValueError + If the log has no iteration line (the run did not start or crashed). + """ + + text = (Path(case_dir) / log_file).read_text(errors="replace") + errors = re.findall(r"iteration\s+\d+\s+error\s*=\s*([-+.\dEe]+)", text) + if not errors: + raise ValueError(f"no iteration in {Path(case_dir) / log_file}") + + return float(errors[-1]) + + def boundary_time_series( forcing: xr.Dataset, node_dim: str = "node", diff --git a/bluemath_tk/wrappers/snapwave/snapwave_wrapper.py b/bluemath_tk/wrappers/snapwave/snapwave_wrapper.py index bf82728..4c6eff0 100644 --- a/bluemath_tk/wrappers/snapwave/snapwave_wrapper.py +++ b/bluemath_tk/wrappers/snapwave/snapwave_wrapper.py @@ -32,7 +32,9 @@ from .._base_wrappers import BaseModelWrapper from .snapwave_utils import ( + MAX_ERROR, POINT_VARS, + read_convergence, read_his_file, write_points_to_txt, write_polygon_vertices_to_txt, @@ -81,6 +83,9 @@ class SnapWaveModelWrapper(BaseModelWrapper): context gets ``n_nodes`` (number of boundary nodes), and ``map_file`` / ``map_dirspr`` from the optional boolean ``save_map`` case parameter. + Postprocessing checks convergence (:func:`~.snapwave_utils.read_convergence`): + a case whose last error is above ``max_error`` is not used. + Attributes ---------- available_launchers : dict @@ -113,6 +118,7 @@ def __init__( point_coords: dict[str, str] | None = None, point_variables: tuple[str, ...] | None = tuple(POINT_VARS.values()), store_parameters: tuple[str, ...] = (), + max_error: float = MAX_ERROR, ) -> None: """ Initialise the SnapWave wrapper. @@ -132,8 +138,8 @@ def __init__( ``(n_nodes, 2)`` boundary node coordinates, in the order of the per-node forcing columns. enclosure_polygon : Polygon or MultiPolygon - Computational enclosure (see - :func:`~bluemath_tk.wrappers.snapwave.snapwave_utils.build_enclosure_polygon`). + Enclosure of the open boundary (see + :func:`~bluemath_tk.wrappers.snapwave.snapwave_utils.open_boundary_enclosure`). output_points : np.ndarray or pd.DataFrame Observation points: ``(n, 2)`` array, or a DataFrame with ``lon`` and ``lat`` columns. @@ -161,6 +167,9 @@ def __init__( Scalar case parameters stored in each postprocessed case as ``f"{name}_forcing"`` variables along the case dimension (e.g. ``("hs", "tp", "dir", "spr", "wl")`` for metamodel training data). + max_error : float, optional + Last-iteration error above which a case did not converge. Default + :data:`~.snapwave_utils.MAX_ERROR`. """ super().__init__( @@ -193,6 +202,31 @@ def __init__( self.point_coords = dict(point_coords or {}) self.point_variables = point_variables self.store_parameters = tuple(store_parameters) + self.max_error = float(max_error) + + def converged(self, case_dir: str) -> bool: + """ + Whether a finished case converged (its last error is below ``max_error``). + + Parameters + ---------- + case_dir : str + Case folder. + + Returns + ------- + bool + True when converged; a warning is logged otherwise. + """ + + error = read_convergence(case_dir) + if error > self.max_error: + self.logger.warning( + f"{os.path.basename(case_dir)} did not converge (error {error:.3g})" + ) + return False + + return True @property def n_nodes(self) -> int: @@ -390,6 +424,8 @@ def postprocess_case( """ Read one case's history output (``hs``, ``tp``, ``dir``, ``spr``). + A case that did not converge keeps its time step with NaN fields. + Parameters ---------- case_num : int @@ -406,6 +442,8 @@ def postprocess_case( """ ds = read_his_file(case_dir, self.his_filename) + if not self.converged(case_dir): + ds = ds.map(lambda v: v.where(False) if "stations" in v.dims else v) return self._format_output(ds, case_context, row_dim="time") @@ -519,6 +557,8 @@ def postprocess_case( """ Read one stationary case, indexed by ``case_num`` instead of ``time``. + A case that did not converge is dropped from the training cases. + Parameters ---------- case_num : int @@ -534,6 +574,8 @@ def postprocess_case( Output on ``(case_num, stations)``. """ + if not self.converged(case_dir): + raise ValueError(f"case {case_num} did not converge") ds = read_his_file(case_dir, self.his_filename) ds = ds.isel(time=[-1]).rename({"time": "case_num"}) ds = ds.assign_coords(case_num=[case_num]) diff --git a/bluemath_tk/wrappers/swan/swan_plots.py b/bluemath_tk/wrappers/swan/swan_plots.py index fc8c8a5..e475376 100644 --- a/bluemath_tk/wrappers/swan/swan_plots.py +++ b/bluemath_tk/wrappers/swan/swan_plots.py @@ -115,7 +115,7 @@ def load_swan_case(case_dir: str | Path, extent: Extent | None = None) -> SwanCa m = re.search(r"POINTS\s+'[^']*'\s+FILE\s+'([^']+)'", text) if m and (case_dir / m.group(1)).is_file(): points = np.loadtxt(case_dir / m.group(1)).reshape(-1, 2) - m = re.search(r"TABLE\s+'[^']*'\s+\w+\s+'([^']+)'\s+([A-Z ]+)", text) + m = re.search(r"TABLE\s+'[^']*'\s+\w+\s+'([^']+)'\s+([A-Z0-9 ]+)", text) if m and (case_dir / m.group(1)).is_file(): df = read_swan_table(str(case_dir / m.group(1)), m.group(2).split()) table = {c: df[c].to_numpy(float) for c in df} diff --git a/bluemath_tk/wrappers/swan/swan_utils.py b/bluemath_tk/wrappers/swan/swan_utils.py index b17b507..6e10f55 100644 --- a/bluemath_tk/wrappers/swan/swan_utils.py +++ b/bluemath_tk/wrappers/swan/swan_utils.py @@ -322,6 +322,7 @@ def vortex2SWAN( "TPS": "tp", "DIR": "dir", "DSPR": "spr", + "TMM10": "tm10", } diff --git a/bluemath_tk/wrappers/swan/swan_wrapper.py b/bluemath_tk/wrappers/swan/swan_wrapper.py index f6d24a5..a30d660 100644 --- a/bluemath_tk/wrappers/swan/swan_wrapper.py +++ b/bluemath_tk/wrappers/swan/swan_wrapper.py @@ -945,7 +945,9 @@ class SwanBoundaryModelWrapper(BaseModelWrapper): - ``map_file`` (``""`` unless the optional boolean ``save_map`` case parameter is true) and ``map_quantities``, for a ``BLOCK 'COMPGRID' ... LAY 3`` output on the whole grid (read it with - :func:`~bluemath_tk.wrappers.swan.swan_utils.read_swan_map`). + :func:`~bluemath_tk.wrappers.swan.swan_utils.read_swan_map`); + - ``spectra_file`` (``""`` without ``spectra_points``) and + ``spectra_points_file``, for ``POINTS`` + ``SPECOUT`` of 2D spectra. Attributes ---------- @@ -962,7 +964,7 @@ class SwanBoundaryModelWrapper(BaseModelWrapper): "geoocean-cluster": f"{GEOOCEAN_SWAN_MODULE}\nswan.exe", } - table_quantities: tuple[str, ...] = ("HSIGN", "TPS", "DIR", "DSPR") + table_quantities: tuple[str, ...] = ("HSIGN", "TPS", "TMM10", "DIR", "DSPR") map_quantities: tuple[str, ...] = ("HSIGN", "TPS", "DIR", "DSPR", "DEPTH") def __init__( @@ -983,6 +985,9 @@ def __init__( point_dim: str = "points", point_coords: dict[str, str] | None = None, store_parameters: tuple[str, ...] = (), + spectra_points: np.ndarray | pd.DataFrame | None = None, + spectra_points_filename: str = "output_spectra.txt", + spectra_filename: str = "output_spectra.spc", ) -> None: """ Initialise the SWAN boundary-forced wrapper. @@ -1024,6 +1029,11 @@ def __init__( store_parameters : tuple of str, optional Scalar case parameters stored as ``f"{name}_forcing"`` along the case dimension. + spectra_points : np.ndarray or pd.DataFrame, optional + Points of the 2D spectra output, as *output_points*. Default no + spectra. + spectra_points_filename, spectra_filename : str, optional + Spectra points file and ``SPECOUT`` output names. """ super().__init__( @@ -1061,6 +1071,15 @@ def __init__( self.point_dim = point_dim self.point_coords = dict(point_coords or {}) self.store_parameters = tuple(store_parameters) + if isinstance(spectra_points, pd.DataFrame): + spectra_points = spectra_points[["lon", "lat"]].to_numpy(float) + self.spectra_points = ( + None + if spectra_points is None + else np.asarray(spectra_points, dtype=float).reshape(-1, 2) + ) + self.spectra_points_filename = spectra_points_filename + self.spectra_filename = spectra_filename @property def n_nodes(self) -> int: @@ -1128,6 +1147,14 @@ def build_case(self, case_context: dict, case_dir: str) -> None: save_map = bool(case_context.get("save_map", False)) case_context["map_file"] = self.map_filename if save_map else "" case_context["map_quantities"] = " ".join(self.map_quantities) + spectra = self.spectra_points is not None and len(self.spectra_points) > 0 + case_context["spectra_points_file"] = self.spectra_points_filename + case_context["spectra_file"] = self.spectra_filename if spectra else "" + if spectra: + write_swan_points( + self.spectra_points, + os.path.join(case_dir, self.spectra_points_filename), + ) def monitor_cases(self, value_counts: str | None = None): """ diff --git a/tests/data/swan/boundary/INPUT b/tests/data/swan/boundary/INPUT index 7cd73eb..dc18273 100644 --- a/tests/data/swan/boundary/INPUT +++ b/tests/data/swan/boundary/INPUT @@ -10,5 +10,7 @@ BOUNDSPEC SIDE {{ boundary_side }} VARIABLE PAR & {% endfor %}POINTS 'sites' FILE '{{ points_file }}' TABLE 'sites' NOHEAD '{{ table_file }}' {{ table_quantities }} {% if map_file %}BLOCK 'COMPGRID' NOHEAD '{{ map_file }}' LAY 3 {{ map_quantities }} +{% endif %}{% if spectra_file %}POINTS 'buoys' FILE '{{ spectra_points_file }}' +SPECOUT 'buoys' SPEC2D ABS '{{ spectra_file }}' {% endif %}COMPUTE STOP diff --git a/tests/topo_bathy/test_mesh.py b/tests/topo_bathy/test_mesh.py index 97f08da..5593c4e 100644 --- a/tests/topo_bathy/test_mesh.py +++ b/tests/topo_bathy/test_mesh.py @@ -12,6 +12,8 @@ depth_at_points, mask_raster_to_mesh, mesh_to_raster, + open_boundary_nodes, + outline_rings, read_mesh, ) from bluemath_tk.waves.hywaves.output_points import ( @@ -48,6 +50,38 @@ def _quad_mesh() -> Mesh: return Mesh(x=x, y=y, faces=faces, z=z) +class TestOpenBoundary(unittest.TestCase): + """Open boundary of the unit square from points on its north edge.""" + + def setUp(self): + """Build the test data.""" + + self.mesh = _quad_mesh() + + def test_outline(self): + """The square has one outline ring with its 40 edge nodes.""" + + rings = outline_rings(self.mesh) + self.assertEqual([len(r) for r in rings], [40]) + + def test_north_edge_out_to_the_coast(self): + """North edge plus both laterals down to the south (coast) corners.""" + + for points in ([[0.2, 1.0], [0.5, 1.0], [0.8, 1.0]], [[0.8, 1.0], [0.5, 1.0], [0.2, 1.0]]): + nodes = open_boundary_nodes(self.mesh, np.array(points)) + xy = np.c_[self.mesh.x[nodes], self.mesh.y[nodes]] + self.assertEqual(len(nodes), 31) + self.assertFalse(((xy[:, 1] == 0) & (xy[:, 0] > 0) & (xy[:, 0] < 1)).any()) + ends = {tuple(xy[0]), tuple(xy[-1])} + self.assertEqual(ends, {(0.0, 0.0), (1.0, 0.0)}) + + def test_points_out_of_order(self): + """Points not in order along the outline are refused.""" + + with self.assertRaises(ValueError): + open_boundary_nodes(self.mesh, np.array([[0.2, 1.0], [0.8, 1.0], [0.5, 1.0], [0.6, 1.0]])) + + class TestMesh(unittest.TestCase): """Depth at points and rasters from a mixed quad/triangle mesh.""" diff --git a/tests/waves/test_calibration.py b/tests/waves/test_calibration.py index e66706b..638d125 100644 --- a/tests/waves/test_calibration.py +++ b/tests/waves/test_calibration.py @@ -3,8 +3,9 @@ import numpy as np import pandas as pd import xarray as xr +from matplotlib import pyplot as plt -from bluemath_tk.waves.calibration import CalVal, sector_index +from bluemath_tk.waves.calibration import CalVal, QuantileMapping, sector_index def _pairs(n: int = 2000, seed: int = 0) -> pd.DataFrame: @@ -90,5 +91,76 @@ def test_bulk_only(self): self.assertGreater(cv.coefficients.query("sector == 12")["factor"].item(), 1.0) + +class TestQuantileMapping(unittest.TestCase): + """Second calibration step: the Hs distribution, tail included.""" + + def _fit(self, n: int = 20000) -> tuple[QuantileMapping, np.ndarray, np.ndarray]: + rng = np.random.default_rng(0) + hs = rng.weibull(1.3, n) * 1.5 + reference = hs + np.maximum(hs - 3.0, 0.0) * 0.3 + return QuantileMapping().fit(hs, reference), hs, reference + + def test_recovers_distribution(self): + """Mapped quantiles match the reference, up to the tail quantile.""" + + mapping, hs, reference = self._fit() + probabilities = [0.5, 0.9, 0.99, 0.999] + np.testing.assert_allclose( + np.quantile(mapping.map(hs), probabilities), + np.quantile(reference, probabilities), + atol=0.02, + ) + + def test_constant_beyond_tail_and_monotonic(self): + """Beyond the last quantile the offset is constant; the map never decreases.""" + + mapping, hs, _ = self._fit() + last = mapping.table.iloc[-1] + offset = last["reference"] - last["hindcast"] + self.assertAlmostEqual(mapping.map(np.array([20.0]))[0], 20.0 + offset) + self.assertTrue((np.diff(mapping.map(np.sort(hs))) >= -1e-9).all()) + + def test_correct_scales_partitions(self): + """Partition Hs follow the bulk ratio; other variables and NaN stay.""" + + mapping, _, _ = self._fit() + data = pd.DataFrame( + {"hs": [0.0, 5.0, np.nan], "phs0": [0.0, 4.0, np.nan], "pdir0": [10.0, 20.0, 30.0]} + ) + out = mapping.correct(data) + ratio = out.loc[1, "hs"] / 5.0 + self.assertGreater(ratio, 1.0) + self.assertAlmostEqual(out.loc[1, "phs0"], 4.0 * ratio) + self.assertEqual(out.loc[0, "hs"], 0.0) + self.assertTrue(np.isnan(out.loc[2, "hs"])) + pd.testing.assert_series_equal(out["pdir0"], data["pdir0"]) + + def test_dataset_and_too_few_pairs(self): + """Datasets give the same result; a tail on too few pairs is refused.""" + + mapping, _, _ = self._fit() + table = pd.DataFrame({"hs": [1.0, 5.0], "phs0": [1.0, 4.0]}) + dataset = xr.Dataset({k: ("time", table[k].to_numpy()) for k in table}) + np.testing.assert_allclose(mapping.correct(dataset)["phs0"], mapping.correct(table)["phs0"]) + with self.assertRaises(ValueError): + QuantileMapping().fit(np.ones(1000), np.ones(1000)) + + def test_quantiles_and_plot(self): + """Diagnostics on the fitted pairs; the samples are not saved with the model.""" + + mapping, hs, _ = self._fit() + table = mapping.quantiles(raw=hs * 0.9) + self.assertEqual(list(table.columns), ["reference", "raw", "hindcast", "mapped", "n_above"]) + np.testing.assert_allclose(table["mapped"], table["reference"], atol=0.02) + self.assertEqual(table.loc[0.999, "n_above"], 20) + fig = mapping.plot(raw=hs * 0.9, labels={"reference": "satellite"}) + self.assertEqual(len(fig.axes), 5) + plt.close(fig) + with self.assertRaises(ValueError): + mapping.evaluated_samples(raw=hs[:10]) + self.assertNotIn("_samples", mapping.__getstate__()) + + if __name__ == "__main__": unittest.main() diff --git a/tests/waves/test_hywaves_offshore.py b/tests/waves/test_hywaves_offshore.py index 564087a..e3bf50b 100644 --- a/tests/waves/test_hywaves_offshore.py +++ b/tests/waves/test_hywaves_offshore.py @@ -10,7 +10,10 @@ ERA5_TO_HYWAVES, EDH_KEY_VARIABLE, era5_at_goals, + jonswap_tm10_ratio, open_era5_waves, + partition_peak_periods, + peak_period_samples, ) @@ -67,5 +70,47 @@ def test_key_required(self): open_era5_waves() +def _partitions(phs: list[list[float]]) -> xr.Dataset: + """Two goals x hours; tp = 11, ptp{i} = 10 (Tm-1,0); hs from the partitions.""" + + phs = np.asarray(phs, dtype=float) # (partition, time) + ds = xr.Dataset({f"phs{i}": ("time", h) for i, h in enumerate(phs)}) + ds = ds.expand_dims(goal=[1, 2]) + ds["hs"] = np.sqrt(sum(ds[f"phs{i}"] ** 2 for i in range(len(phs)))) + ds["tp"] = xr.full_like(ds["hs"], 11.0) + for i in range(len(phs)): + ds[f"ptp{i}"] = xr.full_like(ds["hs"], 10.0) + return ds + + +class TestPeakPeriods(unittest.TestCase): + """Partition mean periods (Tm-1,0) to peak periods.""" + + def test_jonswap_ratio(self): + """Pierson-Moskowitz Tm-1,0 / Tp = 0.857; narrower spectra are closer to 1.""" + + self.assertAlmostEqual(jonswap_tm10_ratio(1.0), 0.857, places=3) + self.assertGreater(jonswap_tm10_ratio(3.3), jonswap_tm10_ratio(1.0)) + + def test_samples_only_where_one_partition_dominates(self): + """Hour 0: sea dominates; hour 1: swell 1; hour 2: mixed (no sample).""" + + ds = _partitions([[3.0, 0.1, 1.0], [0.1, 2.0, 1.0]]) + samples = peak_period_samples(ds, min_fraction=0.9) + self.assertEqual(samples["partition"].value_counts().to_dict(), {"sea": 2, "swell": 2}) + np.testing.assert_allclose(samples["ratio"], 1.1) + + def test_conversion(self): + """ptp0 x sea ratio, swells x swell ratio, recorded in the attributes.""" + + ds = _partitions([[1.0], [1.0], [1.0]]) + out = partition_peak_periods(ds, {"sea": 1.2, "swell": 1.1}) + np.testing.assert_allclose(out["ptp0"], 12.0) + np.testing.assert_allclose(out["ptp2"], 11.0) + np.testing.assert_allclose(ds["ptp0"], 10.0) + self.assertEqual(out["ptp1"].attrs["tp_tm10_ratio"], 1.1) + self.assertIn("peak", out["ptp1"].attrs["long_name"].lower()) + + if __name__ == "__main__": unittest.main() diff --git a/tests/waves/test_hywaves_plots.py b/tests/waves/test_hywaves_plots.py index 02260ad..f7b6af0 100644 --- a/tests/waves/test_hywaves_plots.py +++ b/tests/waves/test_hywaves_plots.py @@ -83,6 +83,53 @@ def test_wind_k_comparison(self): ) self.assertEqual(len(plots.plot_wind_k_comparison(df).axes), 2) + def test_peak_period_ratio(self): + df = pd.DataFrame({"partition": ["sea"] * 50 + ["swell"] * 50, + "ratio": np.linspace(1.0, 1.3, 100)}) + fig = plots.plot_peak_period_ratio(df, {"JONSWAP 3.3": 1.107}) + self.assertEqual(len(fig.axes), 2) + + def test_wind_case_comparison(self): + rng = np.random.default_rng(0) + hs0 = rng.uniform(0.5, 3.0, (2, 30)) + case = {"hs": [3.0, 5.0], "tp": [7.0, 9.0], "dir": [300.0, 200.0], "u10": [15.0, 25.0], + "u10dir": [320.0, 180.0], "wind_wave_angle": [20.0, -20.0], "ws_weight": [1.0, 0.6], + "hs_ws_u2": [0.13, 0.05]} + ds = xr.Dataset( + {"hs0": (("case", "sites"), hs0), "hs1": (("case", "sites"), hs0 * 1.2), + "hs_corrected": (("case", "sites"), hs0 * 1.15)}, + coords={"case": [3, 7], "longitude": ("sites", np.linspace(3, 6, 30)), + "latitude": ("sites", np.linspace(51, 53, 30)), + **{k: ("case", v) for k, v in case.items()}}, + ) + fig = plots.plot_wind_case_comparison(ds, size=2) + self.assertEqual(len(fig.axes), 2 * 4 + 2 * 2) # 4 panels + 2 colorbars per case + + def test_k_response(self): + x = {"u10": np.linspace(1, 40, 5), "angle": np.linspace(-90, 90, 5), "hs_ws_u2": np.geomspace(0.002, 0.6, 5)} + ds = xr.Dataset( + {f"k_{name}": (("sites", name), np.full((2, 5), 1.2)) for name in x}, + coords={"sites": ["a", "b"], **x, "depth": ("sites", [8.0, 20.0]), "fetch": ("sites", [4e4, 2e4])}, + attrs={"wind_dir": 300.0, "hs_ws": 2.0, "u10": 15.0, "angle": 0.0}, + ) + fig = plots.plot_k_response(ds, {"u10": (3.0, 35.0), "hs_ws_u2": (0.004, 0.29)}) + self.assertEqual(len(fig.axes), 3) + + def test_wind_sea_event(self): + time = pd.date_range("2013-12-05", periods=12, freq="h") + ws = xr.Dataset({name: ("time", np.linspace(0.1, 2.0, 12)) for name in ("hs_ws", "u10", "hs_ws_u2", "angle", "k")}, + coords={"time": time}) + hs = {"Buoy": ws["hs_ws"] * 2, "model": ws["hs_ws"] * 1.8} + fig = plots.plot_wind_sea_event(ws, hs, {"hs_ws_u2": (0.004, 0.29), "angle": (-60.0, 60.0)}) + self.assertEqual(len(fig.axes), 6) # 5 rows + the wind twin axis + + def test_event_maps(self): + coords = {"sites": ["a", "b", "c"], "site_longitude": ("sites", [3.0, 4.0, 5.0]), + "site_latitude": ("sites", [51.0, 52.0, 53.0])} + ref = xr.DataArray([1.0, 2.0, 3.0], dims="sites", coords=coords) + fig = plots.plot_event_maps(ref, {"m1": ref * 1.1, "m2": ref * 0.9}, k=ref * 0 + 1.2) + self.assertEqual(len(fig.axes), 4 + 3) # 4 maps + 3 colorbars + if __name__ == "__main__": unittest.main() diff --git a/tests/waves/test_wind_correction.py b/tests/waves/test_wind_correction.py index b5ad83a..2327140 100644 --- a/tests/waves/test_wind_correction.py +++ b/tests/waves/test_wind_correction.py @@ -14,6 +14,7 @@ correction_features, offshore_wind_sea, wind_sea_contribution, + wind_sea_inputs, ) SITES = ["a", "b"] @@ -103,6 +104,19 @@ def test_energy_and_direction(self): np.testing.assert_allclose(out["dir"], 300.0) np.testing.assert_allclose(out["tp"], 7.0) + def test_inputs(self): + """The wind sea and K inputs on (time, sites); K as the contribution's.""" + + corr = _fitted(lambda h: np.full_like(h, 1.5)) + ws = wind_sea_inputs(_cube(), corr, _fetch()) + self.assertEqual(ws["k"].dims, ("time", "sites")) + np.testing.assert_allclose(ws["hs_ws"] ** 2, ws["e_ws"], rtol=1e-6) + np.testing.assert_allclose(ws["hs_ws_u2"], 9.81 * ws["hs_ws"] / ws["u10"] ** 2, rtol=1e-6) + np.testing.assert_allclose(ws["tp_ws"], 7.0) + out = wind_sea_contribution(_cube(), corr, _fetch()) + np.testing.assert_allclose(out["k"].squeeze(["goal", "partition"]), ws["k"]) + self.assertIsNone(wind_sea_inputs(_cube(weight=0.0), corr, _fetch())) + def test_no_wind_sea(self): """No weighted contribution: nothing to add.""" diff --git a/tests/wrappers/test_snapwave_wrapper.py b/tests/wrappers/test_snapwave_wrapper.py index 338eee5..e43f6c8 100644 --- a/tests/wrappers/test_snapwave_wrapper.py +++ b/tests/wrappers/test_snapwave_wrapper.py @@ -14,13 +14,14 @@ import pandas as pd import xarray as xr from matplotlib import pyplot as plt -from shapely.geometry import box +from shapely.geometry import Point from bluemath_tk.wrappers.snapwave import ( SnapWaveDynamicModelWrapper, SnapWaveMetaModelWrapper, boundary_time_series, - build_enclosure_polygon, + open_boundary_enclosure, + read_convergence, read_boundary_nodes, read_enclosure_polygon, read_forcing_table, @@ -32,8 +33,12 @@ POINTS = pd.DataFrame({"lon": [0.2, 0.4], "lat": [0.3, 0.5]}) -def _fake_his(case_dir: str, n_times: int = 1) -> None: - """Write a minimal SnapWave history file into *case_dir*.""" +def _fake_his(case_dir: str, n_times: int = 1, error: float = 1e-4) -> None: + """Write a minimal SnapWave history file and run log into *case_dir*.""" + + with open(op.join(case_dir, "wrapper_out.log"), "w") as f: + f.write(f" iteration 1 error = 0.05000 %ok = 40.00\n") + f.write(f" iteration 2 error = {error:10.5f} %ok = 100.00\n") shape = (n_times, len(POINTS)) ds = xr.Dataset( @@ -71,7 +76,7 @@ def setUp(self): fixed_parameters={"gridfile": "mesh.nc"}, output_dir=op.join(self.test_dir, "cases"), boundary_nodes=NODES, - enclosure_polygon=build_enclosure_polygon(box(0, 0, 1, 1)), + enclosure_polygon=open_boundary_enclosure(NODES), output_points=POINTS, debug=False, ) @@ -111,6 +116,23 @@ def test_postprocess_on_case_num(self): self.assertIn("hs", ds) self.assertNotIn("point_hm0", ds) + def test_unconverged_case_dropped(self): + """A case whose last error is above max_error is left out.""" + + for i, case_dir in enumerate(self.wrapper.cases_dirs): + _fake_his(case_dir, error=0.9 if i == 1 else 1e-4) + self.assertAlmostEqual(read_convergence(self.wrapper.cases_dirs[1]), 0.9) + ds = self.wrapper.postprocess_cases(write_output_nc=False) + self.assertEqual(list(ds.case_num.values), [0, 2]) + + def test_enclosure_holds_only_the_open_boundary(self): + """The enclosure contains the boundary nodes and nothing a segment away.""" + + polygon = read_enclosure_polygon(self.wrapper.cases_dirs[0]) + self.assertTrue(all(polygon.buffer(1e-9).contains(Point(p)) for p in NODES)) + self.assertFalse(polygon.contains(Point(0.5, 0.5))) + self.assertFalse(polygon.contains(Point(0.0, 0.5))) + def test_slurm_array(self): """The SLURM script reads case_dirs.txt and loads the cluster module.""" @@ -159,7 +181,7 @@ def setUp(self): fixed_parameters={"gridfile": "mesh.nc"}, output_dir=op.join(self.test_dir, "cases"), boundary_nodes=NODES[:2], - enclosure_polygon=build_enclosure_polygon(box(0, 0, 1, 1)), + enclosure_polygon=open_boundary_enclosure(NODES), output_points=points, point_dim="sites", point_coords={"sites": "name", "site_depth": "depth"}, @@ -246,7 +268,7 @@ def test_boundary_time_series_and_cases(self): fixed_parameters={"gridfile": "mesh.nc"}, output_dir=op.join(self.test_dir, "cases"), boundary_nodes=NODES, - enclosure_polygon=build_enclosure_polygon(box(0, 0, 1, 1)), + enclosure_polygon=open_boundary_enclosure(NODES), output_points=POINTS.to_numpy(), debug=False, ) @@ -259,8 +281,8 @@ def test_boundary_time_series_and_cases(self): with open(op.join(case_dir, "snapwave.inp")) as f: self.assertIn("tref = 20200101 010000", f.read()) - for d in wrapper.cases_dirs: - _fake_his(d) + for i, d in enumerate(wrapper.cases_dirs): + _fake_his(d, error=0.9 if i == 1 else 1e-4) joined = wrapper.join_postprocessed_files( [ wrapper.postprocess_case(i, d, {}) @@ -268,6 +290,8 @@ def test_boundary_time_series_and_cases(self): ] ) self.assertEqual(joined.sizes["time"], 2) + self.assertTrue(np.isfinite(joined["hs"].isel(time=0)).all()) + self.assertTrue(np.isnan(joined["hs"].isel(time=1)).all()) fig = plot_case(case_dir, mode="dynamic") self.assertEqual(len(fig.axes) >= 5, True) @@ -300,7 +324,7 @@ def test_wind_field_samples(self): fixed_parameters={"gridfile": "mesh.nc"}, output_dir=op.join(self.test_dir, "cases"), boundary_nodes=NODES, - enclosure_polygon=build_enclosure_polygon(box(0, 0, 1, 1)), + enclosure_polygon=open_boundary_enclosure(NODES), output_points=POINTS.to_numpy(), wind=wind, debug=False, diff --git a/tests/wrappers/test_swan_boundary_wrapper.py b/tests/wrappers/test_swan_boundary_wrapper.py index 3589ac5..7d90460 100644 --- a/tests/wrappers/test_swan_boundary_wrapper.py +++ b/tests/wrappers/test_swan_boundary_wrapper.py @@ -199,7 +199,7 @@ def test_case_files(self): ], ) self.assertIn("SET LEVEL=1.0", inp) - self.assertIn("TABLE 'sites' NOHEAD 'output_sites.tab' HSIGN TPS DIR DSPR", inp) + self.assertIn("TABLE 'sites' NOHEAD 'output_sites.tab' HSIGN TPS TMM10 DIR DSPR", inp) self.assertIn( "BLOCK 'COMPGRID' NOHEAD 'output_map.mat' LAY 3 HSIGN TPS DIR DSPR DEPTH", inp, @@ -227,7 +227,7 @@ def test_postprocess_split(self): for case_dir in self.wrapper.cases_dirs: with open(op.join(case_dir, "output_sites.tab"), "w") as f: - f.write(" 1.0 8.0 270.0 20.0\n 0.5 8.1 275.0 15.0\n") + f.write(" 1.0 8.0 7.0 270.0 20.0\n 0.5 8.1 7.1 275.0 15.0\n") written = self.wrapper.postprocess_cases() ds = xr.open_dataset(written[2]) self.assertEqual(list(ds.sites.values), ["a", "b"]) @@ -266,16 +266,21 @@ def test_hourly_cases(self): grid_file=grid, boundary_nodes=GOALS, output_points=POINTS, + spectra_points=POINTS.iloc[:1], debug=False, ) wrapper.build_cases() with open(op.join(wrapper.cases_dirs[1], "INPUT")) as f: inp = f.read() self.assertIn("SET LEVEL=1.5", inp) + self.assertIn("SPECOUT 'buoys' SPEC2D ABS 'output_spectra.spc'", inp) + spectra_points = op.join(wrapper.cases_dirs[1], "output_spectra.txt") + with open(spectra_points) as f: + self.assertEqual(len(f.read().splitlines()), 1) self.assertIn("0.40000000 2.0000 10.0000 281.000 21.000 &", inp) for case_dir in wrapper.cases_dirs: with open(op.join(case_dir, "output_sites.tab"), "w") as f: - f.write(" 1.0 8.0 270.0 20.0\n 0.5 8.1 275.0 15.0\n") + f.write(" 1.0 8.0 7.0 270.0 20.0\n 0.5 8.1 7.1 275.0 15.0\n") ds = wrapper.postprocess_cases() self.assertEqual(list(pd.DatetimeIndex(ds.time.values)), list(time)) self.assertEqual(dict(ds.hs.sizes), {"time": 2, "points": 2}) @@ -345,7 +350,7 @@ def test_wind_field_and_uniform_wind(self): f"INPGRID WIND REGULAR {context['wind_grid']}\nREADINP WIND 1 'wind.dat' 3 0 FREE\n" ) with open(op.join(case_dir, "output_sites.tab"), "w") as f: - f.write(" 1.0 8.0 270.0 20.0\n 0.5 8.1 275.0 15.0\n") + f.write(" 1.0 8.0 7.0 270.0 20.0\n 0.5 8.1 7.1 275.0 15.0\n") case = load_swan_case(case_dir) self.assertEqual(case.boundary.shape[1], 5) self.assertEqual(len(case.wind[0]), 6) From e429cabc8c43a4fba24a967125587bcec138f9da Mon Sep 17 00:00:00 2001 From: tausiaj Date: Fri, 6 Nov 2026 12:57:43 +0100 Subject: [PATCH 2/2] [JTH] remove the hywaves wind correction (k growth factor) The wind correction did not work with the open-boundary enclosure; what was tried and the options are in HyGloW hywaves/WIND.md. Restore point: e280155. Co-Authored-By: Claude Opus 5.5 --- bluemath_tk/waves/hywaves/__init__.py | 2 - bluemath_tk/waves/hywaves/plots.py | 330 +------------ bluemath_tk/waves/hywaves/reconstruction.py | 50 +- bluemath_tk/waves/hywaves/resources.py | 5 +- bluemath_tk/waves/hywaves/wind_correction.py | 484 ------------------- tests/waves/test_hywaves_plots.py | 50 -- tests/waves/test_wind_correction.py | 156 ------ 7 files changed, 12 insertions(+), 1065 deletions(-) delete mode 100644 bluemath_tk/waves/hywaves/wind_correction.py delete mode 100644 tests/waves/test_wind_correction.py diff --git a/bluemath_tk/waves/hywaves/__init__.py b/bluemath_tk/waves/hywaves/__init__.py index 86a5741..b690367 100644 --- a/bluemath_tk/waves/hywaves/__init__.py +++ b/bluemath_tk/waves/hywaves/__init__.py @@ -21,7 +21,6 @@ resources, sectors, store, - wind_correction, ) __all__ = [ @@ -33,5 +32,4 @@ "resources", "sectors", "store", - "wind_correction", ] diff --git a/bluemath_tk/waves/hywaves/plots.py b/bluemath_tk/waves/hywaves/plots.py index 6ec4884..acf1ef0 100644 --- a/bluemath_tk/waves/hywaves/plots.py +++ b/bluemath_tk/waves/hywaves/plots.py @@ -11,7 +11,7 @@ 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`); +- ERA5 partition periods (:func:`plot_peak_period_ratio`); - the MDA-size sweep of the holdout metrics (:func:`plot_mda_sweep`). Goals are mappings ``{goal_id: {"lon", "lat", "angles": [left, right]}}``, @@ -323,7 +323,6 @@ def plot_dynamic_forcing( *, 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", @@ -335,8 +334,7 @@ def plot_dynamic_forcing( 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. + for hours with wind sea; water-level envelope over the goals. Parameters ---------- @@ -350,8 +348,6 @@ def plot_dynamic_forcing( ``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 @@ -437,13 +433,7 @@ def empty(ax: Axes, text: str) -> None: 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)) + ax_angle.set(title="Wind − wave direction, wind-sea hours", 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") @@ -466,64 +456,6 @@ def empty(ax: Axes, text: str) -> None: 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 - - def plot_peak_period_ratio( samples: pd.DataFrame, theory: Mapping[str, float], title: str = "Tp / Tm-1,0" ) -> Figure: @@ -562,268 +494,12 @@ def plot_peak_period_ratio( return fig -def _direction_arrow(ax: Axes, direction: float, xy: tuple[float, float], color: str, label: str) -> None: - """Arrow towards where waves / wind go, from a nautical coming-from *direction* (axes fraction).""" - rad = np.deg2rad(direction) - dx, dy = -0.1 * np.sin(rad), -0.1 * np.cos(rad) - start = (xy[0] - dx / 2, xy[1] - dy / 2) - ax.annotate("", xy=(start[0] + dx, start[1] + dy), xytext=start, xycoords="axes fraction", - arrowprops={"arrowstyle": "-|>", "color": color, "lw": 2}) - ax.text(xy[0], xy[1] - 0.1, label, transform=ax.transAxes, ha="center", va="top", fontsize=8, color=color) - - -def plot_wind_case_comparison(ds: xr.Dataset, size: float = 4) -> Figure: - """ - Training cases with and without wind, and the wind correction, at every site. - - One row per case: what the wind adds in the runs (with - without wind), - what the correction adds (corrected - without wind), on the same scale; - their difference (corrected - with wind) with its bias and RMSE; and the - run with wind (x, reference) against the run without wind and the - corrected one. Arrows: waves (black) and wind (red), towards where they - go. - - Parameters - ---------- - ds : xr.Dataset - ``hs0``, ``hs1``, ``hs_corrected`` on ``(case, sites)``, with site - ``longitude`` / ``latitude`` and case ``hs``, ``tp``, ``u10``, - ``dir``, ``u10dir``, ``wind_wave_angle``, ``ws_weight``, - ``hs_ws_u2`` coordinates. - size : float, optional - Panel size (inches). Default is 4. - - Returns - ------- - Figure - """ - - n = ds.sizes["case"] - fig, axes = plt.subplots(n, 4, figsize=(4 * size, n * size), squeeze=False, constrained_layout=True) - lon, lat = ds["longitude"].values, ds["latitude"].values - aspect = 1.0 / np.cos(np.deg2rad(np.nanmean(lat))) - for row, case in zip(axes, ds["case"].values): - c = ds.sel(case=case) - added = {"wind added, runs": (c["hs1"] - c["hs0"]).values, - "wind added, correction": (c["hs_corrected"] - c["hs0"]).values} - lim = max(float(np.nanpercentile(np.abs(np.r_[tuple(added.values())]), 99)), 0.05) - for ax, (label, value) in zip(row, added.items()): - sc = ax.scatter(lon, lat, c=value, s=2, cmap="RdBu_r", vmin=-lim, vmax=lim, rasterized=True) - ax.set_title(f"{label} (median {np.nanmedian(value):+.2f} m)", fontsize=10) - fig.colorbar(sc, ax=row[:2], label="Hs difference (m)", shrink=0.8) - err = (c["hs_corrected"] - c["hs1"]).values - sc = row[2].scatter(lon, lat, c=err, s=2, cmap="RdBu_r", vmin=-lim, vmax=lim, rasterized=True) - fig.colorbar(sc, ax=row[2], label="corrected - with wind (m)", shrink=0.8) - row[2].set_title(f"bias {np.nanmean(err):+.2f} m, RMSE {np.sqrt(np.nanmean(err**2)):.2f} m", fontsize=10) - for ax in row[:3]: - ax.set_aspect(aspect) - ax.tick_params(labelsize=7) - _direction_arrow(row[0], float(c["dir"]), (0.12, 0.88), "black", "waves") - _direction_arrow(row[0], float(c["u10dir"]), (0.30, 0.88), "#d62728", "wind") - row[0].set_ylabel( - f"case {case}: Hs {float(c['hs']):.1f} m, Tp {float(c['tp']):.1f} s\n" - f"U {float(c['u10']):.1f} m/s, angle {float(c['wind_wave_angle']):+.0f}°, " - f"w {float(c['ws_weight']):.2f}, gHs/U² {float(c['hs_ws_u2']):.2f}", - fontsize=9, - ) - ax = row[3] - ref = c["hs1"].values - ax.scatter(ref, c["hs0"].values, s=2, color="0.6", alpha=0.5, label="without wind", rasterized=True) - ax.scatter(ref, c["hs_corrected"].values, s=2, color="#1f77b4", alpha=0.5, label="corrected", rasterized=True) - top = float(np.nanmax(np.r_[ref, c["hs_corrected"].values])) - ax.plot([0, top], [0, top], "k--", lw=1) - ax.set(xlim=(0, top), ylim=(0, top), xlabel="Hs with wind (m)", ylabel="Hs (m)", title="all sites") - ax.legend(fontsize=8, markerscale=4, loc="upper left") - return fig - - -def plot_k_response( - ds: xr.Dataset, ranges: Mapping[str, tuple[float, float]] | None = None, title: str = "K response" -) -> Figure: - """ - ``K`` of the wind correction against each input, one line per site. - - Parameters - ---------- - ds : xr.Dataset - ``k_u10``, ``k_angle``, ``k_hs_ws_u2`` on ``(sites, )``, site - ``depth`` and ``fetch`` (m) coordinates and the base state - (``wind_dir``, ``hs_ws``, ``u10``, ``angle``) in the attributes. - ranges : Mapping[str, tuple of float], optional - ``u10``, ``angle``, ``hs_ws_u2`` -> range of the training cases; - outside it is shaded (the model keeps its edge value there). - title : str, optional - Figure title. - - Returns - ------- - Figure - """ - - panels = { - "u10": ("wind speed U10 (m/s)", f"Hs_ws {ds.attrs['hs_ws']:g} m, angle {ds.attrs['angle']:g}°"), - "angle": ("wind - wave angle (°)", f"Hs_ws {ds.attrs['hs_ws']:g} m, U {ds.attrs['u10']:g} m/s"), - "hs_ws_u2": ("g Hs_ws / U² (how developed the wind sea is)", - f"U {ds.attrs['u10']:g} m/s, angle {ds.attrs['angle']:g}°"), - } - fig, axes = plt.subplots(1, 3, figsize=(17, 4.8), sharey=True, constrained_layout=True) - fig.suptitle(f"{title} · wind from {ds.attrs['wind_dir']:g}°", fontsize=14, fontweight="bold") - for ax, (name, (xlabel, base)) in zip(axes, panels.items()): - x = ds[name].values - for site in ds["sites"].values: - s = ds.sel(sites=site) - ax.plot(x, s[f"k_{name}"], lw=2, - label=f"{site} ({float(s['depth']):.0f} m, fetch {float(s['fetch']) / 1e3:.0f} km)") - if ranges is not None and name in ranges: - low, high = ranges[name] - ax.axvspan(x.min(), low, color="0.85", zorder=0) - ax.axvspan(high, x.max(), color="0.85", zorder=0, label="outside training") - if name == "hs_ws_u2": - ax.set_xscale("log") - ax.axvline(0.24, color="0.4", ls=":", lw=1) - ax.text(0.24, 1.0, " fully developed", transform=ax.get_xaxis_transform(), va="top", fontsize=8, color="0.4") - ax.axhline(1, color="k", lw=0.8) - ax.set(xlabel=xlabel, xlim=(x.min(), x.max()), title=base) - axes[0].set_ylabel("K = Hs with / without wind") - axes[0].legend(fontsize=8, loc="upper left") - return fig - - -def plot_wind_sea_event( - ws: xr.Dataset, - hs: Mapping[str, xr.DataArray], - ranges: Mapping[str, tuple[float, float]] | None = None, - title: str = "", -) -> Figure: - """ - The wind correction at one site through an event. - - Rows: Hs of every source; the wind sea the correction sees (``hs_ws``) - and the wind (``u10``); how developed the wind sea is - (``g Hs_ws / U^2``); the wind-wave angle; ``K``. Shaded: outside the - training cases (*ranges*). - - Parameters - ---------- - ws : xr.Dataset - ``hs_ws``, ``u10``, ``hs_ws_u2``, ``angle``, ``k`` on ``time`` (one - site of ``wind_sea_inputs``). - hs : Mapping[str, xr.DataArray] - Label -> Hs on ``time`` (e.g. buoy, dynamic run, metamodel without and - with the correction). - ranges : Mapping[str, tuple of float], optional - ``hs_ws_u2``, ``angle`` -> training range. - title : str, optional - Figure title. - - Returns - ------- - Figure - """ - - fig, axes = plt.subplots(5, 1, figsize=(13, 13), sharex=True, constrained_layout=True, - gridspec_kw={"height_ratios": [2, 1, 1, 1, 1]}) - fig.suptitle(title, fontsize=14, fontweight="bold") - times = ws["time"].values - ax = axes[0] - for label, series in hs.items(): - marker = "o" if label.lower().startswith("buoy") else None - ax.plot(series["time"], series, lw=1.5, marker=marker, ms=2.5, ls="-" if marker is None else "none", label=label) - ax.set(ylabel="Hs (m)") - ax.legend(fontsize=9, ncol=2) - ax = axes[1] - ax.plot(times, ws["hs_ws"], color="#2ca02c", lw=1.5, label="wind sea Hs_ws (m)") - ax.set_ylabel("Hs_ws (m)", color="#2ca02c") - twin = ax.twinx() - twin.plot(times, ws["u10"], color="#d62728", lw=1.2, label="U10") - twin.set_ylabel("U10 (m/s)", color="#d62728") - for ax, name, label, log in ((axes[2], "hs_ws_u2", "g Hs_ws / U²", True), (axes[3], "angle", "wind - wave (°)", False)): - values = ws[name].values if name != "angle" else (ws[name].values + 180.0) % 360.0 - 180.0 - ax.plot(times, values, color="k", lw=1.2) - if ranges is not None and name in ranges: - low, high = ranges[name] - ax.axhspan(low, high, color="#1f77b4", alpha=0.12, label="training range") - ax.legend(fontsize=8, loc="upper left") - if log: - ax.set_yscale("log") - ax.axhline(0.24, color="0.4", ls=":", lw=1) - ax.set_ylabel(label) - axes[4].plot(times, ws["k"], color="#9467bd", lw=1.5) - axes[4].axhline(1, color="k", lw=0.8) - axes[4].set_ylabel("K") - axes[4].xaxis.set_major_formatter(mdates.DateFormatter("%d %b %H:%M")) - return fig - - -def plot_event_maps( - reference: xr.DataArray, - estimates: Mapping[str, xr.DataArray], - k: xr.DataArray | None = None, - reference_label: str = "reference", - title: str = "", - size: float = 4.5, -) -> Figure: - """ - Hs at every site at one time: a reference, each estimate minus it, and ``K``. - - Parameters - ---------- - reference : xr.DataArray - Hs on ``sites`` with ``site_longitude`` / ``site_latitude``. - estimates : Mapping[str, xr.DataArray] - Label -> Hs on the same sites. - k : xr.DataArray, optional - ``K`` of the wind correction on the same sites. - reference_label : str, optional - Name of the reference. - title : str, optional - Figure title. - size : float, optional - Panel size (inches). - - Returns - ------- - Figure - """ - - n = 1 + len(estimates) + int(k is not None) - fig, axes = plt.subplots(1, n, figsize=(n * size, size * 1.15), constrained_layout=True, squeeze=False) - axes = axes[0] - fig.suptitle(title, fontsize=14, fontweight="bold") - lon, lat = reference["site_longitude"].values, reference["site_latitude"].values - sc = axes[0].scatter(lon, lat, c=reference.values, s=3, cmap="viridis", vmin=0, rasterized=True) - fig.colorbar(sc, ax=axes[0], label="Hs (m)", shrink=0.8) - axes[0].set_title(f"{reference_label} (median {float(reference.median()):.2f} m)", fontsize=10) - diffs = {label: (est.sel(sites=reference["sites"]) - reference).values for label, est in estimates.items()} - lim = max(float(np.nanpercentile(np.abs(np.concatenate(list(diffs.values()))), 99)), 0.05) if diffs else 1.0 - for ax, (label, diff) in zip(axes[1:], diffs.items()): - sc = ax.scatter(lon, lat, c=diff, s=3, cmap="RdBu_r", vmin=-lim, vmax=lim, rasterized=True) - ax.set_title(f"{label} - {reference_label}\nbias {np.nanmean(diff):+.2f}, RMSE {np.sqrt(np.nanmean(diff**2)):.2f} m", - fontsize=10) - if diffs: - fig.colorbar(sc, ax=axes[1 : 1 + len(diffs)], label="Hs difference (m)", shrink=0.8) - if k is not None: - sc = axes[-1].scatter(lon, lat, c=k.sel(sites=reference["sites"]).values, s=3, cmap="magma_r", vmin=1, rasterized=True) - fig.colorbar(sc, ax=axes[-1], label="K", shrink=0.8) - axes[-1].set_title("K of the wind correction", fontsize=10) - aspect = 1.0 / np.cos(np.deg2rad(np.nanmean(lat))) - for ax in axes: - ax.set_aspect(aspect) - ax.tick_params(labelsize=7) - return fig - - __all__ = [ "SWEEP_METRICS", "goal_colors", "plot_direction_ranges", "plot_dynamic_forcing", - "plot_event_maps", "plot_goal_sectors", - "plot_k_response", "plot_mda_sweep", "plot_peak_period_ratio", - "plot_wind_case_comparison", - "plot_wind_k_comparison", - "plot_wind_sea_event", ] diff --git a/bluemath_tk/waves/hywaves/reconstruction.py b/bluemath_tk/waves/hywaves/reconstruction.py index 43529bf..cdd0b43 100644 --- a/bluemath_tk/waves/hywaves/reconstruction.py +++ b/bluemath_tk/waves/hywaves/reconstruction.py @@ -6,9 +6,7 @@ :func:`~bluemath_tk.waves.hywaves.metamodel.predict_fields`). Every contribution becomes a JONSWAP x Cartwright spectrum, the spectra are summed over goals and partitions, and bulk parameters are computed from the sum -(:func:`sum_partition_spectra`). With a wind correction, one more -contribution, the local wind sea, is added per block before the sum -(see :mod:`.wind_correction`). +(:func:`sum_partition_spectra`). 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* @@ -33,7 +31,6 @@ import logging import time from collections.abc import Mapping -from typing import Any import numpy as np import xarray as xr @@ -190,8 +187,8 @@ def sum_partition_spectra( from wavespectra.construct import construct_partition parts = cube.sel(sites=site_ids).stack(part=("goal", "partition")) - # contributions never predicted (e.g. a goal x partition combination - # created by merging, such as the wind-sea contribution) carry no energy + # contributions never predicted (a goal x partition combination created + # by merging) carry no energy present = parts.hs.notnull().any([d for d in parts.hs.dims if d != "part"]) parts = parts.isel(part=np.flatnonzero(present.values)) # float32 halves the peak array (twice the sites per batch); the relative @@ -232,29 +229,12 @@ def sum_partition_spectra( return result, spectra -def _add_wind_sea( - merged: xr.Dataset, wind_correction: Mapping[str, Any] -) -> tuple[xr.Dataset, xr.DataArray | None]: - """Merge the wind-sea contribution of one site batch; also return its ``K``.""" - - from .wind_correction import wind_sea_contribution - - contrib = wind_sea_contribution(merged, **wind_correction) - if contrib is None: - return merged, None - k = contrib["k"].squeeze(["goal", "partition"], drop=True) - - return xr.merge([merged, contrib.drop_vars("k")]), k - - def _time_axis(chunks: list[xr.Dataset]) -> np.ndarray: return np.unique(np.concatenate([c["time"].values for c in chunks])) -def _n_parts(chunks: list[xr.Dataset], wind_correction: Any) -> int: - return sum(c.sizes["goal"] * c.sizes["partition"] for c in chunks) + int( - wind_correction is not None - ) +def _n_parts(chunks: list[xr.Dataset]) -> int: + return sum(c.sizes["goal"] * c.sizes["partition"] for c in chunks) def _sum_blocks( @@ -265,12 +245,11 @@ def _sum_blocks( budget_bytes: int, load: bool = True, label: str = "", - wind_correction: Mapping[str, Any] | None = None, ) -> tuple[xr.Dataset, xr.Dataset | None]: """Merge contribution chunks and sum them, in site x time blocks within the budget.""" times = _time_axis(chunks) - n_sites, n_steps = batch_shape(len(times), _n_parts(chunks, wind_correction), budget_bytes) + n_sites, n_steps = batch_shape(len(times), _n_parts(chunks), budget_bytes) n_blocks = -(-len(site_ids) // n_sites) * -(-len(times) // n_steps) results, spectras = [], [] t0 = time.monotonic() @@ -284,14 +263,9 @@ def _sum_blocks( for j in range(0, merged.sizes["time"], n_steps): t_block = time.monotonic() block = merged.isel(time=slice(j, j + n_steps)) - wind_k = None - if wind_correction is not None: - block, wind_k = _add_wind_sea(block, wind_correction) result, spectra = sum_partition_spectra( block, sites, variables, save_spectra=save_spectra, load=load ) - if wind_k is not None: - result["wind_k"] = wind_k site_results.append(result) if save_spectra and spectra is not None: site_spectra.append(spectra) @@ -361,7 +335,6 @@ def sum_partition_spectra_batched( save_spectra: bool = False, load: bool = True, log_file: str | None = None, - wind_correction: Mapping[str, Any] | None = None, ) -> tuple[xr.Dataset, xr.Dataset | None]: """ :func:`sum_partition_spectra` over many sites and steps, in memory-sized blocks. @@ -389,13 +362,6 @@ def sum_partition_spectra_batched( Load results into memory (always True in workers). log_file : str, optional Log file for worker progress messages. - wind_correction : mapping, optional - Keyword arguments of - :func:`~bluemath_tk.waves.hywaves.wind_correction.wind_sea_contribution` - (``correction``, ``fetch``, optional ``wind``, ``spr``, ``tp``): adds - the local wind sea per block, and ``wind_k`` (``K``, 1 where there is - no correction) to the bulk output. Chunks then also carry - ``ws_weight`` (and ``u10``, ``u10dir`` without ``wind``). Returns ------- @@ -405,8 +371,7 @@ def sum_partition_spectra_batched( if workers <= 1: return _sum_blocks( - chunks, site_ids, variables, save_spectra, budget_bytes, load, - wind_correction=wind_correction, + chunks, site_ids, variables, save_spectra, budget_bytes, load ) import multiprocessing @@ -442,7 +407,6 @@ def sum_partition_spectra_batched( per_worker, True, f"item {k + 1}/{len(items)}", - wind_correction, ) futures[future] = k for future in as_completed(futures): diff --git a/bluemath_tk/waves/hywaves/resources.py b/bluemath_tk/waves/hywaves/resources.py index ce8d6ba..6a8ff1a 100644 --- a/bluemath_tk/waves/hywaves/resources.py +++ b/bluemath_tk/waves/hywaves/resources.py @@ -191,8 +191,7 @@ def call_memory( months, sites : int Calendar months and sites of the call. parts : int - Contributions summed per site (goals x partitions, + 1 with the wind - correction). + Contributions summed per site (goals x partitions). workers : int Sum processes. save_spectra : bool, optional @@ -295,7 +294,7 @@ def hours(site_months: int, months: int, parts: int, workers: int) -> float: months : int Calendar months they span. parts : int - Contributions per site (goals x partitions, + 1 with wind). + Contributions per site (goals x partitions). workers : int Sum processes. diff --git a/bluemath_tk/waves/hywaves/wind_correction.py b/bluemath_tk/waves/hywaves/wind_correction.py deleted file mode 100644 index 9f52833..0000000 --- a/bluemath_tk/waves/hywaves/wind_correction.py +++ /dev/null @@ -1,484 +0,0 @@ -""" -HyWaves wind correction: local wind growth added to the reconstructed wind sea. - -The metamodel cases run without wind, so the reconstruction misses what the -local wind adds while the waves cross the domain. The correction is applied -once per time step and site, after the goals are predicted and before their -spectra are summed (see :func:`wind_sea_contribution`): - -1. Each goal x partition contribution gets a wind-sea weight ``w`` from its - *offshore* forcing (:func:`~bluemath_tk.waves.wind.wind_sea_weight`), and - the wind it implies (:func:`~bluemath_tk.waves.wind.infer_wind`). -2. The wind-sea energy at the site, ``E_ws = sum w hs^2``, is summed once over - the goals (no goal counted twice). -3. The wind at the site is the data (e.g. ERA5) or the ``w hs^2``-weighted - inferred wind of the contributions. -4. ``K = Hs(with wind) / Hs(without wind)`` comes from :class:`WindCorrection`, - fitted on pairs of full-domain runs with and without wind, on - dimensionless features of the wind sea, the fetch and the depth - (:func:`correction_features`). -5. One extra contribution carries the missing energy, ``(K^2 - 1) E_ws``, - from the wind direction; it is summed with the others as a spectrum. - -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 - -import pickle -from collections.abc import Mapping - -import numpy as np -import pandas as pd -import xarray as xr - -from ...core.constants import GRAVITY - -#: Partition id of the wind-sea contribution added by the correction. -WIND_PARTITION = -2 -#: Goal id of the wind-sea contribution (it belongs to no goal). -WIND_GOAL = 0 -#: Features of :func:`correction_features`, in order. -FEATURES = ("hs_ws", "fetch", "depth", "mean_depth", "cos_angle") - - -def correction_features( - hs_ws: np.ndarray, - u10: np.ndarray, - fetch: np.ndarray, - depth: np.ndarray, - mean_depth: np.ndarray, - angle: np.ndarray, -) -> np.ndarray: - """ - Dimensionless features of the wind correction. - - ``g Hs_ws / U^2`` (how developed the wind sea already is), ``g F / U^2`` - (fetch left), ``g d / U^2`` at the site and along the fetch (depth - limitation) and the cosine of the wind-wave angle. - - Parameters - ---------- - hs_ws : np.ndarray - Wind-sea Hs at the site without wind (m). - u10 : np.ndarray - Wind speed (m/s), positive. - fetch : np.ndarray - Fetch along the wind direction (m). - depth : np.ndarray - Depth at the site (m). - mean_depth : np.ndarray - Mean depth along the fetch (m). - angle : np.ndarray - Wind direction minus wind-sea direction (deg). - - Returns - ------- - np.ndarray - ``(n, 5)`` features, in the order of :data:`FEATURES`. - """ - - u2 = np.asarray(u10, dtype=float) ** 2 / GRAVITY - cols = [ - np.asarray(hs_ws, dtype=float) / u2, - np.asarray(fetch, dtype=float) / u2, - np.asarray(depth, dtype=float) / u2, - np.asarray(mean_depth, dtype=float) / u2, - np.cos(np.deg2rad(np.asarray(angle, dtype=float))), - ] - - return np.column_stack([np.ravel(c) for c in np.broadcast_arrays(*cols)]) - - -class WindCorrection: - """ - ``K = Hs(with wind) / Hs(without wind)`` from :func:`correction_features`. - - A gradient-boosting regression of ``log K`` (scikit-learn), fitted on - pairs of the same full-domain cases run with and without wind. ``K`` - scales the wind-sea part of the energy only (see :meth:`fit`). - - Attributes - ---------- - model : sklearn.ensemble.HistGradientBoostingRegressor - Fitted regressor of ``log K``. - u10_min : float - Below this wind speed (m/s) ``K = 1``. - hs_min : float - Below this wind-sea Hs (m) ``K = 1``. - k_range : tuple of float - Bounds on the predicted ``K``. - """ - - def __init__( - self, - u10_min: float = 1.0, - hs_min: float = 0.05, - k_range: tuple[float, float] = (0.5, 4.0), - **regressor_kwargs, - ) -> None: - """ - Parameters - ---------- - u10_min : float, optional - Wind speed (m/s) below which there is no correction. Default 1. - hs_min : float, optional - Wind-sea Hs (m) below which there is no correction. Default 0.05. - k_range : tuple of float, optional - Bounds on the predicted ``K``. Default ``(0.5, 4.0)``. - **regressor_kwargs - Passed to ``HistGradientBoostingRegressor`` (default - ``max_iter=400, learning_rate=0.05``). - """ - - from sklearn.ensemble import HistGradientBoostingRegressor - - self.u10_min = float(u10_min) - self.hs_min = float(hs_min) - self.k_range = (float(k_range[0]), float(k_range[1])) - self.model = HistGradientBoostingRegressor( - **{"max_iter": 400, "learning_rate": 0.05, **regressor_kwargs} - ) - self.n_samples = 0 - - def _valid(self, hs_ws: np.ndarray, u10: np.ndarray) -> np.ndarray: - return (np.asarray(u10) >= self.u10_min) & (np.asarray(hs_ws) >= self.hs_min) - - def fit( - self, - hs_without: np.ndarray, - hs_with: np.ndarray, - u10: np.ndarray, - fetch: np.ndarray, - depth: np.ndarray, - mean_depth: np.ndarray, - angle: np.ndarray, - weight: float | np.ndarray = 1.0, - ) -> WindCorrection: - """ - Fit on paired runs: site Hs without and with wind, and their conditions. - - With a wind-sea weight ``w`` of the case forcing, only ``w hs^2`` of - the energy is wind sea, as in the reconstruction: the feature is - ``Hs_ws = sqrt(w) hs_without`` and the target the ``K`` of that part, - ``K^2 = 1 + (hs_with^2 - hs_without^2) / (w hs_without^2)``. Pairs - with NaNs, no wind (``u10 < u10_min``) or no wind sea - (``Hs_ws < hs_min``) are left out. - - Parameters - ---------- - hs_without, hs_with : np.ndarray - Site Hs (m) of the run without and with wind. - u10, fetch, depth, mean_depth, angle : np.ndarray - See :func:`correction_features`. - weight : float or np.ndarray, optional - Wind-sea weight of each pair's forcing, in ``(0, 1]``. Default 1. - - Returns - ------- - WindCorrection - Self. - """ - - shape = np.broadcast_shapes( - *(np.shape(a) for a in (hs_without, hs_with, u10, fetch, depth, mean_depth, angle, weight)) - ) - hs0 = np.ravel(np.broadcast_to(hs_without, shape)).astype(float) - hs1 = np.ravel(np.broadcast_to(hs_with, shape)).astype(float) - w = np.ravel(np.broadcast_to(weight, shape)).astype(float) - u = np.ravel(np.broadcast_to(u10, shape)).astype(float) - with np.errstate(invalid="ignore", divide="ignore"): - hs_ws = np.sqrt(w) * hs0 - k2 = 1.0 + (hs1**2 - hs0**2) / (w * hs0**2) - x = correction_features( - hs_ws, - u, - *(np.ravel(np.broadcast_to(a, shape)) for a in (fetch, depth, mean_depth, angle)), - ) - ok = ( - np.isfinite(x).all(axis=1) - & np.isfinite(k2) - & (w > 0) - & self._valid(hs_ws, u) - ) - if not ok.any(): - raise ValueError("no valid pairs to fit the wind correction") - k = np.sqrt(np.clip(k2[ok], self.k_range[0] ** 2, self.k_range[1] ** 2)) - self.model.fit(x[ok], np.log(k)) - self.n_samples = int(ok.sum()) - - return self - - def predict( - self, - hs_ws: np.ndarray, - u10: np.ndarray, - fetch: np.ndarray, - depth: np.ndarray, - mean_depth: np.ndarray, - angle: np.ndarray, - ) -> np.ndarray: - """ - ``K`` for each input (1 without wind, wind sea or valid features). - - Parameters - ---------- - hs_ws, u10, fetch, depth, mean_depth, angle : np.ndarray - See :func:`correction_features` (broadcast together). - - Returns - ------- - np.ndarray - ``K``, with the broadcast shape of the inputs. - """ - - shape = np.broadcast_shapes(*(np.shape(a) for a in (hs_ws, u10, fetch, depth, mean_depth, angle))) - x = correction_features(hs_ws, u10, fetch, depth, mean_depth, angle) - hs = np.ravel(np.broadcast_to(hs_ws, shape)) - u = np.ravel(np.broadcast_to(u10, shape)) - k = np.ones(len(x)) - ok = np.isfinite(x).all(axis=1) & self._valid(hs, u) - if ok.any(): - k[ok] = np.clip(np.exp(self.model.predict(x[ok])), *self.k_range) - - return k.reshape(shape) - - def save_model(self, path: str) -> None: - """Pickle the correction (read back with :func:`bluemath_tk.core.io.load_model`).""" - - with open(path, "wb") as f: - pickle.dump(self, f) - - -def _fetch_along(table: xr.Dataset, name: str, direction: xr.DataArray) -> xr.DataArray: - """``table[name]`` on ``(sites, dir)`` at the nearest direction bin of *direction*.""" - - dirs = table["dir"].values - step = 360.0 / len(dirs) - idx = np.mod(np.rint((direction.values - dirs[0]) / step).astype(int), len(dirs)) - values = table[name].transpose("sites", "dir").values - site_axis = direction.dims.index("sites") - rows = np.arange(values.shape[0]).reshape( - [-1 if a == site_axis else 1 for a in range(direction.ndim)] - ) - - return xr.DataArray(values[rows, idx], dims=direction.dims, coords=direction.coords) - - -def _vector_mean(direction: xr.DataArray, weight: xr.DataArray, dims: list[str]) -> xr.DataArray: - """Weighted vector mean of nautical directions (deg) over *dims*.""" - - rad = np.deg2rad(direction) - s = (weight * np.sin(rad)).sum(dims) - c = (weight * np.cos(rad)).sum(dims) - - return np.mod(np.rad2deg(np.arctan2(s, c)), 360.0) - - -def wind_sea_inputs( - cube: xr.Dataset, - correction: WindCorrection, - fetch: xr.Dataset, - wind: xr.Dataset | None = None, -) -> xr.Dataset | None: - """ - The wind sea at each site and step of one batch, the inputs of ``K`` and ``K``. - - Parameters - ---------- - cube, correction, fetch, wind - As for :func:`wind_sea_contribution`. - - Returns - ------- - xr.Dataset or None - On ``(time, sites)``: ``e_ws`` (wind-sea energy ``sum w hs^2``, m²), - ``hs_ws``, ``tp_ws`` and ``ws_dir`` (its Hs, energy-weighted Tp and - direction), ``u10`` and ``u10dir`` (wind), ``angle`` (wind - wave), - ``fetch``, ``mean_depth`` (along the wind) and ``depth`` (site), - ``hs_ws_u2`` (``g Hs_ws / U^2``) and ``k``; None if nothing in the - batch is wind sea. - """ - - parts = ["goal", "partition"] - weight = cube["ws_weight"].fillna(0.0) - energy = (weight * cube["hs"].fillna(0.0) ** 2).transpose("time", "sites", ...) - e_ws = energy.sum(parts) - if not bool((e_ws > 0).any()): - return None - - with np.errstate(invalid="ignore", divide="ignore"): - ws_dir = _vector_mean(cube["dir"].fillna(0.0), energy, parts) - tp_ws = (energy * cube["tp"].fillna(0.0)).sum(parts) / e_ws - if wind is None: - u10 = (energy * cube["u10"].fillna(0.0)).sum(parts) / e_ws - u10dir = _vector_mean(cube["u10dir"].fillna(0.0), energy * cube["u10"].fillna(0.0), parts) - else: - u10 = wind["u10"].sel(time=cube["time"], sites=cube["sites"]) - u10dir = wind["u10dir"].sel(time=cube["time"], sites=cube["sites"]) - u10 = u10.fillna(0.0).transpose("time", "sites") - u10dir = u10dir.transpose("time", "sites") - hs_ws = np.sqrt(e_ws) - - fetch = fetch.sel(sites=cube["sites"].values) - name = "fetch_eff" if "fetch_eff" in fetch else "fetch" - fetch_u = _fetch_along(fetch, name, u10dir) - depth_u = _fetch_along(fetch, "mean_depth", u10dir) - site_depth = fetch["site_depth"].broadcast_like(hs_ws) - angle = u10dir - ws_dir - k = correction.predict( - hs_ws.values, u10.values, fetch_u.values, site_depth.values, depth_u.values, angle.values - ) - with np.errstate(invalid="ignore", divide="ignore"): - hs_ws_u2 = GRAVITY * hs_ws / u10**2 - - return xr.Dataset( - { - "e_ws": e_ws, "hs_ws": hs_ws, "tp_ws": tp_ws, "ws_dir": ws_dir, - "u10": u10, "u10dir": u10dir, "angle": angle, "fetch": fetch_u, - "mean_depth": depth_u, "depth": site_depth.reset_coords(drop=True), "hs_ws_u2": hs_ws_u2, - "k": xr.DataArray(k, dims=hs_ws.dims, coords=hs_ws.coords), - } - ) - - -def wind_sea_contribution( - cube: xr.Dataset, - correction: WindCorrection, - fetch: xr.Dataset, - wind: xr.Dataset | None = None, - spr: float = 30.0, - tp: str = "wind_sea", -) -> xr.Dataset | None: - """ - The extra contribution of the local wind for one batch of sites. - - Parameters - ---------- - cube : xr.Dataset - Contributions on ``(goal, partition, time, sites)``: ``hs``, ``tp``, - ``dir``, ``spr`` (nearshore) plus, from each contribution's offshore - forcing, ``ws_weight`` (wind-sea weight) and, without *wind*, ``u10`` - and ``u10dir`` (inferred wind), on ``(goal, partition, time)``. - correction : WindCorrection - Fitted ``K`` model. - fetch : xr.Dataset - ``fetch_eff`` (or ``fetch``) and ``mean_depth`` on ``(sites, dir)`` - and ``site_depth`` on ``sites`` (see - :func:`bluemath_tk.waves.fetch.fetch_table`), with the cube's sites. - wind : xr.Dataset, optional - ``u10`` and ``u10dir`` on ``(time, sites)`` from data. Default: the - ``w hs^2``-weighted inferred wind of the contributions. - spr : float, optional - Directional spreading of the added wind sea (deg). Default is 30. - tp : str, optional - Period of the added wind sea: ``"wind_sea"`` (energy-weighted Tp of - the wind-sea contributions, default) or ``"growth"`` (Kahma & Calkoen - period after growing over the fetch). - - Returns - ------- - xr.Dataset or None - ``hs``, ``tp``, ``dir``, ``spr`` and ``k`` on - ``(goal, partition, time, sites)`` with goal :data:`WIND_GOAL` and - partition :data:`WIND_PARTITION`, ready to merge with *cube*; None - if nothing in the batch is wind sea. - """ - - ws = wind_sea_inputs(cube, correction, fetch, wind) - if ws is None: - return None - hs_add = np.sqrt(np.clip(ws["k"] ** 2 - 1.0, 0.0, None) * ws["e_ws"]) - - if tp == "growth": - from ..wind import equivalent_fetch, growth_tp - - with np.errstate(invalid="ignore", divide="ignore"): - grown = equivalent_fetch(ws["hs_ws"].values, ws["u10"].values) + ws["fetch"].values - tp_add = growth_tp(ws["u10"].values, np.where(np.isfinite(grown), grown, 1e12), ws["mean_depth"].values) - tp_add = xr.DataArray(tp_add, dims=hs_add.dims, coords=hs_add.coords) - elif tp == "wind_sea": - tp_add = ws["tp_ws"] - else: - raise ValueError(f"tp must be 'wind_sea' or 'growth', got {tp!r}") - - keep = (hs_add > 0) & np.isfinite(tp_add) & (tp_add > 0) - out = xr.Dataset( - { - "hs": hs_add.where(keep), - "tp": tp_add.where(keep), - "dir": ws["u10dir"].where(keep), - "spr": xr.full_like(hs_add, float(spr)).where(keep), - "k": ws["k"], - } - ) - - return out.expand_dims(goal=[WIND_GOAL], partition=[WIND_PARTITION]).transpose( - "goal", "partition", "time", "sites" - ) - - -def offshore_wind_sea( - forcing: pd.DataFrame, - steepness_range: tuple[float, float] | None = None, - spr_range: tuple[float, float] | None = None, - speed_factor: float = 1.0, -) -> pd.DataFrame: - """ - Wind-sea weight and inferred wind of one goal x partition offshore forcing. - - Parameters - ---------- - forcing : pd.DataFrame - ``hs``, ``tp``, ``dir``, ``spr`` per time step (see - :func:`bluemath_tk.waves.hywaves.metamodel.goal_forcing`). - steepness_range, spr_range : tuple of float, optional - Ramps of :func:`~bluemath_tk.waves.wind.wind_sea_weight` (defaults - there). - speed_factor : float, optional - See :func:`~bluemath_tk.waves.wind.infer_wind`. - - Returns - ------- - pd.DataFrame - ``ws_weight``, ``u10``, ``u10dir`` indexed like *forcing*. - """ - - from ..wind import WIND_SEA_SPREAD, WIND_SEA_STEEPNESS, infer_wind, wind_sea_weight - - ramps: Mapping[str, tuple[float, float]] = { - "steepness_range": steepness_range or WIND_SEA_STEEPNESS, - "spr_range": spr_range or WIND_SEA_SPREAD, - } - weight = wind_sea_weight(forcing["hs"], forcing["tp"], forcing["spr"], **ramps) - u10, u10dir = infer_wind( - forcing["hs"], forcing["tp"], forcing["dir"], forcing["spr"], - speed_factor=speed_factor, **ramps, - ) - - return pd.DataFrame( - {"ws_weight": np.nan_to_num(weight), "u10": u10, "u10dir": u10dir}, - index=forcing.index, - ) diff --git a/tests/waves/test_hywaves_plots.py b/tests/waves/test_hywaves_plots.py index f7b6af0..f9d847f 100644 --- a/tests/waves/test_hywaves_plots.py +++ b/tests/waves/test_hywaves_plots.py @@ -74,62 +74,12 @@ def test_dynamic_forcing_without_wind_any_dim_order(self): 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) - def test_peak_period_ratio(self): df = pd.DataFrame({"partition": ["sea"] * 50 + ["swell"] * 50, "ratio": np.linspace(1.0, 1.3, 100)}) fig = plots.plot_peak_period_ratio(df, {"JONSWAP 3.3": 1.107}) self.assertEqual(len(fig.axes), 2) - def test_wind_case_comparison(self): - rng = np.random.default_rng(0) - hs0 = rng.uniform(0.5, 3.0, (2, 30)) - case = {"hs": [3.0, 5.0], "tp": [7.0, 9.0], "dir": [300.0, 200.0], "u10": [15.0, 25.0], - "u10dir": [320.0, 180.0], "wind_wave_angle": [20.0, -20.0], "ws_weight": [1.0, 0.6], - "hs_ws_u2": [0.13, 0.05]} - ds = xr.Dataset( - {"hs0": (("case", "sites"), hs0), "hs1": (("case", "sites"), hs0 * 1.2), - "hs_corrected": (("case", "sites"), hs0 * 1.15)}, - coords={"case": [3, 7], "longitude": ("sites", np.linspace(3, 6, 30)), - "latitude": ("sites", np.linspace(51, 53, 30)), - **{k: ("case", v) for k, v in case.items()}}, - ) - fig = plots.plot_wind_case_comparison(ds, size=2) - self.assertEqual(len(fig.axes), 2 * 4 + 2 * 2) # 4 panels + 2 colorbars per case - - def test_k_response(self): - x = {"u10": np.linspace(1, 40, 5), "angle": np.linspace(-90, 90, 5), "hs_ws_u2": np.geomspace(0.002, 0.6, 5)} - ds = xr.Dataset( - {f"k_{name}": (("sites", name), np.full((2, 5), 1.2)) for name in x}, - coords={"sites": ["a", "b"], **x, "depth": ("sites", [8.0, 20.0]), "fetch": ("sites", [4e4, 2e4])}, - attrs={"wind_dir": 300.0, "hs_ws": 2.0, "u10": 15.0, "angle": 0.0}, - ) - fig = plots.plot_k_response(ds, {"u10": (3.0, 35.0), "hs_ws_u2": (0.004, 0.29)}) - self.assertEqual(len(fig.axes), 3) - - def test_wind_sea_event(self): - time = pd.date_range("2013-12-05", periods=12, freq="h") - ws = xr.Dataset({name: ("time", np.linspace(0.1, 2.0, 12)) for name in ("hs_ws", "u10", "hs_ws_u2", "angle", "k")}, - coords={"time": time}) - hs = {"Buoy": ws["hs_ws"] * 2, "model": ws["hs_ws"] * 1.8} - fig = plots.plot_wind_sea_event(ws, hs, {"hs_ws_u2": (0.004, 0.29), "angle": (-60.0, 60.0)}) - self.assertEqual(len(fig.axes), 6) # 5 rows + the wind twin axis - - def test_event_maps(self): - coords = {"sites": ["a", "b", "c"], "site_longitude": ("sites", [3.0, 4.0, 5.0]), - "site_latitude": ("sites", [51.0, 52.0, 53.0])} - ref = xr.DataArray([1.0, 2.0, 3.0], dims="sites", coords=coords) - fig = plots.plot_event_maps(ref, {"m1": ref * 1.1, "m2": ref * 0.9}, k=ref * 0 + 1.2) - self.assertEqual(len(fig.axes), 4 + 3) # 4 maps + 3 colorbars - if __name__ == "__main__": unittest.main() diff --git a/tests/waves/test_wind_correction.py b/tests/waves/test_wind_correction.py deleted file mode 100644 index 2327140..0000000 --- a/tests/waves/test_wind_correction.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Tests for bluemath_tk.waves.hywaves.wind_correction and its reconstruction hook.""" - -import unittest - -import numpy as np -import pandas as pd -import xarray as xr - -from bluemath_tk.waves.hywaves.reconstruction import sum_partition_spectra_batched -from bluemath_tk.waves.hywaves.wind_correction import ( - WIND_GOAL, - WIND_PARTITION, - WindCorrection, - correction_features, - offshore_wind_sea, - wind_sea_contribution, - wind_sea_inputs, -) - -SITES = ["a", "b"] -TIMES = pd.date_range("2020-01-01", periods=3, freq="h") - - -def _fitted(k_of: callable) -> WindCorrection: - """A correction fitted on synthetic pairs with K = k_of(g Hs / U^2).""" - - rng = np.random.default_rng(0) - n = 1500 - u = rng.uniform(5, 25, n) - hs0 = rng.uniform(0.3, 6, n) - hs1 = hs0 * k_of(9.81 * hs0 / u**2) - return WindCorrection(max_iter=100).fit( - hs0, hs1, u, rng.uniform(1e4, 1e5, n), 20.0, 25.0, 0.0 - ) - - -def _cube(hs: float = 2.0, weight: float = 1.0) -> xr.Dataset: - """Two goals, bulk partition; goal 2 is swell (weight 0).""" - - shape = (2, 1, len(TIMES), len(SITES)) - coords = {"goal": [1, 2], "partition": [-1], "time": TIMES, "sites": SITES} - dims = ("goal", "partition", "time", "sites") - forcing_dims = ("goal", "partition", "time") - return xr.Dataset( - { - "hs": (dims, np.full(shape, hs)), - "tp": (dims, np.full(shape, 7.0)), - "dir": (dims, np.full(shape, 300.0)), - "spr": (dims, np.full(shape, 25.0)), - "ws_weight": (forcing_dims, np.stack([np.full((1, 3), weight), np.zeros((1, 3))])), - "u10": (forcing_dims, np.stack([np.full((1, 3), 15.0), np.zeros((1, 3))])), - "u10dir": (forcing_dims, np.full((2, 1, 3), 300.0)), - }, - coords=coords, - ) - - -def _fetch() -> xr.Dataset: - dirs = np.arange(0.0, 360.0, 5.0) - return xr.Dataset( - { - "fetch_eff": (("sites", "dir"), np.full((2, len(dirs)), 5e4)), - "mean_depth": (("sites", "dir"), np.full((2, len(dirs)), 25.0)), - "site_depth": ("sites", [20.0, 20.0]), - }, - coords={"sites": SITES, "dir": dirs}, - ) - - -class TestWindCorrection(unittest.TestCase): - """K model: features, fit/predict and its limits.""" - - def test_features(self): - """Dimensionless groups scale with U^2 / g; cos of the angle last.""" - - x = correction_features(2.0, 10.0, 1e4, 20.0, 30.0, 60.0) - np.testing.assert_allclose( - x[0], [2 * 9.81 / 100, 1e4 * 9.81 / 100, 20 * 9.81 / 100, 30 * 9.81 / 100, 0.5] - ) - - def test_fit_predict_and_limits(self): - """Learns K from g Hs / U^2; K = 1 without wind or wind sea.""" - - corr = _fitted(lambda h: 1 + 2 * np.clip(0.2 - h, 0, None)) - k = corr.predict(np.array([0.5, 0.5, 0.5, 0.01]), np.array([20.0, 0.5, 20.0, 20.0]), 5e4, 20.0, 25.0, 0.0) - self.assertAlmostEqual(float(k[0]), 1 + 2 * (0.2 - 9.81 * 0.5 / 400), delta=0.03) - self.assertEqual(float(k[1]), 1.0) # no wind - self.assertEqual(float(k[3]), 1.0) # no wind sea - - -class TestWindSeaContribution(unittest.TestCase): - """The extra contribution carries (K^2 - 1) E_ws from the wind direction.""" - - def test_energy_and_direction(self): - """Only weighted contributions count; hs_add^2 = (K^2 - 1) sum w hs^2.""" - - corr = _fitted(lambda h: np.full_like(h, 1.5)) - out = wind_sea_contribution(_cube(), corr, _fetch()) - self.assertEqual(out["goal"].values.tolist(), [WIND_GOAL]) - self.assertEqual(out["partition"].values.tolist(), [WIND_PARTITION]) - k = float(out["k"].mean()) - self.assertAlmostEqual(k, 1.5, delta=0.02) - np.testing.assert_allclose(out["hs"], 2.0 * np.sqrt(k**2 - 1), rtol=1e-6) - np.testing.assert_allclose(out["dir"], 300.0) - np.testing.assert_allclose(out["tp"], 7.0) - - def test_inputs(self): - """The wind sea and K inputs on (time, sites); K as the contribution's.""" - - corr = _fitted(lambda h: np.full_like(h, 1.5)) - ws = wind_sea_inputs(_cube(), corr, _fetch()) - self.assertEqual(ws["k"].dims, ("time", "sites")) - np.testing.assert_allclose(ws["hs_ws"] ** 2, ws["e_ws"], rtol=1e-6) - np.testing.assert_allclose(ws["hs_ws_u2"], 9.81 * ws["hs_ws"] / ws["u10"] ** 2, rtol=1e-6) - np.testing.assert_allclose(ws["tp_ws"], 7.0) - out = wind_sea_contribution(_cube(), corr, _fetch()) - np.testing.assert_allclose(out["k"].squeeze(["goal", "partition"]), ws["k"]) - self.assertIsNone(wind_sea_inputs(_cube(weight=0.0), corr, _fetch())) - - def test_no_wind_sea(self): - """No weighted contribution: nothing to add.""" - - corr = _fitted(lambda h: np.full_like(h, 1.5)) - self.assertIsNone(wind_sea_contribution(_cube(weight=0.0), corr, _fetch())) - - def test_offshore_wind_sea(self): - """Steep broad seas are wind sea with wind from their direction; swell not.""" - - forcing = pd.DataFrame( - {"hs": [3.0, 1.0], "tp": [7.0, 14.0], "dir": [310.0, 250.0], "spr": [30.0, 15.0]} - ) - ws = offshore_wind_sea(forcing) - self.assertEqual(ws["ws_weight"].tolist(), [1.0, 0.0]) - self.assertGreater(ws["u10"].iloc[0], 10.0) - self.assertEqual(ws["u10"].iloc[1], 0.0) - - def test_reconstruction_adds_the_energy(self): - """Summed spectra gain (K^2 - 1) E_ws; wind_k is reported; NaN parts dropped.""" - - cube = _cube() - chunks = [cube.sel(goal=[g]) for g in (1, 2)] - variables = {"hs": "hs"} - base, _ = sum_partition_spectra_batched(chunks, SITES, variables, budget_bytes=2**30) - corr = _fitted(lambda h: np.full_like(h, 1.5)) - wind, _ = sum_partition_spectra_batched( - chunks, SITES, variables, budget_bytes=2**30, - wind_correction={"correction": corr, "fetch": _fetch()}, - ) - k = float(wind["wind_k"].mean()) - expected = np.sqrt(base["hs"] ** 2 + (k**2 - 1) * 2.0**2) - np.testing.assert_allclose(wind["hs"], expected, rtol=0.02) - self.assertNotIn("wind_k", base) - - -if __name__ == "__main__": - unittest.main()