diff --git a/bluemath_tk/core/decorators.py b/bluemath_tk/core/decorators.py index 3fd3189..6f7b40d 100644 --- a/bluemath_tk/core/decorators.py +++ b/bluemath_tk/core/decorators.py @@ -459,62 +459,6 @@ def wrapper( return wrapper -def validate_data_calval(func): - """ - Validate data in CalVal class fit method. - - Parameters - ---------- - func : callable - The function to be decorated - - Returns - ------- - callable - The decorated function - """ - - @functools.wraps(func) - def wrapper( - self, - data: pd.DataFrame, - data_longitude: float, - data_latitude: float, - data_to_calibrate: pd.DataFrame, - max_time_diff: int = 2, - ): - if not isinstance(data, pd.DataFrame): - raise TypeError("Data must be a pandas DataFrame") - if not isinstance(data_longitude, float): - raise TypeError("Longitude must be a float") - data_longitude = data_longitude % 360 - if not isinstance(data_latitude, float): - raise TypeError("Latitude must be a float") - if not isinstance(data_to_calibrate, pd.DataFrame): - raise TypeError("Data to calibrate must be a pandas DataFrame") - if "LONGITUDE" not in data_to_calibrate.columns: - raise ValueError( - "Data to calibrate must contain a column named 'LONGITUDE'" - ) - if "LATITUDE" not in data_to_calibrate.columns: - raise ValueError("Data to calibrate must contain a column named 'LATITUDE'") - if "Hs_CAL" not in data_to_calibrate.columns: - raise ValueError("Data to calibrate must contain a column named 'Hs_CAL'") - if not isinstance(max_time_diff, int) or max_time_diff <= 0: - raise ValueError("Maximum time difference must be an integer and > 0") - - return func( - self, - data, - data_longitude, - data_latitude, - data_to_calibrate, - max_time_diff, - ) - - return wrapper - - def validate_gp_data(func): """ Validate data in ExactGPInterpolation class fit method. diff --git a/bluemath_tk/core/plotting/cases.py b/bluemath_tk/core/plotting/cases.py new file mode 100644 index 0000000..e93e1f2 --- /dev/null +++ b/bluemath_tk/core/plotting/cases.py @@ -0,0 +1,135 @@ +""" +Shared pieces of the numerical-case diagnostic plots. + +Used by :mod:`bluemath_tk.wrappers.snapwave.snapwave_plots` and +:mod:`bluemath_tk.wrappers.swan.swan_plots`, so a case plot looks the same +whatever model ran it: :class:`CasePlotStyle` (colormaps per variable and +mode), colour limits and colorbars, the wave-direction arrow and the wind +summary line. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker +import numpy as np +from matplotlib.axes import Axes + +from .wind import wind_components + +#: ``(xmin, xmax, ymin, ymax)`` in plot coordinates. +Extent = tuple[float, float, float, float] + +#: Default observation-point panels (2x2 under the domain panel). +TOP_VARS: tuple[str, ...] = ("hs", "tp", "dir", "spr") + + +@dataclass +class CasePlotStyle: + """Colormaps and marker size for numerical-case plots (SnapWave, SWAN).""" + + cmaps: dict[str, str] = field( + default_factory=lambda: { + "hs": "viridis", + "tp": "plasma", + "dir": "twilight_shifted", + "spr": "cividis", + } + ) + metamodel_cmaps: dict[str, str] = field( + default_factory=lambda: { + "hs": "Blues_r", + "tp": "plasma", + "dir": "twilight_shifted", + "spr": "cividis", + } + ) + site_marker_size: float = 28 + top_vars: tuple[str, ...] = TOP_VARS + + def cmap_for(self, var: str, mode: str = "dynamic") -> str: + """ + Colormap name for *var* in *mode*. + + Parameters + ---------- + var : str + Short variable name. + mode : str, optional + "metamodel" or "dynamic". Default is "dynamic". + + Returns + ------- + str + Colormap name. + """ + + if mode == "metamodel": + return self.metamodel_cmaps.get(var, self.cmaps.get(var, "viridis")) + + return self.cmaps.get(var, "viridis") + + +DEFAULT_CASE_PLOT_STYLE = CasePlotStyle() + + +def color_limits(values: np.ndarray) -> tuple[float, float]: + """Return the 2nd-98th percentile of the finite values.""" + + finite = values[np.isfinite(values)] + if finite.size == 0: + return 0.0, 1.0 + vmin, vmax = (float(v) for v in np.percentile(finite, [2, 98])) + scale = max(abs(vmin), abs(vmax), 1.0) + if vmax - vmin < max(1e-6, 1e-4 * scale): + mid = 0.5 * (vmin + vmax) + pad = max(0.5, 0.05 * abs(mid)) + return mid - pad, mid + pad + + return vmin, vmax + + +def add_colorbar(ax: Axes, mappable, vmin: float, vmax: float, fraction: float): + """Colorbar with a few ticks and two decimals.""" + + cb = plt.colorbar(mappable, ax=ax, fraction=fraction, pad=0.02) + cb.locator = mticker.MaxNLocator(nbins=5 if vmax - vmin < 1 else 6) + cb.formatter = mticker.FormatStrFormatter("%.2f") + cb.update_ticks() + + +def draw_dir_arrow(ax: Axes, lon: float, lat: float, dir_deg: float, extent: Extent): + """Arrow pointing at ``(lon, lat)`` from the nautical wave-from direction.""" + + theta = np.deg2rad(float(dir_deg)) + length = 0.08 * max(extent[1] - extent[0], extent[3] - extent[2]) + ax.annotate( + "", + xy=(lon, lat), + xytext=(lon + np.sin(theta) * length, lat + np.cos(theta) * length), + arrowprops=dict( + arrowstyle="-|>", + color="black", + lw=2.0, + shrinkA=0, + shrinkB=0, + mutation_scale=22, + ), + zorder=7, + ) + + +def wind_summary(wind: tuple[np.ndarray, ...] | None) -> str | None: + """``u10=[min, max] m/s from °`` of a ``(lon, lat, speed, dir)`` tuple.""" + + if wind is None: + return None + _, _, speed, direction = wind + u, v = wind_components(speed, direction) + mean_dir = np.mod(np.rad2deg(np.arctan2(-np.nanmean(u), -np.nanmean(v))), 360.0) + + return ( + f"u10=[{np.nanmin(speed):.3g}, {np.nanmax(speed):.3g}] m/s from {mean_dir:.0f}°" + ) diff --git a/bluemath_tk/waves/altimetry.py b/bluemath_tk/waves/altimetry.py new file mode 100644 index 0000000..58af27e --- /dev/null +++ b/bluemath_tk/waves/altimetry.py @@ -0,0 +1,299 @@ +""" +Satellite altimeter significant wave height near a point, ready to compare +with a wave hindcast. + +The altimeter store is a zarr of along-track observations on one ``obs`` +dimension, sorted by 1-degree cell and then time, with a cell index +(``cell_start``, ``cell_count`` on ``lat_cell``, ``lon_cell``), so the +observations near a point are read as a few contiguous slices: the layout of +the IMOS multi-mission product (``imos_altimeter_1985_2020.zarr``): + +- ``hs`` (calibrated Ku-band Hs), ``time``, ``latitude``, ``longitude`` + (0-360), ``flag`` (1 good, 2 probably good, 3 correctable), ``depth`` + (negative in the ocean), ``dist2coast`` (km), ``satellite`` (mission code). + +Typical use, for one hindcast node:: + + ds = open_altimeter(path) + obs = altimeter_near(ds, lon, lat, radius_km=50, min_depth=20) + passes = altimeter_passes(obs) + pairs = collocate(hindcast_node, passes) # hindcast at the pass times + +Along-track 1 Hz observations are strongly correlated, so they are first +reduced to one value per satellite pass (:func:`altimeter_passes`) before any +comparison or fit. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import xarray as xr + +#: Variables read from the altimeter store. +ALTIMETER_VARIABLES: tuple[str, ...] = ( + "time", + "latitude", + "longitude", + "hs", + "flag", + "depth", + "dist2coast", + "satellite", +) + +EARTH_RADIUS_KM = 6371.0 + +#: Deeper "depths" (m) are fill values and read as missing. +MAX_OCEAN_DEPTH = 11000.0 + + +def open_altimeter(path: str) -> xr.Dataset: + """ + Open a cell-indexed altimeter zarr (local path or URL), lazily. + + Parameters + ---------- + path : str + Zarr store path or URL (``https://``, ``s3://``, ... through fsspec). + + Returns + ------- + xr.Dataset + The store, with ``cell_start`` / ``cell_count`` loaded. + + Raises + ------ + ValueError + If the store has no cell index. + """ + + ds = xr.open_zarr(path) + if "cell_start" not in ds or "cell_count" not in ds: + raise ValueError( + f"{path} has no cell index (cell_start, cell_count on lat_cell, lon_cell)" + ) + ds["cell_start"] = ds["cell_start"].load() + ds["cell_count"] = ds["cell_count"].load() + + return ds + + +def haversine_km( + lon1: np.ndarray | float, lat1: np.ndarray | float, lon2: float, lat2: float +) -> np.ndarray: + """Great-circle distance (km) from ``(lon1, lat1)`` to ``(lon2, lat2)``.""" + + lon1, lat1 = np.deg2rad(lon1), np.deg2rad(lat1) + lon2, lat2 = np.deg2rad(lon2), np.deg2rad(lat2) + a = ( + np.sin((lat1 - lat2) / 2) ** 2 + + np.cos(lat1) * np.cos(lat2) * np.sin((lon1 - lon2) / 2) ** 2 + ) + + return 2 * EARTH_RADIUS_KM * np.arcsin(np.sqrt(a)) + + +def _cells_around(lon: float, lat: float, radius_km: float) -> list[tuple[int, int]]: + """``(lat_cell, lon_cell)`` of the 1-degree cells a circle can touch.""" + + dlat = np.rad2deg(radius_km / EARTH_RADIUS_KM) + lat0, lat1 = max(lat - dlat, -90.0), min(lat + dlat, 89.999) + coslat = max(np.cos(np.deg2rad(max(abs(lat0), abs(lat1)))), 1e-6) + dlon = min(dlat / coslat, 180.0) + lats = range(int(np.floor(lat0)), int(np.floor(lat1)) + 1) + lons = {int(np.floor(x)) % 360 for x in np.arange(lon - dlon, lon + dlon + 1, 1.0)} + lons |= {int(np.floor(lon + dlon)) % 360} + + return [(la, lo) for la in lats for lo in sorted(lons)] + + +def altimeter_near( + ds: xr.Dataset, + lon: float, + lat: float, + radius_km: float = 50.0, + *, + time: slice | None = None, + flags: tuple[int, ...] = (1, 2), + min_depth: float | None = None, + min_dist2coast: float | None = None, + variables: tuple[str, ...] = ALTIMETER_VARIABLES, +) -> pd.DataFrame: + """ + Altimeter observations within *radius_km* of a point. + + Parameters + ---------- + ds : xr.Dataset + Output of :func:`open_altimeter`. + lon, lat : float + Point (degrees; longitude in -180..180 or 0..360). + radius_km : float, optional + Search radius. Default is 50 km. + time : slice, optional + Time window, e.g. ``slice("2000", "2020")``. + flags : tuple of int, optional + Quality flags kept. Default is (1, 2): good and probably good. + min_depth : float, optional + Keep only observations in water deeper than this (m, positive): the + altimeter footprint degrades in shallow water and near the coast. + Observations without a valid depth are then dropped. + min_dist2coast : float, optional + Keep only observations farther than this from the coast (km). + Observations without a distance (e.g. HY-2A) are dropped. + variables : tuple of str, optional + Variables read. Default is :data:`ALTIMETER_VARIABLES`. + + Returns + ------- + pd.DataFrame + One row per observation, with ``lon`` (-180..180) and ``distance_km`` + to the point, sorted by time. + """ + + present = [v for v in variables if v in ds.variables] + chunks = [] + for la, lo in _cells_around(lon, lat, radius_km): + sel = {"lat_cell": la, "lon_cell": lo} + n = int(ds["cell_count"].sel(sel)) + if n == 0: + continue + start = int(ds["cell_start"].sel(sel)) + chunks.append(ds[present].isel(obs=slice(start, start + n))) + if not chunks: + return pd.DataFrame(columns=[*present, "lon", "distance_km"]) + + obs = xr.concat(chunks, dim="obs").load().to_dataframe().reset_index(drop=True) + obs["lon"] = (obs["longitude"] + 180.0) % 360.0 - 180.0 + obs["distance_km"] = haversine_km( + obs["lon"].to_numpy(), obs["latitude"].to_numpy(), lon, lat + ) + if "depth" in obs: + # a few tracks carry fill values far below any ocean depth + obs.loc[obs["depth"] < -MAX_OCEAN_DEPTH, "depth"] = np.nan + keep = obs["distance_km"] <= radius_km + if flags is not None and "flag" in obs: + keep &= obs["flag"].isin(flags) + if min_depth is not None and "depth" in obs: + keep &= obs["depth"] <= -abs(min_depth) + if min_dist2coast is not None and "dist2coast" in obs: + keep &= obs["dist2coast"] >= min_dist2coast + if time is not None: + t = pd.to_datetime(obs["time"]) + if time.start is not None: + keep &= t >= pd.Timestamp(time.start) + if time.stop is not None: + stop = pd.Timestamp(time.stop) + if len(str(time.stop)) <= 4: # a bare year includes the whole year + stop = stop + pd.DateOffset(years=1) + keep &= t < stop + + return obs[keep].sort_values("time").reset_index(drop=True) + + +def altimeter_passes( + obs: pd.DataFrame, max_gap: str = "5min", min_points: int = 3 +) -> pd.DataFrame: + """ + One value per satellite pass: the along-track median of the observations. + + Consecutive observations of the same satellite less than *max_gap* apart + belong to one pass (a track crosses a 50 km circle in about 15 s). + + Parameters + ---------- + obs : pd.DataFrame + Output of :func:`altimeter_near`. + max_gap : str, optional + Largest time gap inside a pass. Default is "5min". + min_points : int, optional + Passes with fewer observations are dropped. Default is 3. + + Returns + ------- + pd.DataFrame + Indexed by the mean pass ``time``: ``hs`` (median), ``hs_std``, + ``n``, ``lon``, ``lat``, ``distance_km``, ``depth``, ``satellite``. + """ + + columns = ["hs", "hs_std", "n", "lon", "lat", "distance_km", "depth", "satellite"] + if obs.empty: + return pd.DataFrame(columns=columns, index=pd.DatetimeIndex([], name="time")) + + obs = obs.sort_values(["satellite", "time"]).reset_index(drop=True) + t = pd.to_datetime(obs["time"]) + new = (obs["satellite"].diff() != 0) | (t.diff() > pd.Timedelta(max_gap)) + pass_id = new.cumsum() + grouped = obs.assign(time=t).groupby(pass_id) + passes = pd.DataFrame( + { + "time": grouped["time"].apply(lambda x: x.mean()), + "hs": grouped["hs"].median(), + "hs_std": grouped["hs"].std(), + "n": grouped["hs"].size(), + "lon": grouped["lon"].mean(), + "lat": grouped["latitude"].mean(), + "distance_km": grouped["distance_km"].mean(), + "depth": grouped["depth"].median() if "depth" in obs else np.nan, + "satellite": grouped["satellite"].first(), + } + ) + passes = passes[passes["n"] >= min_points] + + return passes.set_index("time").sort_index()[columns] + + +def collocate( + hindcast: xr.Dataset | pd.DataFrame, + passes: pd.DataFrame, + max_time_diff: str = "30min", +) -> pd.DataFrame: + """ + Hindcast values at the satellite pass times (nearest time step). + + The nearest step is used rather than interpolating, so directions and + partitions stay consistent; with an hourly hindcast the default 30 min + tolerance keeps every pass. + + Parameters + ---------- + hindcast : xr.Dataset or pd.DataFrame + Time series at one node (``time`` dimension or a datetime index), + e.g. ``hs``, ``dir`` and partitions ``phs0``, ``pdir0``, ... + passes : pd.DataFrame + Output of :func:`altimeter_passes`. + max_time_diff : str, optional + Largest time difference kept. Default is "30min". + + Returns + ------- + pd.DataFrame + Indexed by pass time: the hindcast variables and the pass columns + prefixed ``sat_`` (``sat_hs`` is the reference Hs). + """ + + if isinstance(hindcast, xr.Dataset): + hindcast = hindcast.to_dataframe() + hindcast = hindcast.sort_index() + times = hindcast.index.to_numpy() + if passes.empty or times.size == 0: + return pd.DataFrame( + columns=[*hindcast.columns, *(f"sat_{c}" for c in passes.columns)] + ) + + targets = passes.index.to_numpy() + right = np.clip(np.searchsorted(times, targets), 0, times.size - 1) + left = np.clip(right - 1, 0, times.size - 1) + nearest = np.where( + np.abs(times[left] - targets) <= np.abs(times[right] - targets), left, right + ) + ok = ( + np.abs(times[nearest] - targets) <= pd.Timedelta(max_time_diff).to_timedelta64() + ) + + out = hindcast.iloc[nearest[ok]].reset_index(drop=True) + sat = passes[ok].add_prefix("sat_").reset_index() + out = pd.concat([sat, out], axis=1).set_index("time") + + return out diff --git a/bluemath_tk/waves/calibration.py b/bluemath_tk/waves/calibration.py index d6ade17..79659e5 100644 --- a/bluemath_tk/waves/calibration.py +++ b/bluemath_tk/waves/calibration.py @@ -1,862 +1,451 @@ -import cartopy.crs as ccrs -import cartopy.feature as cfeature -import matplotlib as mpl -import numpy as np -import pandas as pd -import statsmodels.api as sm -import xarray as xr -from matplotlib import pyplot as plt -from matplotlib.axes import Axes -from matplotlib.figure import Figure - -from ..core.decorators import validate_data_calval -from ..core.models import BlueMathModel -from ..core.plotting.scatter import density_scatter, validation_scatter +""" +Directional calibration of hindcast significant wave height against a +reference (satellite altimetry, buoys). +Method (Mínguez et al., 2011, *Directional calibration of wave reanalysis +databases using instrumental data*): the reference energy is fitted as a sum +of the hindcast energy of each family (wind sea, swell) in each direction +sector, -def get_matching_times_between_arrays( - times1: np.ndarray, - times2: np.ndarray, - max_time_diff: int, -) -> tuple[np.ndarray, np.ndarray]: - """ - Finds matching time indices between two arrays of timestamps. +.. math:: - For each time in `times1`, finds the closest time in `times2` that is within `max_time_diff` hours. - Returns the indices of matching times in both arrays. + H_{s,ref}^2 = \\sum_k a^{sea}_k E^{sea}_k + \\sum_k a^{swell}_k E^{swell}_k, + \\qquad E_k = \\sum_{i \\in k} H_{s,i}^2, - Parameters - ---------- - times1 : np.ndarray - First array of timestamps (reference times, e.g., from model data). - times2 : np.ndarray - Second array of timestamps (e.g., from satellite or validation data). - max_time_diff : int - Maximum time difference in hours for considering times as matching. +without intercept, so calm seas stay calm. A partition contributes to the +sector of its mean direction; partition 0 is the wind sea and the others are +swells (the HyWaves ``phs{i}`` / ``pdir{i}`` naming). Without partitions the +bulk ``hs`` / ``dir`` is calibrated with one family. Each partition Hs is then +multiplied by :math:`\\sqrt{a}` of its family and sector; the bulk Hs is +scaled by the change in total partition energy. - Returns - ------- - Tuple[np.ndarray, np.ndarray] - Two arrays containing the indices of matching times: - - First array: indices in times1 that have matches - - Second array: corresponding indices in times2 that match +Coefficients that are not supported by the data (fewer than ``min_samples`` +pairs in the sector, p-value above ``p_value`` or not positive) are fixed at +1 (no change) and the others refitted, until all are accepted. - Example - ------- - >>> idx1, idx2 = get_matching_times_between_arrays( - ... model_df.index.values, - ... sat_df.index.values, - ... max_time_diff=2, - ... ) - """ +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`). +""" - indices1 = np.array([], dtype=int) - indices2 = np.array([], dtype=int) +from __future__ import annotations - for i in range(len(times1)): - # Find minimum time difference for current time1 - time_diffs = np.abs(times2 - times1[i]) - min_diff = np.min(time_diffs) +import numpy as np +import pandas as pd +import statsmodels.api as sm +import xarray as xr +from matplotlib import pyplot as plt +from matplotlib.figure import Figure - # If minimum difference is within threshold, record the indices - if min_diff < np.timedelta64(max_time_diff, "h"): - min_index = np.argmin(time_diffs) - indices1 = np.append(indices1, i) - indices2 = np.append(indices2, min_index) +from ..core.metrics import compute_metrics +from ..core.models import BlueMathModel +from ..core.plotting.scatter import validation_scatter - return indices1, indices2 +#: Metrics reported by :meth:`CalVal.metrics` (reference first: bias = hindcast - reference). +CALIBRATION_METRICS: tuple[str, ...] = ("bias", "rmse", "si", "r2", "slope") -def process_imos_satellite_data( - satellite_df: pd.DataFrame, - ini_lat: float, - end_lat: float, - ini_lon: float, - end_lon: float, - depth_threshold: float = -200, -) -> pd.DataFrame: +def sector_index(direction: np.ndarray, sector_size: float) -> np.ndarray: """ - Processes IMOS satellite data for calibration. + Direction sector of nautical directions, sectors centred on North. - This function filters and processes IMOS satellite altimeter data to be used as - reference data for calibration (e.g., as `data_to_calibrate` in CalVal.fit). + Sector 0 is ``[360 - size/2, size/2)``, sector 1 the next one clockwise, ... Parameters ---------- - satellite_df : pd.DataFrame - IMOS satellite data. Must contain columns: - - 'LATITUDE' (float): Latitude in decimal degrees - - 'LONGITUDE' (float): Longitude in decimal degrees - - 'SWH_KU_quality_control' (float): Quality control flag for Ku-band - - 'SWH_KA_quality_control' (float): Quality control flag for Ka-band - - 'SWH_KU_CAL' (float): Calibrated significant wave height (Ku-band) - - 'SWH_KA_CAL' (float): Calibrated significant wave height (Ka-band) - - 'BOT_DEPTH' (float): Bathymetry (negative values for ocean) - ini_lat : float - Minimum latitude (southern boundary) for filtering. - end_lat : float - Maximum latitude (northern boundary) for filtering. - ini_lon : float - Minimum longitude (western boundary) for filtering. - end_lon : float - Maximum longitude (eastern boundary) for filtering. - depth_threshold : float, optional - Only include points with BOT_DEPTH < depth_threshold. Default is -200. + direction : np.ndarray + Directions (degrees); NaN gives -1. + sector_size : float + Sector width (degrees), dividing 360. Returns ------- - pd.DataFrame - Filtered and processed satellite data, suitable for use as `data_to_calibrate` in CalVal.fit. - Includes a new column 'Hs_CAL' (combination of Ku-band and Ka-band calibrated significant wave heights). - - Notes - ----- - - The returned DataFrame can be used directly as the `data_to_calibrate` argument in CalVal.fit. + np.ndarray + Integer sector per value (-1 where the direction is NaN). """ - # Filter satellite data by coordinates - satellite_df = satellite_df[ - (satellite_df.LATITUDE > ini_lat) - & (satellite_df.LATITUDE < end_lat) - & (satellite_df.LONGITUDE > ini_lon) - & (satellite_df.LONGITUDE < end_lon) - & (satellite_df.BOT_DEPTH < depth_threshold) - ] - - # Process quality control - wave_height_qlt = np.nansum( - np.concatenate( - ( - satellite_df["SWH_KU_quality_control"].values[:, np.newaxis], - satellite_df["SWH_KA_quality_control"].values[:, np.newaxis], - ), - axis=1, - ), - axis=1, - ) - good_qlt = np.where(wave_height_qlt < 1.5) - - # Process wave heights - satellite_df["Hs_CAL"] = np.nansum( - np.concatenate( - ( - satellite_df["SWH_KU_CAL"].values[:, np.newaxis], - satellite_df["SWH_KA_CAL"].values[:, np.newaxis], - ), - axis=1, - ), - axis=1, + direction = np.asarray(direction, dtype=float) + n_sectors = int(round(360.0 / sector_size)) + idx = np.floor( + ((np.nan_to_num(direction) + sector_size / 2.0) % 360.0) / sector_size ) + idx = np.where(np.isfinite(direction), idx % n_sectors, -1) - return satellite_df.iloc[good_qlt] + return idx.astype(int) class CalVal(BlueMathModel): """ - Calibrates wave data using reference data. + Directional calibration of hindcast Hs against reference Hs. - This class provides a framework for calibrating wave model outputs (e.g., hindcast or reanalysis) - using reference data (e.g., satellite or buoy observations). - It supports directionally-dependent calibration for both sea and swell components. + Parameters + ---------- + sector_size : float, optional + Direction sector width (degrees). Default is 22.5 (16 sectors). + hs_prefix, dir_prefix : str, optional + Partition variable prefixes. Default "phs" / "pdir" (``phs0`` = wind + sea, ``phs1``, ... = swells). + min_samples : int, optional + Pairs with energy in a sector needed to fit its coefficient. Default 20. + p_value : float, optional + Largest accepted p-value of a coefficient. Default 0.05. Attributes ---------- - direction_bin_size : int - Size of directional bins in degrees. - direction_bins : np.ndarray - Array of bin edges for directions. - calibration_model : sm.OLS - The calibration model, more details in `statsmodels.api.OLS`. - calibrated_data : pd.DataFrame - DataFrame with columns ['Hs', 'Hs_CORR', 'Hs_CAL'] after calibration. - The time domain is the same as the model data. - calibration_params : dict - Dictionary with 'sea_correction' and 'swell_correction' correction coefficients. + coefficients : pd.DataFrame + One row per family and sector: ``family``, ``sector``, ``center``, + ``n`` (training pairs with energy there), ``a`` (energy factor), + ``factor`` (Hs factor, ``sqrt(a)``), ``fitted`` (False when fixed at + 1) and ``p_value``. """ - direction_bin_size: int = 22.5 - direction_bins: np.ndarray = np.arange( - direction_bin_size, 360.5, direction_bin_size - ) - - def __init__(self) -> None: - """ - Initialize the CalVal class. - """ - + def __init__( + self, + sector_size: float = 22.5, + hs_prefix: str = "phs", + dir_prefix: str = "pdir", + min_samples: int = 20, + p_value: float = 0.05, + ) -> None: super().__init__() self.set_logger_name(name="CalVal", level="INFO", console=True) - - # Save input data - self._data: pd.DataFrame = None - self._data_longitude: float = None - self._data_latitude: float = None - self._data_to_calibrate: pd.DataFrame = None - self._max_time_diff: int = None - - # Initialize calibration results - self._data_to_fit: tuple[pd.DataFrame, pd.DataFrame] = (None, None) - self._calibration_model: sm.OLS = None - self._calibrated_data: pd.DataFrame = None - self._calibration_params: pd.Series = None - - # Exclude large attributes from model saving - self._exclude_attributes += [ - "_data", - "_data_to_calibrate", - "_data_to_fit", - ] - - @property - def calibration_model(self) -> sm.OLS: - """Returns the calibration model.""" - - if self._calibration_model is None: - raise ValueError( - "Calibration model is not available. Please run the fit method first." - ) - - return self._calibration_model - - @property - def calibrated_data(self) -> pd.DataFrame: - """Returns the calibrated data.""" - - if self._calibrated_data is None: - raise ValueError( - "Calibrated data is not available. Please run the fit method first." - ) - - return self._calibrated_data - - @property - def calibration_params(self) -> pd.Series: - """Returns the calibration parameters.""" - - if self._calibration_params is None: - raise ValueError( - "Calibration parameters are not available. Please run the fit method first." - ) - - return self._calibration_params - - def _plot_data_domains(self) -> tuple[Figure, Axes]: - """ - Plots the domains of the data points. - - Returns - ------- - Tuple[Figure, Axes] - A tuple containing the figure and axes objects. - """ - - fig, ax = plt.subplots( - figsize=(10, 10), - subplot_kw={ - "projection": ccrs.PlateCarree(central_longitude=self._data_longitude) - }, - ) - land_10m = cfeature.NaturalEarthFeature( - "physical", - "land", - "10m", - edgecolor="face", - facecolor=cfeature.COLORS["land"], - ) - # Plot calibration data - ax.scatter( - self._data_to_calibrate.LONGITUDE, - self._data_to_calibrate.LATITUDE, - s=0.01, - c="k", - transform=ccrs.PlateCarree(), + if (360.0 / sector_size) % 1: + raise ValueError(f"sector_size must divide 360, got {sector_size}") + self.sector_size = float(sector_size) + self.n_sectors = int(round(360.0 / sector_size)) + self.hs_prefix = hs_prefix + self.dir_prefix = dir_prefix + self.min_samples = int(min_samples) + self.p_value = float(p_value) + + self.families: dict[str, list[tuple[str, str]]] = {} + self.reference: str = "" + self.coefficients: pd.DataFrame | None = None + self.train_index: pd.Index | None = None + self.test_index: pd.Index | None = None + self._pairs: pd.DataFrame | None = None + self._exclude_attributes += ["_pairs"] + + # ------------------------------------------------------------------ + # Families and design matrix + # ------------------------------------------------------------------ + + def _families(self, columns) -> dict[str, list[tuple[str, str]]]: + """``{"sea": [(phs0, pdir0)], "swell": [(phs1, pdir1), ...]}`` or bulk.""" + + columns = list(columns) + n = len(self.hs_prefix) + parts = sorted( + int(c[n:]) + for c in columns + if c.startswith(self.hs_prefix) + and c[n:].isdigit() + and f"{self.dir_prefix}{c[n:]}" in columns ) - # Plot main data point - ax.scatter( - self._data_longitude, - self._data_latitude, - s=50, - c="red", - zorder=10, - transform=ccrs.PlateCarree(), + if parts: + pair = lambda i: (f"{self.hs_prefix}{i}", f"{self.dir_prefix}{i}") # noqa: E731 + families = {"sea": [pair(parts[0])]} + if len(parts) > 1: + families["swell"] = [pair(i) for i in parts[1:]] + return families + if "hs" in columns and "dir" in columns: + return {"total": [("hs", "dir")]} + raise ValueError( + f"No partitions ({self.hs_prefix}N/{self.dir_prefix}N) nor hs/dir in {columns}" ) - # Set plot extent - ax.set_extent( - [ - self._data_longitude - 2, - self._data_longitude + 2, - self._data_latitude - 2, - self._data_latitude + 2, - ] - ) - ax.set_facecolor("lightblue") - ax.add_feature(land_10m) - - return fig, ax - def _create_vec_direc(self, waves: np.ndarray, direcs: np.ndarray) -> np.ndarray: - """ - Creates a vector of wave heights for each directional bin. + def _energy(self, data: pd.DataFrame) -> np.ndarray: + """Energy per family and sector, ``(n, n_families * n_sectors)``.""" - Parameters - ---------- - waves : np.ndarray - Wave heights. - direcs : np.ndarray - Wave directions in degrees. - - Returns - ------- - np.ndarray - Matrix of wave heights for each directional bin. - """ + blocks = [] + for members in self.families.values(): + block = np.zeros((len(data), self.n_sectors)) + for hs_name, dir_name in members: + hs = np.nan_to_num(data[hs_name].to_numpy(float)) + idx = sector_index(data[dir_name].to_numpy(float), self.sector_size) + ok = (idx >= 0) & (hs > 0) + np.add.at(block, (np.flatnonzero(ok), idx[ok]), hs[ok] ** 2) + blocks.append(block) - data = np.zeros((len(waves), len(self.direction_bins))) - for i in range(len(waves)): - if direcs[i] < 0: - direcs[i] = direcs[i] + 360 - if direcs[i] > 0 and waves[i] > 0: - # Handle direction = 360° case by mapping to the first bin (0-22.5°) - if direcs[i] >= 360: - bin_idx = 0 - else: - bin_idx = int(direcs[i] / self.direction_bin_size) - data[i, bin_idx] = waves[i] - - return data - - @staticmethod - def _get_nparts(data: pd.DataFrame) -> int: - """ - Gets the number of parts in the wave data. - - Parameters - ---------- - data : pd.DataFrame - Wave data. - - Returns - ------- - int - The number of parts in the wave data. - """ + return np.hstack(blocks) - return len([col for col in data.columns if col.startswith("Hswell")]) + # ------------------------------------------------------------------ + # Fit + # ------------------------------------------------------------------ - def _get_joined_sea_swell_data(self, data: pd.DataFrame) -> np.ndarray: + def fit( + self, + pairs: pd.DataFrame, + reference: str = "sat_hs", + train_fraction: float = 0.75, + split: str = "random", + seed: int = 0, + ) -> CalVal: """ - Joins the sea and swell data. + Fit the directional coefficients on part of the pairs. Parameters ---------- - data : pd.DataFrame - Wave data. + pairs : pd.DataFrame + Hindcast variables (``hs`` and partitions ``phs{i}``, ``pdir{i}``, + or ``hs``, ``dir``) and the reference column, one row per + collocation (e.g. :func:`bluemath_tk.waves.altimetry.collocate`). + reference : str, optional + Reference Hs column. Default is "sat_hs". + train_fraction : float, optional + Share of the pairs used for the fit; the rest validates it. + Default is 0.75. Use 1.0 to fit on everything. + split : {"random", "time"}, optional + "random" draws the training pairs at random (seeded); "time" takes + the first *train_fraction* of the record. Default "random". + seed : int, optional + Random seed for ``split="random"``. Default is 0. Returns ------- - np.ndarray - The joined sea and swell matrix. - """ - - # Process sea waves - Hsea = self._create_vec_direc(data["Hsea"], data["Dirsea"]) ** 2 - - # Process swells - Hs_swells = np.zeros(Hsea.shape) - for part in range(1, self._get_nparts(data) + 1): - Hs_swells += ( - self._create_vec_direc(data[f"Hswell{part}"], data[f"Dirswell{part}"]) - ) ** 2 - - # Combine sea and swell matrices - sea_swell_matrix = np.concatenate([Hsea, Hs_swells], axis=1) - - return sea_swell_matrix - - @validate_data_calval - def fit( - self, - data: pd.DataFrame, - data_longitude: float, - data_latitude: float, - data_to_calibrate: pd.DataFrame, - max_time_diff: int = 2, - ) -> None: + CalVal + The fitted instance. """ - Calibrate the model data using reference (calibration) data. - - This method matches the model data and calibration data in time, - constructs directionally-binned sea and swell matrices, - and fits a linear regression to obtain correction coefficients - for each direction bin. - - Parameters - ---------- - data : pd.DataFrame - Model data to calibrate. Must contain columns: - - 'Hs' (float): Significant wave height - - 'Hsea' (float): Sea component significant wave height - - 'Dirsea' (float): Sea component mean direction (degrees) - - 'Hswell1', 'Dirswell1', ... (float): Swell components (at least one required) - The index must be datetime-like. - data_longitude : float - Longitude of the model location (used for plotting and filtering). - data_latitude : float - Latitude of the model location (used for plotting and filtering). - data_to_calibrate : pd.DataFrame - Reference data for calibration. Must contain column: - - 'Hs_CAL' (float): Calibrated significant wave height (e.g., from satellite) - The index must be datetime-like. - max_time_diff : int, optional - Maximum time difference (in hours) allowed when matching model and calibration data. - Default is 2. - - Notes - ----- - - After calling this method, the calibration parameters are stored in `self.calibration_params` - and the calibrated data is available in `self.calibrated_data`. - - The calibration is directionally dependent, meaning it uses different correction coefficients - for different wave directions. - - The coefficients with p-values greater than 0.05 or negative values are set to 1.0, - indicating no correction is applied for those directions. - """ - - self.logger.info("Starting calibration fit procedure.") - - # Save input data - self._data = data.copy() - self._data_longitude = data_longitude - self._data_latitude = data_latitude - self._data_to_calibrate = data_to_calibrate.copy() - self._max_time_diff = max_time_diff - # Plot data domains - self.logger.info("Plotting data domains.") - self._plot_data_domains() + if not 0 < train_fraction <= 1: + raise ValueError(f"train_fraction must be in (0, 1], got {train_fraction}") + pairs = pairs[np.isfinite(pairs[reference]) & np.isfinite(pairs["hs"])] + if pairs.empty: + raise ValueError("No finite pairs to fit") + pairs = pairs.reset_index() + self.reference = reference + self.families = self._families(pairs.columns) + self._pairs = pairs + + n_train = int(round(train_fraction * len(pairs))) + if split == "random": + order = np.random.default_rng(seed).permutation(len(pairs)) + elif split == "time": + order = np.argsort(pairs.iloc[:, 0].to_numpy(), kind="stable") + else: + raise ValueError(f"split must be 'random' or 'time', got {split!r}") + self.train_index = pd.Index(np.sort(order[:n_train])) + self.test_index = pd.Index(np.sort(order[n_train:])) + train = pairs.loc[self.train_index] + + X = self._energy(train) + y = train[reference].to_numpy(float) ** 2 + n_obs = (X > 0).sum(axis=0) + a = np.ones(X.shape[1]) + fitted = n_obs >= self.min_samples + p_values = np.full(X.shape[1], np.nan) + for _ in range(X.shape[1]): + if not fitted.any(): + break + offset = X[:, ~fitted] @ a[~fitted] + result = sm.OLS(y - offset, X[:, fitted]).fit() + params, pvals = np.asarray(result.params), np.asarray(result.pvalues) + bad = (params <= 0) | ~(pvals < self.p_value) + a[fitted] = np.where(bad, 1.0, params) + p_values[fitted] = pvals + if not bad.any(): + break + fitted[np.flatnonzero(fitted)[bad]] = False + a[~fitted] = 1.0 + + rows = [] + for f, family in enumerate(self.families): + for k in range(self.n_sectors): + j = f * self.n_sectors + k + rows.append( + { + "family": family, + "sector": k, + "center": k * self.sector_size, + "n": int(n_obs[j]), + "a": float(a[j]), + "factor": float(np.sqrt(a[j])), + "fitted": bool(fitted[j]), + "p_value": float(p_values[j]) if fitted[j] else np.nan, + } + ) + self.coefficients = pd.DataFrame(rows) + self.logger.info( + "Fitted %d of %d coefficients on %d pairs (%d kept for validation)", + int(fitted.sum()), + fitted.size, + len(train), + len(self.test_index), + ) - # Construct matrices for calibration - self.logger.info("Matching times and constructing matrices for calibration.") + return self - # Get matching times - times_data_to_fit, times_data_to_calibrate = get_matching_times_between_arrays( - self._data.index.values, - self._data_to_calibrate.index.values, - max_time_diff=self._max_time_diff, - ) - self._data_to_fit = ( - self._data.iloc[times_data_to_fit], - self._data_to_calibrate.iloc[times_data_to_calibrate], - ) + # ------------------------------------------------------------------ + # Correction + # ------------------------------------------------------------------ - # Get joined sea and swell data - sea_swell_matrix = self._get_joined_sea_swell_data(self._data_to_fit[0]) - - # Perform calibration - self.logger.info("Fitting OLS regression for calibration.") - X = sm.add_constant(sea_swell_matrix) - self._calibration_model = sm.OLS(self._data_to_fit[1]["Hs_CAL"] ** 2, X) - calibrated_model_results = self._calibration_model.fit() - - # Get significant correction coefficients - significant_model_params = [ - model_param - if calibrated_model_results.pvalues[imp] < 0.05 and model_param > 0 - else 1.0 - for imp, model_param in enumerate(calibrated_model_results.params) - ] - - # Save sea and swell correction coefficients - self._calibration_params = { - "sea_correction": { - ip: param - for ip, param in enumerate( - np.sqrt(significant_model_params[: len(self.direction_bins)]) - ) - }, - "swell_correction": { - ip: param - for ip, param in enumerate( - np.sqrt(significant_model_params[len(self.direction_bins) :]) - ) - }, - } + def _factor(self, family: str, direction: np.ndarray) -> np.ndarray: + """Hs factor of *family* at each direction (1 where unknown).""" - # Save calibrated data to be used in plot_calibration_results() - self._calibrated_data = self.correct(self._data_to_fit[0]) - self._calibrated_data["Hs_CAL"] = self._data_to_fit[1]["Hs_CAL"].values + table = self.coefficients + factors = table.loc[table["family"] == family, "factor"].to_numpy() + idx = sector_index(direction, self.sector_size) - self.logger.info("Calibration fit procedure completed.") + return np.where(idx >= 0, factors[np.clip(idx, 0, None)], 1.0) def correct(self, data: pd.DataFrame | xr.Dataset) -> pd.DataFrame | xr.Dataset: """ - Apply the calibration correction to new data. + Apply the calibration: partition Hs times its factor, bulk Hs rescaled. Parameters ---------- data : pd.DataFrame or xr.Dataset - Data to correct. If DataFrame, must contain columns: - - 'Hs', 'Hsea', 'Dirsea', 'Hswell1', 'Dirswell1', ... - If xarray.Dataset, must have variable 'efth' and dimension 'part'. + Hindcast with the variables used in :meth:`fit` (any dimensions + for a Dataset, e.g. ``time``). Returns ------- pd.DataFrame or xr.Dataset - Corrected data. For DataFrame, returns columns ['Hs', 'Hs_CORR'] (original and corrected SWH). - For Dataset, adds variables 'corr_coeffs' and 'corr_efth'. - - Notes - ----- - - The correction is directionally dependent and uses the coefficients obtained from `fit`. + A copy with corrected ``hs`` and partition Hs; other variables + unchanged. """ - if self._calibration_params is None: - raise ValueError( - "Calibration parameters are not available. Run fit() first." - ) - - if isinstance(data, xr.Dataset): - self.logger.info( - "Input is xarray.Dataset. Applying correction to spectra data." + if self.coefficients is None: + raise ValueError("Run fit() first") + out = data.copy() + if "total" in self.families: + out["hs"] = data["hs"] * self._factor("total", np.asarray(data["dir"])) + return out + + energy_before = 0.0 + energy_after = 0.0 + for family, members in self.families.items(): + for hs_name, dir_name in members: + hs = data[hs_name] + corrected = hs * self._factor(family, np.asarray(data[dir_name])) + out[hs_name] = corrected + energy_before = energy_before + np.nan_to_num(np.asarray(hs)) ** 2 + energy_after = energy_after + np.nan_to_num(np.asarray(corrected)) ** 2 + energy_before = np.asarray(energy_before, dtype=float) + ratio = np.sqrt( + np.divide( + np.asarray(energy_after, dtype=float), + energy_before, + out=np.ones_like(energy_before), + where=energy_before > 0, ) + ) + out["hs"] = data["hs"] * ratio - corrected_data = data.copy() # Copy data to avoid modifying original data - peak_directions = corrected_data.spec.stats(["dp"]).load() - correction_coeffs = np.ones(peak_directions.dp.shape) - for n_part in peak_directions.part: - if n_part == 0: - correction_coeffs[n_part, :] = np.array( - [ - self.calibration_params["sea_correction"][ - int(peak_direction / self.direction_bin_size) - if peak_direction < 360 - else 0 # TODO: Check if this with Javi - ] - for peak_direction in peak_directions.isel( - part=n_part - ).dp.values - ] - ) - else: - correction_coeffs[n_part, :] = np.array( - [ - self.calibration_params["swell_correction"][ - int(peak_direction / self.direction_bin_size) - if peak_direction < 360 - else 0 # TODO: Check if this with Javi - ] - for peak_direction in peak_directions.isel( - part=n_part - ).dp.values - ] - ) - corrected_data["corr_coeffs"] = (("part", "time"), correction_coeffs) - corrected_data["corr_efth"] = ( - corrected_data.efth * corrected_data.corr_coeffs - ) - self.logger.info("Spectra correction complete.") + return out - return corrected_data + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ - elif isinstance(data, pd.DataFrame): - self.logger.info( - "Input is pandas.DataFrame. Applying correction to wave data." - ) + def evaluated_pairs(self) -> pd.DataFrame: + """ + The fitted pairs with the calibrated Hs (``hs_cal``) and their ``set``. - corrected_data = data.copy() - corrected_data["Hsea"] = ( - corrected_data["Hsea"].fillna(0) ** 2 - * np.array( - [ - self.calibration_params["sea_correction"][ - int(peak_direction / self.direction_bin_size) - if peak_direction < 360 - else 0 - ] - for peak_direction in corrected_data["Dirsea"] - ] - ) - ** 2 - ) - corrected_data["Hs_CORR"] = corrected_data["Hsea"] - for n_part in range(1, self._get_nparts(corrected_data) + 1): - corrected_data[f"Hswell{n_part}"] = ( - corrected_data[f"Hswell{n_part}"].fillna(0) ** 2 - * np.array( - [ - self.calibration_params["swell_correction"][ - int(peak_direction / self.direction_bin_size) - if peak_direction < 360 - else 0 - ] - for peak_direction in corrected_data[f"Dirswell{n_part}"] - ] - ) - ** 2 - ) - corrected_data["Hs_CORR"] += corrected_data[f"Hswell{n_part}"] + Returns + ------- + pd.DataFrame + The pairs given to :meth:`fit`, plus ``hs_cal`` and ``set`` + ("train" or "test"). + """ - corrected_data["Hs_CORR"] = np.sqrt(corrected_data["Hs_CORR"]) - self.logger.info("Wave data correction complete.") + if self._pairs is None: + raise ValueError("Run fit() first (the pairs are not stored after loading)") + pairs = self._pairs.copy() + pairs["hs_cal"] = np.asarray(self.correct(pairs)["hs"]) + pairs["set"] = np.where(pairs.index.isin(self.test_index), "test", "train") - return corrected_data[["Hs", "Hs_CORR"]] + return pairs - def plot_calibration_results(self) -> tuple[Figure, list]: + def metrics(self, metrics: tuple[str, ...] = CALIBRATION_METRICS) -> pd.DataFrame: """ - Plot the calibration results, including: - - Pie charts of correction coefficients for sea and swell - - Scatter plots of model vs. reference (before and after correction) - - Polar density plots of sea and swell wave climate + Hindcast against the reference, before and after the calibration. + + Parameters + ---------- + metrics : tuple of str, optional + Metric names (see :func:`bluemath_tk.core.metrics.compute_metrics`). Returns ------- - Tuple[Figure, list] - The matplotlib Figure and a list of Axes objects for all subplots. + pd.DataFrame + Rows ``set`` (train, test, all) x ``stage`` (raw, calibrated): + ``n`` and the metrics (reference first: bias = hindcast - reference). """ - self.logger.info("Plotting calibration results.") - - fig = plt.figure(figsize=(10, 15)) - gs = fig.add_gridspec(8, 2, wspace=0.4, hspace=0.7) - - # Create subplots with proper projections - ax1 = fig.add_subplot(gs[:2, 0]) # Sea correction pie - ax2 = fig.add_subplot(gs[:2, 1]) # Swell correction pie - ax1_cbar = fig.add_subplot(gs[2, 0]) # Sea correction colorbar - ax2_cbar = fig.add_subplot(gs[2, 1]) # Swell correction colorbar - ax3 = fig.add_subplot(gs[3:5, 0]) # No correction scatter - ax4 = fig.add_subplot(gs[3:5, 1]) # With correction scatter - ax5 = fig.add_subplot(gs[6:, 0], projection="polar") # Sea climate - ax6 = fig.add_subplot(gs[6:, 1], projection="polar") # Swell climate - - # Plot sea correction pie chart - sea_norm = 0.35 # Normalization factor for sea correction - sea_fracs = np.repeat(10, len(self.calibration_params["sea_correction"])) - sea_norm = mpl.colors.Normalize(1 - sea_norm, 1 + sea_norm) - sea_cmap = mpl.cm.get_cmap( - "bwr", len(self.calibration_params["sea_correction"]) - ) - sea_colors = sea_cmap( - sea_norm(list(self.calibration_params["sea_correction"].values())) - ) - ax1.pie( - sea_fracs, - labels=None, - colors=sea_colors, - startangle=90, - counterclock=False, - radius=1.2, - ) - ax1.set_title("SEA $Correction$", fontweight="bold") - # Add colorbar for sea correction below the pie chart, shrink it - _sea_cbar = mpl.colorbar.ColorbarBase( - ax1_cbar, - cmap=sea_cmap, - norm=sea_norm, - orientation="horizontal", - label="Correction Factor", - ) - box = ax1_cbar.get_position() - ax1_cbar.set_position( - [ - box.x0 + 0.15 * box.width, - box.y0 + 0.3 * box.height, - 0.7 * box.width, - 0.4 * box.height, - ] - ) - ax1_cbar.set_frame_on(False) - ax1_cbar.tick_params( - left=False, right=False, labelleft=False, labelbottom=True, bottom=True - ) - - # Plot swell correction pie chart - swell_norm = 0.35 # Normalization factor for swell correction - swell_fracs = np.repeat(10, len(self.calibration_params["swell_correction"])) - swell_norm = mpl.colors.Normalize(1 - swell_norm, 1 + swell_norm) - swell_cmap = mpl.cm.get_cmap( - "bwr", len(self.calibration_params["swell_correction"]) - ) - swell_colors = swell_cmap( - swell_norm(list(self.calibration_params["swell_correction"].values())) - ) - ax2.pie( - swell_fracs, - labels=None, - colors=swell_colors, - startangle=90, - counterclock=False, - radius=1.2, - ) - ax2.set_title("SWELL $Correction$", fontweight="bold") - # Add colorbar for swell correction below the pie chart, shrink it - _swell_cbar = mpl.colorbar.ColorbarBase( - ax2_cbar, - cmap=swell_cmap, - norm=swell_norm, - orientation="horizontal", - label="Correction Factor", - ) - box = ax2_cbar.get_position() - ax2_cbar.set_position( - [ - box.x0 + 0.15 * box.width, - box.y0 + 0.3 * box.height, - 0.7 * box.width, - 0.4 * box.height, - ] - ) - ax2_cbar.set_frame_on(False) - ax2_cbar.tick_params( - left=False, right=False, labelleft=False, labelbottom=True, bottom=True - ) - - # Plot no correction scatter - validation_scatter( - ax=ax3, - reference=self.calibrated_data["Hs_CAL"].values, - estimate=self.calibrated_data["Hs"].values, - reference_label="Satellite", - estimate_label="Hindcast", - title="No Correction", - ) + pairs = self.evaluated_pairs() + rows = [] + for name, subset in ( + ("train", pairs[pairs["set"] == "train"]), + ("test", pairs[pairs["set"] == "test"]), + ("all", pairs), + ): + for stage, column in (("raw", "hs"), ("calibrated", "hs_cal")): + values = compute_metrics( + reference=subset[self.reference], + estimate=subset[column], + metrics=metrics, + ) + rows.append({"set": name, "stage": stage, **values}) - # Plot with correction scatter - validation_scatter( - ax=ax4, - reference=self.calibrated_data["Hs_CAL"].values, - estimate=self.calibrated_data["Hs_CORR"].values, - reference_label="Satellite", - estimate_label="Hindcast", - title="With Correction", - ) + return pd.DataFrame(rows) - # Plot sea wave climate - sea_dirs = self._data["Dirsea"].iloc[::10] * np.pi / 180 - sea_heights = self._data["Hsea"].iloc[::10] - # Filter out NaN and infinite values - valid_mask = np.isfinite(sea_dirs) & np.isfinite(sea_heights) - sea_dirs_valid = sea_dirs[valid_mask] - sea_heights_valid = sea_heights[valid_mask] - - if len(sea_dirs_valid) > 0: - x, y, z = density_scatter(sea_dirs_valid, sea_heights_valid) - ax5.scatter(x, y, c=z, s=3, cmap="jet") - ax5.set_theta_zero_location("N", offset=0) - ax5.set_xticklabels(["N", "NE", "E", "SE", "S", "SW", "W", "NW"]) - ax5.xaxis.grid(True, color="lavender", linestyle="-") - ax5.yaxis.grid(True, color="lavender", linestyle="-") - ax5.set_theta_direction(-1) - ax5.set_xlabel("$\u03b8_{m}$ ($\degree$)") - ax5.set_ylabel("$H_{s}$ (m)", labelpad=20) - ax5.set_title("SEA $Wave$ $Climate$", pad=35, fontweight="bold") - - # Plot swell wave climate - swell_dirs = self._data["Dirswell1"].iloc[::10] * np.pi / 180 - swell_heights = self._data["Hswell1"].iloc[::10] - # Filter out NaN and infinite values - valid_mask = np.isfinite(swell_dirs) & np.isfinite(swell_heights) - swell_dirs_valid = swell_dirs[valid_mask] - swell_heights_valid = swell_heights[valid_mask] - - if len(swell_dirs_valid) > 0: - x, y, z = density_scatter(swell_dirs_valid, swell_heights_valid) - ax6.scatter(x, y, c=z, s=3, cmap="jet") - ax6.set_theta_zero_location("N", offset=0) - ax6.set_xticklabels(["N", "NE", "E", "SE", "S", "SW", "W", "NW"]) - ax6.xaxis.grid(True, color="lavender", linestyle="-") - ax6.yaxis.grid(True, color="lavender", linestyle="-") - ax6.set_theta_direction(-1) - ax6.set_xlabel("$\u03b8_{m}$ ($\degree$)") - ax6.set_ylabel("$H_{s}$ (m)", labelpad=20) - ax6.set_title("SWELL 1 $Wave$ $Climate$", pad=35, fontweight="bold") - - return fig, [ax1, ax2, ax1_cbar, ax2_cbar, ax3, ax4, ax5, ax6] - - def validate_calibration( - self, data_to_validate: pd.DataFrame - ) -> tuple[Figure, list]: + def plot(self, title: str | None = None) -> Figure: """ - Validate the calibration using independent validation data. - - This method compares the original and corrected model data to the validation data, - both as time series and with scatter plots. + Test pairs before and after the calibration, and the factors per sector. Parameters ---------- - data_to_validate : pd.DataFrame - Validation data. Must contain column: - - 'Hs_VAL' (float): Validation significant wave height (e.g., from buoy) - The index must be datetime-like. + title : str, optional + Figure title. Returns ------- - Tuple[Figure, list] - The matplotlib Figure and a list of Axes objects: - [time series axis, scatter (no correction), scatter (corrected)]. + Figure + Two scatters (reference on x) and one polar panel per family. """ - if "Hs_VAL" not in data_to_validate.columns: - raise ValueError("Validation data is missing required column: 'Hs_VAL'") - - data_corr = self.correct(data=self._data) - data_times, data_to_validate_times = get_matching_times_between_arrays( - times1=data_corr.index, - times2=data_to_validate.index, - max_time_diff=1, - ) - - # Create figure with a 2-row, 2-column grid, top row spans both columns - fig = plt.figure(figsize=(12, 8)) - gs = fig.add_gridspec(2, 2, height_ratios=[2, 3], hspace=0.4, wspace=0.3) - - # Top row: time series plot (spans both columns) - ax_ts = fig.add_subplot(gs[0, :]) - t = data_corr.index[data_times] - ax_ts.plot( - t, - data_to_validate["Hs_VAL"].iloc[data_to_validate_times], - label="Validation", - color="k", - lw=1.5, - ) - ax_ts.plot( - t, - data_corr["Hs"].iloc[data_times], - label="Model (No Correction)", - color="tab:blue", - alpha=0.7, - ) - ax_ts.plot( - t, - data_corr["Hs_CORR"].iloc[data_times], - label="Model (Corrected)", - color="tab:orange", - alpha=0.7, - ) - ax_ts.set_ylabel("$H_s$ (m)") - ax_ts.set_xlabel("Time") - ax_ts.set_title("Time Series Comparison") - ax_ts.legend(loc="upper right") - ax_ts.grid(True, linestyle=":", alpha=0.5) - - # Bottom row: scatter plots - ax_sc1 = fig.add_subplot(gs[1, 0]) - ax_sc2 = fig.add_subplot(gs[1, 1]) - validation_scatter( - ax=ax_sc1, - reference=data_to_validate["Hs_VAL"].iloc[data_to_validate_times].values, - estimate=data_corr["Hs"].iloc[data_times].values, - reference_label="Validation", - estimate_label="Model (No Correction)", - title="No Correction", - ) - validation_scatter( - ax=ax_sc2, - reference=data_to_validate["Hs_VAL"].iloc[data_to_validate_times].values, - estimate=data_corr["Hs_CORR"].iloc[data_times].values, - reference_label="Validation", - estimate_label="Model (Corrected)", - title="With Correction", - ) + pairs = self.evaluated_pairs() + test = pairs[pairs["set"] == "test"] + label_set = "test" + if test.empty: + test, label_set = pairs, "all" + n_fam = len(self.families) + fig = plt.figure(figsize=(4.5 * (2 + n_fam), 4.8)) + for i, (column, label) in enumerate((("hs", "raw"), ("hs_cal", "calibrated"))): + ax = fig.add_subplot(1, 2 + n_fam, i + 1) + validation_scatter( + ax, + reference=test[self.reference].to_numpy(), + estimate=test[column].to_numpy(), + reference_label="Reference Hs", + estimate_label=f"Hindcast Hs ({label})", + title=f"{label_set} pairs · {label}", + units="m", + ) + theta = np.deg2rad(np.arange(self.n_sectors) * self.sector_size) + width = np.deg2rad(self.sector_size) * 0.9 + for j, family in enumerate(self.families): + ax = fig.add_subplot(1, 2 + n_fam, 3 + j, projection="polar") + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + table = self.coefficients[self.coefficients["family"] == family] + colors = np.where( + ~table["fitted"].to_numpy(), + "0.75", + np.where(table["factor"].to_numpy() >= 1, "tab:red", "tab:blue"), + ) + ax.bar( + theta, + table["factor"].to_numpy() - 1.0, + width=width, + bottom=1.0, + color=colors, + ) + ax.plot(np.linspace(0, 2 * np.pi, 100), np.ones(100), "k-", lw=0.8) + ax.set_title(f"{family}: Hs factor (grey = not fitted)", fontsize=10) + if title: + fig.suptitle(title) + fig.tight_layout() - return fig, [ax_ts, ax_sc1, ax_sc2] + return fig diff --git a/bluemath_tk/waves/hywaves/__init__.py b/bluemath_tk/waves/hywaves/__init__.py index 506faf9..5adfeea 100644 --- a/bluemath_tk/waves/hywaves/__init__.py +++ b/bluemath_tk/waves/hywaves/__init__.py @@ -6,16 +6,17 @@ Status: Under development (Working) HyWaves: hybrid downscaling of offshore waves to the nearshore. Stationary -SnapWave runs from each offshore *goal* (see :mod:`.sectors`) train a +SnapWave or SWAN runs from each offshore *goal* (see :mod:`.sectors`) train a metamodel; nearshore conditions at the output points (see :mod:`.output_points`) are reconstructed by linear summation of the goal contributions. """ -from . import metamodel, output_points, reconstruction, sectors, wind_correction +from . import metamodel, offshore, output_points, reconstruction, sectors, wind_correction __all__ = [ "metamodel", + "offshore", "output_points", "reconstruction", "sectors", diff --git a/bluemath_tk/waves/hywaves/metamodel.py b/bluemath_tk/waves/hywaves/metamodel.py index 2daa12b..039edde 100644 --- a/bluemath_tk/waves/hywaves/metamodel.py +++ b/bluemath_tk/waves/hywaves/metamodel.py @@ -1,11 +1,12 @@ """ HyWaves metamodel for one goal: MDA -> per-variable PCA -> stacked exact GP. -Training data are stationary SnapWave runs from one goal (sampled Hs, Tp, Dir, -Spr and WL), on ``(case_num, sites)``. The metamodel maps offshore forcing to -nearshore fields, ``hs`` included, in absolute terms: +Training data are stationary model runs (SnapWave or SWAN) from one goal +(sampled Hs, Tp, Dir, Spr and WL), on ``(case_num, sites)``. The metamodel maps +offshore forcing to nearshore fields, ``hs`` included, in absolute terms: -1. :func:`detect_removed_sites` / :func:`drop_sites` - optional site filter. +1. :func:`mask_dry_sites` -> :func:`detect_removed_sites` / :func:`drop_sites` + -> :func:`fill_dry_sites` - site filter and dry-case handling. 2. :func:`fit_mda` - MDA centroids of the forcing are the training cases (the rest is held out). 3. :func:`fit_pcas` - one PCA per predicted variable over the site fields @@ -55,8 +56,9 @@ # summed. MIN_HS = 0.0 -#: Default PCA NaN handling: fraction of NaN cases above which a site is -#: dropped, and the fill value for the remaining NaNs. +#: Default PCA NaN handling: fraction of valid cases a site needs to be kept, +#: and the fill value for the remaining NaNs. After :func:`fill_dry_sites` the +#: training fields have no NaN, so this is only a safety net. DEFAULT_PCA_PREPROCESSING: dict[str, dict[str, float]] = { "nan_threshold_to_drop": {"hs": 0.99, "dir_u": 0.99, "dir_v": 0.99, "spr": 0.99}, "value_to_replace_nans": {"hs": 0.0, "dir_u": 0.0, "dir_v": 0.0, "spr": 0.0}, @@ -83,6 +85,12 @@ def detect_removed_sites( - ``{var}_ratio_max``: bound on ``var / {var}_forcing`` (the case forcing), e.g. ``{"hs_ratio_max": 3}`` drops sites where ``hs`` exceeds three times the offshore Hs of a case (numerical anomalies). + - ``{var}_nan_max``: bound on the fraction of cases where ``var`` is NaN, + e.g. ``{"hs_nan_max": 0.01}`` drops sites outside the model grid + (always NaN) and sites that are dry in more than 1 % of the cases. + Apply :func:`mask_dry_sites` first so that dry means NaN for every model. + + NaNs are skipped by the bound rules, so a dry case never trips them. Parameters ---------- @@ -100,13 +108,14 @@ def detect_removed_sites( bad: set[str] = set() for key, threshold in site_filter.items(): - if key.endswith("_ratio_max"): - var, rule = key[: -len("_ratio_max")], "ratio_max" + for rule in ("ratio_max", "nan_max", "max", "min"): + if key.endswith(f"_{rule}"): + var = key[: -len(rule) - 1] + break else: - var, _, rule = key.rpartition("_") - if rule not in ("max", "min", "ratio_max"): logger.warning( - "Ignoring site_filter key %r (expected {var}_max, _min or _ratio_max)", + "Ignoring site_filter key %r (expected {var}_max, _min, _ratio_max " + "or _nan_max)", key, ) continue @@ -119,6 +128,9 @@ def detect_removed_sites( continue agg = (cases[var] / cases[f"{var}_forcing"]).max("case_num") offenders = agg.where(agg > float(threshold), drop=True) + elif rule == "nan_max": + agg = cases[var].isnull().mean("case_num") + offenders = agg.where(agg > float(threshold), drop=True) elif rule == "max": agg = cases[var].max("case_num") offenders = agg.where(agg > float(threshold), drop=True) @@ -130,6 +142,90 @@ def detect_removed_sites( return sorted(bad) +def _site_fields(cases: xr.Dataset) -> list[str]: + """Variables on ``(case_num, sites)``: the model output fields.""" + + return [ + v + for v, da in cases.data_vars.items() + if "case_num" in da.dims and "sites" in da.dims + ] + + +def mask_dry_sites(cases: xr.Dataset, var: str = "hs") -> xr.Dataset: + """ + Set every output field to NaN where a site is dry or outside the grid. + + Models flag such points differently: SWAN writes its exception value (NaN + after reading) in every quantity, SnapWave writes ``hs = 0`` with + placeholder ``tp``/``dir`` and NaN ``spr``. Here a site-case is dry when + *var* is NaN or ``<= 0``, and then all its fields become NaN, so both + models look the same to :func:`detect_removed_sites` and + :func:`fill_dry_sites`. + + Parameters + ---------- + cases : xr.Dataset + Model cases on ``(case_num, sites)``. + var : str, optional + Wave height variable that defines dry. Default is "hs". + + Returns + ------- + xr.Dataset + A copy with dry site-cases set to NaN. + """ + + dry = cases[var].isnull() | (cases[var] <= 0) + out = cases.copy() + for name in _site_fields(cases): + out[name] = cases[name].where(~dry) + + return out + + +def fill_dry_sites(cases: xr.Dataset, var: str = "hs") -> xr.Dataset: + """ + Fill the dry site-cases left after the site filter, for PCA training. + + *var* becomes 0 (no water, no waves). Every other field takes the case + forcing (``{field}_forcing``, e.g. ``tp_forcing``): no value is physical + for a dry point, the forcing is a smooth function of the GP inputs (a + constant would add a jump the GP has to learn), and it does not matter in + the reconstruction, where a contribution with ``hs = 0`` adds no energy. + Fields without a forcing counterpart keep their NaNs (the PCA fills them). + + Parameters + ---------- + cases : xr.Dataset + Output of :func:`mask_dry_sites` after :func:`drop_sites`. + var : str, optional + Wave height variable. Default is "hs". + + Returns + ------- + xr.Dataset + A copy with the dry site-cases filled. + """ + + dry = cases[var].isnull() + if not bool(dry.any()): + return cases + out = cases.copy() + out[var] = cases[var].fillna(0.0) + for name in _site_fields(cases): + if name != var and f"{name}_forcing" in cases: + out[name] = cases[name].where(~dry, cases[f"{name}_forcing"]) + logger.info( + "Filled %d dry site-cases at %d sites (%s = 0, other fields = forcing)", + int(dry.sum()), + int(dry.any("case_num").sum()), + var, + ) + + return out + + def drop_sites(cases: xr.Dataset, site_ids: Iterable[str]) -> xr.Dataset: """ Remove sites by label. diff --git a/bluemath_tk/waves/hywaves/offshore.py b/bluemath_tk/waves/hywaves/offshore.py new file mode 100644 index 0000000..54bb241 --- /dev/null +++ b/bluemath_tk/waves/hywaves/offshore.py @@ -0,0 +1,177 @@ +""" +Offshore hindcasts at the HyWaves goals, in the HyWaves variable names. + +Every offshore source is stored on ``(time, goal)`` with the same names, so +the rest of the method (dynamic runs, reconstruction, satellite calibration) +does not depend on where it came from: + +- bulk ``hs`` (m), ``tp`` (s), ``dir`` (deg, nautical, coming from), + ``spr`` (deg); +- partitions ``phs{i}``, ``ptp{i}``, ``pdir{i}``, ``pspr{i}``, partition 0 + the wind sea and 1, 2, ... the swells. + +This module reads ERA5 ocean waves from the Earth Data Hub zarr +(: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``). +""" + +from __future__ import annotations + +import os + +import dask +import numpy as np +import xarray as xr + +#: Earth Data Hub (DestinE) ERA5 single-level ocean-wave variables. +EDH_ERA5_WAVES_URL = ( + "https://data.earthdatahub.destine.eu/era5/era5-single-levels-ocean-waves-v0.zarr" +) + +#: Environment variable holding the Earth Data Hub key. +EDH_KEY_VARIABLE = "EDH_API_KEY" + +#: HyWaves name -> ERA5 short name. +ERA5_TO_HYWAVES: dict[str, str] = { + "hs": "swh", + "tp": "pp1d", + "dir": "mwd", + "spr": "wdw", + "phs0": "shww", + "ptp0": "mpww", + "pdir0": "mdww", + "pspr0": "dwww", + **{ + name: f"{era5}{i}" + for i in (1, 2, 3) + for name, era5 in ((f"phs{i}", "swh"), (f"ptp{i}", "mwp"), (f"pdir{i}", "mwd")) + }, + **{f"pspr{i}": "dwps" for i in (1, 2, 3)}, +} + +_UNITS = {"hs": "m", "tp": "s", "dir": "degree", "spr": "degree"} +_LONG_NAMES = { + "hs": "Significant wave height", + "tp": "Peak wave period", + "dir": "Mean wave direction (coming from)", + "spr": "Directional spread", +} +_PART_NAMES = { + "phs": ("Significant wave height", "m"), + "ptp": ("Mean wave period Tm-1,0 (ERA5 has no partition peak period)", "s"), + "pdir": ("Mean wave direction (coming from)", "degree"), + "pspr": ("Directional spread", "degree"), +} + + +def open_era5_waves( + key: str | None = None, url: str = EDH_ERA5_WAVES_URL +) -> xr.Dataset: + """ + Open the ERA5 ocean-wave zarr of Earth Data Hub, lazily. + + Parameters + ---------- + key : str, optional + Earth Data Hub key. Default: the ``EDH_API_KEY`` environment variable + (keep keys out of notebooks and repositories). + url : str, optional + Store URL (without credentials). Default :data:`EDH_ERA5_WAVES_URL`. + + Returns + ------- + xr.Dataset + ERA5 waves on ``(valid_time, latitude, longitude)``. + + Raises + ------ + ValueError + If no key is given or found in the environment. + """ + + key = key or os.environ.get(EDH_KEY_VARIABLE) + if not key: + raise ValueError( + f"No Earth Data Hub key: pass key= or set the {EDH_KEY_VARIABLE} environment variable" + ) + scheme, rest = url.split("://", 1) + + return xr.open_dataset(f"{scheme}://edh:{key}@{rest}", chunks={}, engine="zarr") + + +def era5_at_goals( + ds: xr.Dataset, + lon: np.ndarray, + lat: np.ndarray, + goals: np.ndarray, + time: slice | None = None, + num_workers: int = 16, +) -> xr.Dataset: + """ + ERA5 waves at the goal nodes, in HyWaves names, on ``(time, goal)``. + + Parameters + ---------- + ds : xr.Dataset + Output of :func:`open_era5_waves`. + lon, lat : np.ndarray + Node of each goal (degrees; the nearest ERA5 grid point is taken). + goals : np.ndarray + Goal ids (the ``goal`` coordinate). + time : slice, optional + Time window, e.g. ``slice("1979", "2023")``. Default: everything. + num_workers : int, optional + Parallel chunk reads. Default is 16. + + Returns + ------- + xr.Dataset + ``hs``, ``tp``, ``dir``, ``spr`` and ``phs{i}``, ``ptp{i}``, + ``pdir{i}``, ``pspr{i}`` (i = 0 wind sea, 1-3 swells), loaded, with + the grid point of each goal as ``longitude`` / ``latitude`` + coordinates. Directional widths are converted from radians to + degrees. + """ + + lon = xr.DataArray(np.mod(np.asarray(lon, dtype=float), 360.0), dims="goal") + lat = xr.DataArray(np.asarray(lat, dtype=float), dims="goal") + era5_names = sorted(set(ERA5_TO_HYWAVES.values())) + points = ds[era5_names].sel(longitude=lon, latitude=lat, method="nearest") + if time is not None: + points = points.sel(valid_time=time) + with dask.config.set(scheduler="threads", num_workers=num_workers): + points = points.compute() + + out = xr.Dataset( + { + name: points[era5] + * (180.0 / np.pi if name.rstrip("0123456789").endswith("spr") else 1.0) + for name, era5 in ERA5_TO_HYWAVES.items() + } + ) + out = out.rename(valid_time="time").assign_coords(goal=np.asarray(goals)) + out = out.assign_coords( + longitude=("goal", ((points["longitude"].values + 180.0) % 360.0) - 180.0), + latitude=("goal", points["latitude"].values), + ) + for name in out.data_vars: + stem = name.rstrip("0123456789") + if stem in _PART_NAMES: + long_name, units = _PART_NAMES[stem] + label = "wind sea" if name.endswith("0") else f"swell {name[-1]}" + out[name].attrs = {"long_name": f"{long_name}, {label}", "units": units} + else: + out[name].attrs = {"long_name": _LONG_NAMES[name], "units": _UNITS[name]} + out[name].attrs["era5_name"] = ERA5_TO_HYWAVES[name] + out.attrs = { + "source": "ERA5 single levels, ocean waves (ECMWF), via Earth Data Hub", + "url": EDH_ERA5_WAVES_URL, + "partitions": "0 = wind sea, 1-3 = swell partitions", + "note": "ptp* are partition mean periods (Tm-1,0), not peak periods; " + "the swells share the total-swell directional width (dwps)", + } + + return out diff --git a/bluemath_tk/wrappers/snapwave/snapwave_plots.py b/bluemath_tk/wrappers/snapwave/snapwave_plots.py index cf4f409..f732ef4 100644 --- a/bluemath_tk/wrappers/snapwave/snapwave_plots.py +++ b/bluemath_tk/wrappers/snapwave/snapwave_plots.py @@ -25,12 +25,11 @@ from __future__ import annotations from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass from pathlib import Path import cartopy.crs as ccrs import matplotlib.pyplot as plt -import matplotlib.ticker as mticker import matplotlib.tri as mtri import numpy as np import xarray as xr @@ -41,6 +40,15 @@ from shapely.geometry import Polygon from ...core.plotting.base_plotting import DefaultStaticPlotting +from ...core.plotting.cases import ( + DEFAULT_CASE_PLOT_STYLE, + CasePlotStyle, + Extent, + add_colorbar, + color_limits, + draw_dir_arrow, + wind_summary, +) from .snapwave_utils import ( FORCING_FILES, parse_snapwave_inp, @@ -49,7 +57,6 @@ read_forcing_table, ) -Extent = tuple[float, float, float, float] Background = Callable[..., None] #: Native SnapWave names accepted for each short variable name. @@ -60,61 +67,9 @@ "spr": ("point_dirspr", "dirspr", "spr"), } -#: Default observation-point panels (2x2 under the domain panel). -TOP_VARS: tuple[str, ...] = ("hs", "tp", "dir", "spr") - _STATIC = DefaultStaticPlotting() -@dataclass -class CasePlotStyle: - """Colormaps, limits and marker size for SnapWave case plots.""" - - cmaps: dict[str, str] = field( - default_factory=lambda: { - "hs": "viridis", - "tp": "plasma", - "dir": "twilight_shifted", - "spr": "cividis", - } - ) - metamodel_cmaps: dict[str, str] = field( - default_factory=lambda: { - "hs": "Blues_r", - "tp": "plasma", - "dir": "twilight_shifted", - "spr": "cividis", - } - ) - site_marker_size: float = 28 - top_vars: tuple[str, ...] = TOP_VARS - - def cmap_for(self, var: str, mode: str = "dynamic") -> str: - """ - Colormap name for *var* in *mode*. - - Parameters - ---------- - var : str - Short variable name. - mode : str, optional - "metamodel" or "dynamic". Default is "dynamic". - - Returns - ------- - str - Colormap name. - """ - - if mode == "metamodel": - return self.metamodel_cmaps.get(var, self.cmaps.get(var, "viridis")) - - return self.cmaps.get(var, "viridis") - - -DEFAULT_CASE_PLOT_STYLE = CasePlotStyle() - - @dataclass class CaseContext: """Paths and settings found in a SnapWave case folder.""" @@ -257,7 +212,9 @@ def _forcing_row(case_dir: Path, var: str, time_index: int) -> tuple[float, np.n return float(table.index[idx]), table.iloc[idx].to_numpy(float) -def read_case_wind(ctx: CaseContext, extent: Extent | None = None) -> tuple[np.ndarray, ...] | None: +def read_case_wind( + ctx: CaseContext, extent: Extent | None = None +) -> tuple[np.ndarray, ...] | None: """ The wind a case was given: ``lon, lat, speed, direction`` (None without wind). @@ -278,7 +235,10 @@ def read_case_wind(ctx: CaseContext, extent: Extent | None = None) -> tuple[np.n Coordinates, speed (m/s) and nautical direction (deg, from). """ - u10, u10dir = ctx.snapwave_inp.get("u10", "").strip(), ctx.snapwave_inp.get("u10dir", "").strip() + u10, u10dir = ( + ctx.snapwave_inp.get("u10", "").strip(), + ctx.snapwave_inp.get("u10dir", "").strip(), + ) if not u10: return None speed_file, dir_file = ctx.case_dir / u10, ctx.case_dir / u10dir @@ -291,25 +251,14 @@ def read_case_wind(ctx: CaseContext, extent: Extent | None = None) -> tuple[np.n return None if value <= 0 or extent is None: return None - lon, lat = np.meshgrid(np.linspace(extent[0], extent[1], 8)[1:-1], np.linspace(extent[2], extent[3], 6)[1:-1]) + lon, lat = np.meshgrid( + np.linspace(extent[0], extent[1], 8)[1:-1], + np.linspace(extent[2], extent[3], 6)[1:-1], + ) return lon.ravel(), lat.ravel(), np.full(lon.size, value), np.full(lon.size, angle) -def wind_summary(wind: tuple[np.ndarray, ...] | None) -> str | None: - """``u10=[min, max] m/s from °`` of :func:`read_case_wind` output.""" - - if wind is None: - return None - from ...core.plotting.wind import wind_components - - _, _, speed, direction = wind - u, v = wind_components(speed, direction) - mean_dir = np.mod(np.rad2deg(np.arctan2(-np.nanmean(u), -np.nanmean(v))), 360.0) - - return f"u10=[{np.nanmin(speed):.3g}, {np.nanmax(speed):.3g}] m/s from {mean_dir:.0f}°" - - def active_boundary_node(case_dir: str | Path, time_index: int = 0) -> int: """ 0-based index of a metamodel case's active boundary node (largest Hs). @@ -340,31 +289,6 @@ def active_boundary_node(case_dir: str | Path, time_index: int = 0) -> int: return idx -def _color_limits(values: np.ndarray) -> tuple[float, float]: - """Return the 2nd-98th percentile of the finite values.""" - - finite = values[np.isfinite(values)] - if finite.size == 0: - return 0.0, 1.0 - vmin, vmax = (float(v) for v in np.percentile(finite, [2, 98])) - scale = max(abs(vmin), abs(vmax), 1.0) - if vmax - vmin < max(1e-6, 1e-4 * scale): - mid = 0.5 * (vmin + vmax) - pad = max(0.5, 0.05 * abs(mid)) - return mid - pad, mid + pad - - return vmin, vmax - - -def _add_colorbar(ax: Axes, mappable, vmin: float, vmax: float, fraction: float): - """Colorbar with a few ticks and two decimals.""" - - cb = plt.colorbar(mappable, ax=ax, fraction=fraction, pad=0.02) - cb.locator = mticker.MaxNLocator(nbins=5 if vmax - vmin < 1 else 6) - cb.formatter = mticker.FormatStrFormatter("%.2f") - cb.update_ticks() - - def _default_background( show_satellite: bool, satellite_source: str, gridfile: Path | None ) -> Background: @@ -387,27 +311,6 @@ def _extent(enclosure: Polygon, pad: float = 0.08) -> Extent: return minx - pad, maxx + pad, miny - pad, maxy + pad -def _draw_dir_arrow(ax: Axes, lon: float, lat: float, dir_deg: float, extent: Extent): - """Arrow pointing at ``(lon, lat)`` from the nautical wave-from direction.""" - - theta = np.deg2rad(float(dir_deg)) - length = 0.08 * max(extent[1] - extent[0], extent[3] - extent[2]) - ax.annotate( - "", - xy=(lon, lat), - xytext=(lon + np.sin(theta) * length, lat + np.cos(theta) * length), - arrowprops=dict( - arrowstyle="-|>", - color="black", - lw=2.0, - shrinkA=0, - shrinkB=0, - mutation_scale=22, - ), - zorder=7, - ) - - def _plot_domain_panel( ax: Axes, ctx: CaseContext, @@ -429,7 +332,7 @@ def _plot_domain_panel( ) vals = np.asarray(values.values, dtype=float) if np.isfinite(vals).any(): - vmin, vmax = _color_limits(vals) + vmin, vmax = color_limits(vals) mesh = ax.tripcolor( map_triangulation(map_ds), vals, @@ -442,7 +345,7 @@ def _plot_domain_panel( {"transform": ccrs.PlateCarree()} if isinstance(ax, GeoAxes) else {} ), ) - _add_colorbar(ax, mesh, vmin, vmax, fraction=0.025) + add_colorbar(ax, mesh, vmin, vmax, fraction=0.025) ax.set_title(f"domain map · {map_variable}", fontsize=11) else: ax.set_title("domain overview", fontsize=11) @@ -475,7 +378,7 @@ def _plot_domain_panel( zorder=6, label="active boundary", ) - _draw_dir_arrow(ax, *nodes[active], dirs[active], extent) + draw_dir_arrow(ax, *nodes[active], dirs[active], extent) else: peak = np.nanmax(hs) hs_norm = hs / peak if peak > 0 else hs @@ -490,7 +393,7 @@ def _plot_domain_panel( ) for (lon, lat), h, d in zip(nodes, hs, dirs): if h > 0: - _draw_dir_arrow(ax, lon, lat, d, extent) + draw_dir_arrow(ax, lon, lat, d, extent) wind = read_case_wind(ctx, extent) if wind is not None: @@ -529,7 +432,7 @@ def _plot_site_panel( background(ax, extent, panel="site", has_map=False) vals = np.asarray(values.values, dtype=float) ok = np.isfinite(vals) - vmin, vmax = _color_limits(vals) + vmin, vmax = color_limits(vals) sc = ax.scatter( his_ds["station_x"].values.astype(float)[ok], his_ds["station_y"].values.astype(float)[ok], @@ -547,11 +450,15 @@ def _plot_site_panel( ax.set_ylim(extent[2], extent[3]) ax.set_aspect("equal", adjustable="box") ax.tick_params(labelsize=8) - _add_colorbar(ax, sc, vmin, vmax, fraction=0.05) + add_colorbar(ax, sc, vmin, vmax, fraction=0.05) def _subtitle( - ctx: CaseContext, mode: str, time_index: int, has_map: bool, extent: Extent | None = None + ctx: CaseContext, + mode: str, + time_index: int, + has_map: bool, + extent: Extent | None = None, ) -> str: """One-line forcing summary.""" diff --git a/bluemath_tk/wrappers/swan/swan_plots.py b/bluemath_tk/wrappers/swan/swan_plots.py index ecf8f67..fc8c8a5 100644 --- a/bluemath_tk/wrappers/swan/swan_plots.py +++ b/bluemath_tk/wrappers/swan/swan_plots.py @@ -22,23 +22,26 @@ from matplotlib.axes import Axes from matplotlib.figure import Figure from matplotlib.gridspec import GridSpec +from matplotlib.tri import Triangulation -from ..snapwave.snapwave_plots import ( +from ...core.plotting.cases import ( DEFAULT_CASE_PLOT_STYLE, CasePlotStyle, - _add_colorbar, - _color_limits, - _draw_dir_arrow, + Extent, + add_colorbar, + color_limits, + draw_dir_arrow, + wind_summary, ) from .swan_utils import ( boundary_distances, + boundary_hs_profile, + read_swan_map, read_adcirc_grid, read_swan_table, swan_boundary_side, ) -Extent = tuple[float, float, float, float] - @dataclass class SwanCase: @@ -50,6 +53,7 @@ class SwanCase: points: np.ndarray | None # (n, 2) table: dict[str, np.ndarray] | None wind: tuple[np.ndarray, ...] | None # lon, lat, speed, dir + map_file: Path | None # BLOCK output on the grid (save_map cases) def _read_wind_field(path: Path, grid: str) -> tuple[np.ndarray, ...]: @@ -68,7 +72,7 @@ def _read_wind_field(path: Path, grid: str) -> tuple[np.ndarray, ...]: def load_swan_case(case_dir: str | Path, extent: Extent | None = None) -> SwanCase: """ - Parse a SWAN case folder: water level, boundary rows, points, table, wind. + Parse a SWAN case folder: water level, boundary rows, points, table, wind, map. Parameters ---------- @@ -85,7 +89,11 @@ def load_swan_case(case_dir: str | Path, extent: Extent | None = None) -> SwanCa case_dir = Path(case_dir) text = (case_dir / "INPUT").read_text() - lines = [l.strip() for l in text.splitlines() if l.strip() and not l.strip().startswith("$")] + lines = [ + l.strip() + for l in text.splitlines() + if l.strip() and not l.strip().startswith("$") + ] level = re.search(r"LEVEL\s*=\s*([-\d.eE+]+)", text) rows = [] @@ -120,10 +128,20 @@ def load_swan_case(case_dir: str | Path, extent: Extent | None = None) -> SwanCa wind = _read_wind_field(case_dir / readinp.group(1), grid.group(1)) elif uniform and float(uniform.group(1)) > 0 and extent is not None: lon, lat = np.meshgrid( - np.linspace(extent[0], extent[1], 8)[1:-1], np.linspace(extent[2], extent[3], 6)[1:-1] + np.linspace(extent[0], extent[1], 8)[1:-1], + np.linspace(extent[2], extent[3], 6)[1:-1], + ) + wind = ( + lon.ravel(), + lat.ravel(), + np.full(lon.size, float(uniform.group(1))), + np.full(lon.size, float(uniform.group(2))), ) - wind = (lon.ravel(), lat.ravel(), np.full(lon.size, float(uniform.group(1))), - np.full(lon.size, float(uniform.group(2)))) + + map_file = None + m = re.search(r"^BLOCK\s+'COMPGRID'\s+\w+\s+'([^']+)'", text, re.MULTILINE) + if m and (case_dir / m.group(1)).is_file(): + map_file = case_dir / m.group(1) return SwanCase( case_dir=case_dir, @@ -132,34 +150,64 @@ def load_swan_case(case_dir: str | Path, extent: Extent | None = None) -> SwanCa points=points, table=table, wind=wind, + map_file=map_file, ) +def _active_row(rows: np.ndarray) -> int: + """``PAR`` row of a metamodel case's goal (largest Hs, not the end padding). + + :func:`~.swan_utils.boundary_par_table` repeats the first and last nodes + at the boundary ends, so the largest Hs can appear twice; the goal is the + row that is not at an end. + """ + + hs = rows[:, 1] + peak = np.flatnonzero(np.isclose(hs, hs.max())) + inner = [i for i in peak if 0 < i < len(hs) - 1] + + return int(inner[0] if inner else peak[0]) + + def _boundary_xy(side_xy: np.ndarray, lengths: np.ndarray) -> np.ndarray: """Positions along the open boundary at ``[len]`` distances.""" dist = boundary_distances(side_xy) - return np.c_[np.interp(lengths, dist, side_xy[:, 0]), np.interp(lengths, dist, side_xy[:, 1])] + return np.c_[ + np.interp(lengths, dist, side_xy[:, 0]), np.interp(lengths, dist, side_xy[:, 1]) + ] def _plot_site_panel( - ax: Axes, points: np.ndarray, values: np.ndarray, var: str, mode: str, - style: CasePlotStyle, extent: Extent, + ax: Axes, + points: np.ndarray, + values: np.ndarray, + var: str, + mode: str, + style: CasePlotStyle, + extent: Extent, ) -> None: ok = np.isfinite(values) - vmin, vmax = _color_limits(values) + vmin, vmax = color_limits(values) sc = ax.scatter( - points[ok, 0], points[ok, 1], c=values[ok], s=style.site_marker_size, - cmap=style.cmap_for(var, mode), vmin=vmin, vmax=vmax, - edgecolors="white", linewidths=0.25, zorder=3, + points[ok, 0], + points[ok, 1], + c=values[ok], + s=style.site_marker_size, + cmap=style.cmap_for(var, mode), + vmin=vmin, + vmax=vmax, + edgecolors="white", + linewidths=0.25, + zorder=3, ) ax.set_title(var, fontsize=11) ax.set_xlim(extent[0], extent[1]) ax.set_ylim(extent[2], extent[3]) ax.set_aspect("equal", adjustable="box") ax.tick_params(labelsize=8) - _add_colorbar(ax, sc, vmin, vmax, fraction=0.05) + add_colorbar(ax, sc, vmin, vmax, fraction=0.05) def plot_case( @@ -169,10 +217,15 @@ def plot_case( title: str | None = None, show_satellite: bool = False, satellite_source: str = "arcgis", + map_variable: str = "hs", ) -> Figure: """ Diagnostic figure of a finished (or only built) SWAN case. + When the case wrote a grid map (``save_map``), the domain panel shows + *map_variable* on the whole grid; the open boundary is coloured by the Hs + SWAN imposes there. + Parameters ---------- case_dir : str or Path @@ -185,6 +238,8 @@ def plot_case( Figure title. show_satellite, satellite_source : bool, str, optional Draw a satellite image under the panels. + map_variable : str, optional + Grid map variable (``hs``, ``tp``, ``dir``, ``spr``, ``depth``). Default "hs". Returns ------- @@ -194,7 +249,6 @@ def plot_case( from ...core.plotting.base_plotting import DefaultStaticPlotting from ...core.plotting.wind import plot_wind_arrows - from ..snapwave.snapwave_plots import wind_summary if mode not in ("metamodel", "dynamic"): raise ValueError(f"mode must be 'metamodel' or 'dynamic', got {mode!r}") @@ -209,41 +263,108 @@ def plot_case( static = DefaultStaticPlotting() fig = plt.figure(figsize=(11, 11)) - gs = GridSpec(3, 2, figure=fig, height_ratios=[1.25, 1.0, 1.0], hspace=0.32, wspace=0.22) + gs = GridSpec( + 3, 2, figure=fig, height_ratios=[1.25, 1.0, 1.0], hspace=0.32, wspace=0.22 + ) axes_kw = {"projection": ccrs.PlateCarree()} if show_satellite else {} ax_map = fig.add_subplot(gs[0, :], **axes_kw) if show_satellite: static.plot_satellite(ax_map, area=extent, source=satellite_source, zorder=0) + if case.map_file is not None: + field = read_swan_map(str(case.map_file))[map_variable].to_numpy() + tri = Triangulation(x, y, grid["triangles"]) + tri.set_mask(np.isnan(field)[grid["triangles"]].any(axis=1)) + vmin, vmax = color_limits(field) + if map_variable == "hs": + vmin = 0.0 + mesh = ax_map.tripcolor( + tri, + np.nan_to_num(field, nan=vmin), + shading="gouraud", + cmap=style.cmap_for(map_variable), + vmin=vmin, + vmax=vmax, + zorder=1, + ) + add_colorbar(ax_map, mesh, vmin, vmax, fraction=0.03) + mesh.colorbar.set_label(map_variable, fontsize=8) for b in grid["land_boundaries"]: ax_map.plot(x[b], y[b], color="#555555", lw=0.6, zorder=2) ax_map.plot(*side_xy.T, color="#d62728", lw=1.6, label="open boundary", zorder=3) + if case.boundary is not None: + imposed = boundary_hs_profile(side_xy, case.boundary) + sc = ax_map.scatter( + *side_xy.T, + c=imposed, + s=6, + cmap=style.cmap_for("hs"), + vmin=0, + vmax=max(float(imposed.max()), 1e-6), + zorder=4, + ) + if case.map_file is None: + add_colorbar(ax_map, sc, 0.0, float(imposed.max()), fraction=0.03) + sc.colorbar.set_label("imposed Hs (m)", fontsize=8) if case.boundary is not None: bxy = _boundary_xy(side_xy, case.boundary[:, 0]) hs, dirs = case.boundary[:, 1], case.boundary[:, 3] if mode == "metamodel": - ax_map.scatter(*bxy.T, c="#888888", s=24, edgecolors="white", linewidths=0.4, - zorder=5, label="boundary points") - k = int(np.argmax(hs)) - ax_map.scatter(*bxy[k], c="#d62728", s=120, marker="*", edgecolors="black", - linewidths=0.5, zorder=6, label="active boundary") - _draw_dir_arrow(ax_map, *bxy[k], dirs[k], extent) + ax_map.scatter( + *bxy.T, + c="#888888", + s=24, + edgecolors="white", + linewidths=0.4, + zorder=5, + label="boundary points", + ) + k = _active_row(case.boundary) + ax_map.scatter( + *bxy[k], + c="#d62728", + s=120, + marker="*", + edgecolors="black", + linewidths=0.5, + zorder=6, + label="active boundary", + ) + draw_dir_arrow(ax_map, *bxy[k], dirs[k], extent) else: peak = np.nanmax(hs) norm = hs / peak if peak > 0 else hs - ax_map.scatter(*bxy.T, c=plt.cm.viridis(norm), s=20 + 80 * norm, edgecolors="white", - linewidths=0.4, zorder=5, label="boundary points") + ax_map.scatter( + *bxy.T, + c=plt.cm.viridis(norm), + s=20 + 80 * norm, + edgecolors="white", + linewidths=0.4, + zorder=5, + label="boundary points", + ) for (lon, lat), h, d in zip(bxy, hs, dirs): if h > 0: - _draw_dir_arrow(ax_map, lon, lat, d, extent) + draw_dir_arrow(ax_map, lon, lat, d, extent) if case.wind is not None: plot_wind_arrows(ax_map, *case.wind) if case.points is not None: - ax_map.scatter(*case.points.T, s=8, c="white", edgecolors="black", linewidths=0.2, - alpha=0.7, label="output sites", zorder=8) + ax_map.scatter( + *case.points.T, + s=8, + c="white", + edgecolors="black", + linewidths=0.2, + alpha=0.7, + label="output sites", + zorder=8, + ) ax_map.legend(loc="upper right", fontsize=7, framealpha=0.9) - ax_map.set_title("domain overview", fontsize=11) + ax_map.set_title( + f"domain map · {map_variable}" if case.map_file else "domain overview", + fontsize=11, + ) ax_map.set_xlim(extent[0], extent[1]) ax_map.set_ylim(extent[2], extent[3]) ax_map.set_aspect("equal", adjustable="datalim") @@ -252,31 +373,57 @@ def plot_case( ax = fig.add_subplot(gs[1 + i // 2, i % 2], **axes_kw) if case.table is not None and var in case.table and case.points is not None: if show_satellite: - static.plot_satellite(ax, area=extent, source=satellite_source, zorder=0) + static.plot_satellite( + ax, area=extent, source=satellite_source, zorder=0 + ) _plot_site_panel(ax, case.points, case.table[var], var, mode, style, extent) else: - ax.text(0.5, 0.5, f"{var}: no TABLE output", ha="center", va="center", - transform=ax.transAxes) + ax.text( + 0.5, + 0.5, + f"{var}: no TABLE output", + ha="center", + va="center", + transform=ax.transAxes, + ) ax.set_axis_off() parts = [] if case.boundary is not None: b = case.boundary if mode == "metamodel": - k = int(np.argmax(b[:, 1])) - parts.append(f"hs={b[k, 1]:.3g}m · tp={b[k, 2]:.3g}s · dir={b[k, 3]:.0f}° · spr={b[k, 4]:.0f}°") + k = _active_row(b) + parts.append( + f"hs={b[k, 1]:.3g}m · tp={b[k, 2]:.3g}s · dir={b[k, 3]:.0f}° · spr={b[k, 4]:.0f}°" + ) else: - for j, (var, unit) in enumerate((("hs", "m"), ("tp", "s"), ("dir", "°"), ("spr", "°")), start=1): - parts.append(f"{var}=[{np.nanmin(b[:, j]):.3g}, {np.nanmax(b[:, j]):.3g}]{unit}") + for j, (var, unit) in enumerate( + (("hs", "m"), ("tp", "s"), ("dir", "°"), ("spr", "°")), start=1 + ): + parts.append( + f"{var}=[{np.nanmin(b[:, j]):.3g}, {np.nanmax(b[:, j]):.3g}]{unit}" + ) if case.level is not None: parts.append(f"wl={case.level:.3g}m") if case.wind is not None: parts.append(wind_summary(case.wind)) parts.append("output=yes" if case.table is not None else "output=no") - fig.suptitle(title or f"SWAN {mode} · case {case_dir.name}", fontsize=14, - fontweight="bold", y=0.98) - fig.text(0.5, 0.94, " · ".join(parts), ha="center", va="top", fontsize=10, - color="#444444", wrap=True) + fig.suptitle( + title or f"SWAN {mode} · case {case_dir.name}", + fontsize=14, + fontweight="bold", + y=0.98, + ) + fig.text( + 0.5, + 0.94, + " · ".join(parts), + ha="center", + va="top", + fontsize=10, + color="#444444", + wrap=True, + ) fig.subplots_adjust(top=0.875) return fig diff --git a/bluemath_tk/wrappers/swan/swan_utils.py b/bluemath_tk/wrappers/swan/swan_utils.py index a172819..b17b507 100644 --- a/bluemath_tk/wrappers/swan/swan_utils.py +++ b/bluemath_tk/wrappers/swan/swan_utils.py @@ -51,6 +51,7 @@ EOF """ + def generate_forcing_file_GreenWaves( case_context: dict, case_dir: str, @@ -68,7 +69,7 @@ def generate_forcing_file_GreenWaves( new_date = ref_date + pd.Timedelta(hours=int(case_context.get("simul_time"))) d_date = pd.Timedelta(seconds=int(case_context.get("time_step_comp"))) - time = np.arange(ref_date, new_date + d_date, step = d_date, dtype = "datetime64[s]") + time = np.arange(ref_date, new_date + d_date, step=d_date, dtype="datetime64[s]") connectivity = ds_GFD_info.triangle_forcing_connectivity triangle_longitude = ds_GFD_info.node_forcing_longitude.isel( @@ -110,12 +111,12 @@ def generate_forcing_file_GreenWaves( windx = np.zeros((n_time, n_points)) windy = np.zeros((n_time, n_points)) - - windx[:, triangle_mask] = wind_u windy[:, triangle_mask] = wind_v - with open(op.join(case_dir, case_context.get("forcing_file", "wind_forcing.wind")), "w") as f: + with open( + op.join(case_dir, case_context.get("forcing_file", "wind_forcing.wind")), "w" + ) as f: for t in range(n_time): time_str = np.datetime_as_string(time[t], unit="s") time_str = time_str.replace("T", ".").replace("-", "").replace(":", "") @@ -129,18 +130,18 @@ def generate_forcing_file_GreenWaves( def mat_to_xr_dataset(path, variables=None, time_range=None): """ Convert a MATLAB file to xarray.Dataset, loading only requested data. - + Parameters ---------- path : str Path to .mat file variables : list of str, optional - Variable names to load (e.g., ['Hsig', 'Windv_x']). + Variable names to load (e.g., ['Hsig', 'Windv_x']). If None, loads all variables. time_range : tuple of (start, end), optional Only load timesteps within this range. e.g., ('2000-01-01', '2000-01-02') - + Returns ------- xr.Dataset @@ -149,7 +150,7 @@ def mat_to_xr_dataset(path, variables=None, time_range=None): all_keys = [name for name, shape, dtype in mat_info] vars_with_keys = {} - + for key in all_keys: if key.startswith(("Xp", "Yp", "__")): continue @@ -165,16 +166,17 @@ def mat_to_xr_dataset(path, variables=None, time_range=None): if variables is not None: vars_with_keys = {k: v for k, v in vars_with_keys.items() if k in variables} - + if not vars_with_keys: - raise ValueError(f"No matching variables found. Available: {list(vars_with_keys.keys())}") + raise ValueError( + f"No matching variables found. Available: {list(vars_with_keys.keys())}" + ) if time_range is not None: t_start, t_end = pd.to_datetime(time_range[0]), pd.to_datetime(time_range[1]) for varname in vars_with_keys: vars_with_keys[varname] = [ - (k, t) for k, t in vars_with_keys[varname] - if t_start <= t <= t_end + (k, t) for k, t in vars_with_keys[varname] if t_start <= t <= t_end ] keys_to_load = {"Xp", "Yp"} @@ -184,7 +186,7 @@ def mat_to_xr_dataset(path, variables=None, time_range=None): with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="Duplicate variable name") mat = loadmat(path, variable_names=list(keys_to_load)) - + lon = mat["Xp"].ravel().astype(np.float32) lat = mat["Yp"].ravel().astype(np.float32) @@ -192,22 +194,22 @@ def mat_to_xr_dataset(path, variables=None, time_range=None): for varname, key_times in vars_with_keys.items(): key_times.sort(key=lambda x: x[1]) # Sort by timestamp - + n_times = len(key_times) n_points = lon.size data_array = np.empty((n_times, n_points), dtype=np.float32) - + for i, (key, _) in enumerate(key_times): data_array[i] = mat[key].ravel() - + time_index = pd.DatetimeIndex([t for _, t in key_times]) - + unique_times, time_indices = np.unique(time_index, return_index=True) data_array = data_array[time_indices] time_index = pd.DatetimeIndex(unique_times) - + ds_vars[varname] = (("time", "points"), data_array) - + return xr.Dataset( data_vars=ds_vars, coords={ @@ -241,7 +243,7 @@ def vortex2SWAN( The extension for the forcing file, default is "GreenSurge_GFDcase_wind.ext". fill_value : float The fill value to use for missing data, default is 0. - + Returns ------- xarray.Dataset @@ -308,6 +310,7 @@ def vortex2SWAN( f.write("wind y component\n") np.savetxt(f, windy[t, :].reshape(1, -1), fmt="%.3f") + # -------------------------------------------------------------------------------------- # Unstructured SWAN cases forced along the open boundary (HyWaves-style), see # SwanMetaModelWrapper / SwanDynamicModelWrapper. @@ -516,10 +519,12 @@ def boundary_par_table( ``BOUNDSPEC SIDE ... VARIABLE PAR`` rows for values given at boundary nodes. Each node (e.g. a HyWaves goal) is snapped to its nearest boundary - vertex; SWAN interpolates linearly between them along the boundary. The - first and last vertices repeat the values of the nearest node, so the - whole boundary is forced (beyond the last ``[len]`` SWAN would leave it - unforced). + vertex; SWAN interpolates the spectra (energy) linearly between them along + the boundary (see :func:`boundary_hs_profile`). The first and last + vertices repeat the values of the nearest node, so the whole boundary is + forced (beyond the last ``[len]`` SWAN would leave it unforced): the + stretch between a boundary end and its nearest node takes that node's + forcing in full. Parameters ---------- @@ -570,6 +575,37 @@ def boundary_par_table( return rows +def boundary_hs_profile(open_xy: np.ndarray, rows: np.ndarray) -> np.ndarray: + """ + Hs that SWAN imposes at each open-boundary vertex for ``VARIABLE PAR`` rows. + + SWAN builds one spectrum per ``PAR`` row and interpolates the *spectra* + linearly in ``[len]`` between consecutive rows (``SINTRP``), so the energy + (``Hs**2``) is linear along the boundary and Hs is not: halfway between a + goal with Hs ``H`` and a neighbour with Hs 0 it is ``H / sqrt(2)``. This + is what makes goal contributions add up as energy, as SnapWave does with + its two-nearest-goals weights. + + Parameters + ---------- + open_xy : np.ndarray + ``(n, 2)`` boundary vertices in the order SWAN walks them (see + :func:`swan_boundary_side`). + rows : np.ndarray + ``(m, >=2)`` ``PAR`` rows, ``[len, hs, ...]``. + + Returns + ------- + np.ndarray + ``(n,)`` imposed Hs at the vertices. + """ + + rows = np.asarray(rows, dtype=float) + energy = np.interp(boundary_distances(open_xy), rows[:, 0], rows[:, 1] ** 2) + + return np.sqrt(energy) + + def write_swan_points(points: np.ndarray, out_path: str) -> str: """ Write ``x y`` output points for ``POINTS 'name' FILE``. @@ -623,7 +659,55 @@ def write_swan_wind(path: str, wind: xr.Dataset) -> str: dx = float(np.diff(lon).mean()) if lon.size > 1 else 1.0 dy = float(np.diff(lat).mean()) if lat.size > 1 else 1.0 - return f"{lon[0]:.6f} {lat[0]:.6f} 0 {lon.size - 1} {lat.size - 1} {dx:.6f} {dy:.6f}" + return ( + f"{lon[0]:.6f} {lat[0]:.6f} 0 {lon.size - 1} {lat.size - 1} {dx:.6f} {dy:.6f}" + ) + + +#: ``BLOCK ... LAY 3`` (.mat) variable names -> HyWaves short names. +SWAN_MAP_NAMES: dict[str, str] = { + "Hsig": "hs", + "TPsmoo": "tp", + "Dir": "dir", + "Dspr": "spr", + "Depth": "depth", +} + + +def read_swan_map(path: str, grid: dict | None = None) -> xr.Dataset: + """ + Read a stationary ``BLOCK 'COMPGRID' ... LAY 3`` output on an unstructured grid. + + Parameters + ---------- + path : str + The ``.mat`` file SWAN wrote. + grid : dict, optional + Output of :func:`read_adcirc_grid` for the same grid; adds ``lon`` / + ``lat`` coordinates on ``node``. + + Returns + ------- + xr.Dataset + One variable per quantity on ``node`` (``hs``, ``tp``, ``dir``, ``spr``, + ``depth``; other SWAN names kept as they are), exception values + (dry or outside points, ``<= -9``) as NaN. + """ + + data = { + SWAN_MAP_NAMES.get(k, k): np.asarray(v, dtype=float).ravel() + for k, v in loadmat(path).items() + if not k.startswith("__") + } + ds = xr.Dataset( + {k: ("node", np.where(v <= -9, np.nan, v)) for k, v in data.items()} + ) + if grid is not None: + ds = ds.assign_coords( + lon=("node", grid["vertices"][:, 0]), lat=("node", grid["vertices"][:, 1]) + ) + + return ds def read_swan_table(path: str, quantities: list[str]) -> pd.DataFrame: diff --git a/bluemath_tk/wrappers/swan/swan_wrapper.py b/bluemath_tk/wrappers/swan/swan_wrapper.py index 009a77f..f6d24a5 100644 --- a/bluemath_tk/wrappers/swan/swan_wrapper.py +++ b/bluemath_tk/wrappers/swan/swan_wrapper.py @@ -916,6 +916,11 @@ def _tide_value(ds: xr.Dataset) -> float: """ _SWAN_NUMBER = (int, float, np.integer, np.floating) +_SAVE_MAP = { + "type": (bool, np.bool_), + "value": None, + "description": "Also write the grid output (BLOCK, map_file) for this case.", +} class SwanBoundaryModelWrapper(BaseModelWrapper): @@ -926,8 +931,8 @@ class SwanBoundaryModelWrapper(BaseModelWrapper): requires) with counterclockwise triangles and boundary node strings covering the whole outline (open + land). Boundary conditions are given at ``boundary_nodes`` (e.g. HyWaves goals), each snapped to the nearest - vertex of the grid's open boundary, and SWAN interpolates linearly - between them along it (see + vertex of the grid's open boundary, and SWAN interpolates the spectra + linearly between them along it (see :func:`~bluemath_tk.wrappers.swan.swan_utils.swan_boundary_side` and :func:`~bluemath_tk.wrappers.swan.swan_utils.boundary_par_table`). @@ -936,12 +941,18 @@ class SwanBoundaryModelWrapper(BaseModelWrapper): - ``boundary_side``: e.g. ``"1 CCW"``, for ``BOUNDSPEC SIDE``; - ``boundary_par``: list of ``"len hs per dir dd"`` rows, for ``VARIABLE PAR``; - - ``n_nodes``, ``points_file``, ``table_file``, ``table_quantities``. + - ``n_nodes``, ``points_file``, ``table_file``, ``table_quantities``; + - ``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`). Attributes ---------- table_quantities : tuple of str SWAN quantities written by the templates' ``TABLE`` command, in order. + map_quantities : tuple of str + SWAN quantities written by the ``BLOCK`` command of ``save_map`` cases. """ default_parameters: dict = {} @@ -952,6 +963,7 @@ class SwanBoundaryModelWrapper(BaseModelWrapper): } table_quantities: tuple[str, ...] = ("HSIGN", "TPS", "DIR", "DSPR") + map_quantities: tuple[str, ...] = ("HSIGN", "TPS", "DIR", "DSPR", "DEPTH") def __init__( self, @@ -967,6 +979,7 @@ def __init__( open_boundary: int = 0, points_filename: str = "output_sites.txt", table_filename: str = "output_sites.tab", + map_filename: str = "output_map.mat", point_dim: str = "points", point_coords: dict[str, str] | None = None, store_parameters: tuple[str, ...] = (), @@ -1001,6 +1014,8 @@ def __init__( Which open boundary of *grid_file* is forced. Default is 0. points_filename, table_filename : str, optional Output-points file and ``TABLE`` output names. + map_filename : str, optional + ``BLOCK`` output of ``save_map`` cases. Default is "output_map.mat". point_dim : str, optional Name of the output-points dimension. Default is "points". point_coords : dict, optional @@ -1042,6 +1057,7 @@ def __init__( self.output_points = np.asarray(output_points, dtype=float).reshape(-1, 2) self.points_filename = points_filename self.table_filename = table_filename + self.map_filename = map_filename self.point_dim = point_dim self.point_coords = dict(point_coords or {}) self.store_parameters = tuple(store_parameters) @@ -1109,6 +1125,9 @@ def build_case(self, case_context: dict, case_dir: str) -> None: case_context["points_file"] = self.points_filename case_context["table_file"] = self.table_filename case_context["table_quantities"] = " ".join(self.table_quantities) + 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) def monitor_cases(self, value_counts: str | None = None): """ @@ -1291,9 +1310,9 @@ class SwanMetaModelWrapper(SwanBoundaryModelWrapper): The case ``hs`` is applied at ``active_node`` (1-based) and zero at the other boundary nodes, with the case ``tp``, ``dir``, ``spr`` everywhere, - so along the boundary Hs decays linearly from the active node to zero at - its two neighbouring nodes (as in SnapWave): each goal's metamodel learns - its own contribution, to be summed with the others. Templates read ``wl`` + so along the boundary the energy decays linearly (Hs as its square root) + from the active node to zero at its two neighbouring nodes: each goal's + metamodel learns its own contribution, to be summed with the others. Templates read ``wl`` for ``SET LEVEL``. Each case is postprocessed to a single ``case_num`` row. """ @@ -1320,6 +1339,7 @@ class SwanMetaModelWrapper(SwanBoundaryModelWrapper): "value": None, "description": "1-based boundary node where the case Hs is applied.", }, + "save_map": _SAVE_MAP, } def __init__( @@ -1468,6 +1488,7 @@ class SwanDynamicModelWrapper(SwanBoundaryModelWrapper): "value": None, "description": "Uniform wind direction (deg, nautical, coming from).", }, + "save_map": _SAVE_MAP, } def __init__(self, *args, wind: xr.Dataset | None = None, **kwargs) -> None: diff --git a/tests/data/swan/boundary/INPUT b/tests/data/swan/boundary/INPUT index 2b104e0..7cd73eb 100644 --- a/tests/data/swan/boundary/INPUT +++ b/tests/data/swan/boundary/INPUT @@ -9,5 +9,6 @@ BOUNDSPEC SIDE {{ boundary_side }} VARIABLE PAR & {% for row in boundary_par %} {{ row }}{% if not loop.last %} &{% endif %} {% endfor %}POINTS 'sites' FILE '{{ points_file }}' TABLE 'sites' NOHEAD '{{ table_file }}' {{ table_quantities }} -COMPUTE +{% if map_file %}BLOCK 'COMPGRID' NOHEAD '{{ map_file }}' LAY 3 {{ map_quantities }} +{% endif %}COMPUTE STOP diff --git a/tests/waves/test_altimetry.py b/tests/waves/test_altimetry.py new file mode 100644 index 0000000..72308fe --- /dev/null +++ b/tests/waves/test_altimetry.py @@ -0,0 +1,118 @@ +import shutil +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import pandas as pd +import xarray as xr + +from bluemath_tk.waves.altimetry import ( + altimeter_near, + altimeter_passes, + collocate, + haversine_km, + open_altimeter, +) + + +def _write_store(path: Path) -> None: + """Two 1-degree cells (lat 53, lon 3 and 4), obs sorted by cell then time.""" + + t0 = np.datetime64("2010-01-01T10:00:00") + # cell (53, 3): one pass of 5 points near (3.9, 53.5), one far point + lon_a = np.array([3.90, 3.91, 3.92, 3.93, 3.94, 3.05]) + lat_a = np.array([53.50, 53.51, 53.52, 53.53, 53.54, 53.05]) + time_a = t0 + np.array([0, 1, 2, 3, 4, 3600], dtype="timedelta64[s]") + # cell (53, 4): a pass 2 h later (one bad flag, one shallow, one fill depth) + lon_b = np.array([4.05, 4.06, 4.07, 4.08, 4.09]) + lat_b = np.array([53.50, 53.50, 53.50, 53.50, 53.50]) + time_b = t0 + np.array([7200, 7201, 7202, 7203, 7204], dtype="timedelta64[s]") + obs = { + "longitude": np.r_[lon_a, lon_b].astype("float32"), + "latitude": np.r_[lat_a, lat_b].astype("float32"), + "time": np.r_[time_a, time_b].astype("datetime64[ns]"), + "hs": np.r_[[2.0, 2.1, 2.2, 1.9, 2.0, 5.0], [3.0, 3.1, 9.0, 2.9, 3.0]].astype( + "float32" + ), + "flag": np.r_[[1, 1, 2, 1, 1, 1], [1, 1, 3, 1, 1]].astype("int8"), + "depth": np.r_[[-30] * 6, [-30, -30, -30, -5, -32000]].astype("float32"), + "dist2coast": np.full(11, 40.0, dtype="float32"), + "satellite": np.r_[[7] * 6, [11] * 5].astype("int8"), + } + lat_cell = np.arange(-90, 90, dtype="int16") + lon_cell = np.arange(0, 360, dtype="int16") + start = np.zeros((180, 360), dtype="int64") + count = np.zeros((180, 360), dtype="int32") + start[53 + 90, 3], count[53 + 90, 3] = 0, 6 + start[53 + 90, 4], count[53 + 90, 4] = 6, 5 + ds = xr.Dataset( + { + k: ("obs", v) + for k, v in obs.items() + if k not in ("time", "latitude", "longitude") + }, + coords={ + "time": ("obs", obs["time"]), + "latitude": ("obs", obs["latitude"]), + "longitude": ("obs", obs["longitude"]), + "lat_cell": lat_cell, + "lon_cell": lon_cell, + }, + ) + ds["cell_start"] = (("lat_cell", "lon_cell"), start) + ds["cell_count"] = (("lat_cell", "lon_cell"), count) + ds.to_zarr(path) + + +class TestAltimetry(unittest.TestCase): + """Reading, pass averaging and collocation of altimeter data.""" + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.path = Path(self.tmp) / "alt.zarr" + _write_store(self.path) + self.ds = open_altimeter(str(self.path)) + + def tearDown(self): + shutil.rmtree(self.tmp) + + def test_near_radius_flags_depth(self): + """Radius, quality flags and depth (fill values as missing) filter obs.""" + + obs = altimeter_near(self.ds, 4.0, 53.5, radius_km=20, flags=None) + self.assertEqual(len(obs), 10) # the far point (3.05, 53.05) is out + self.assertTrue((obs["distance_km"] <= 20).all()) + self.assertEqual(int(obs["depth"].isna().sum()), 1) # -32000 fill value + obs = altimeter_near(self.ds, 4.0, 53.5, radius_km=20, min_depth=10) + self.assertEqual(len(obs), 7) # flag 3, the 5 m and the fill-depth points out + obs = altimeter_near(self.ds, 4.0, 53.5, radius_km=20, time=slice("2011", None)) + self.assertTrue(obs.empty) + + def test_passes_and_collocation(self): + """One median per pass; hindcast at the nearest step within tolerance.""" + + obs = altimeter_near(self.ds, 4.0, 53.5, radius_km=20, min_depth=10) + passes = altimeter_passes(obs, min_points=2) + self.assertEqual(len(passes), 2) + self.assertAlmostEqual(float(passes["hs"].iloc[0]), 2.0, places=5) + self.assertEqual(list(passes["n"]), [5, 2]) + + time = pd.date_range("2010-01-01", periods=48, freq="h") + hindcast = xr.Dataset({"hs": ("time", np.arange(48.0))}, coords={"time": time}) + pairs = collocate(hindcast, passes, max_time_diff="30min") + self.assertEqual(list(pairs["hs"]), [10.0, 12.0]) + self.assertIn("sat_hs", pairs) + none = collocate(hindcast.isel(time=slice(0, 5)), passes, max_time_diff="30min") + self.assertTrue(none.empty) + + def test_haversine(self): + """One degree of latitude is ~111 km.""" + + self.assertAlmostEqual( + float(haversine_km(0.0, 1.0, 0.0, 0.0)), 111.19, places=1 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/waves/test_calibration.py b/tests/waves/test_calibration.py new file mode 100644 index 0000000..e66706b --- /dev/null +++ b/tests/waves/test_calibration.py @@ -0,0 +1,94 @@ +import unittest + +import numpy as np +import pandas as pd +import xarray as xr + +from bluemath_tk.waves.calibration import CalVal, sector_index + + +def _pairs(n: int = 2000, seed: int = 0) -> pd.DataFrame: + """Hindcast underestimating the westerly wind sea by 10 %, swell right.""" + + rng = np.random.default_rng(seed) + sea = rng.uniform(0.5, 4.0, n) + swell = rng.uniform(0.2, 2.0, n) + dir_sea = rng.choice([270.0, 0.0], n) + dir_swell = np.full(n, 315.0) + factor = np.where(dir_sea == 270.0, 1.1, 1.0) + ref = np.sqrt((factor * sea) ** 2 + swell**2) * rng.normal(1.0, 0.02, n) + + return pd.DataFrame( + { + "hs": np.sqrt(sea**2 + swell**2), + "phs0": sea, + "pdir0": dir_sea, + "phs1": swell, + "pdir1": dir_swell, + "sat_hs": ref, + }, + index=pd.date_range("2000", periods=n, freq="6h", name="time"), + ) + + +class TestCalVal(unittest.TestCase): + """Directional calibration against a reference.""" + + def test_sectors(self): + """Sectors are centred on North; NaN gives -1.""" + + np.testing.assert_array_equal( + sector_index(np.array([0.0, 11.0, 11.3, 349.0, 270.0, np.nan]), 22.5), + [0, 0, 1, 0, 12, -1], + ) + + def test_fit_recovers_factor(self): + """The westerly sea sector gets ~1.1, the rest stays ~1 or unfitted.""" + + cv = CalVal().fit(_pairs(), train_fraction=0.75) + table = cv.coefficients.set_index(["family", "sector"]) + self.assertAlmostEqual(table.loc[("sea", 12), "factor"], 1.1, delta=0.02) + self.assertAlmostEqual(table.loc[("sea", 0), "factor"], 1.0, delta=0.02) + self.assertAlmostEqual(table.loc[("swell", 14), "factor"], 1.0, delta=0.03) + self.assertFalse(table.loc[("sea", 4), "fitted"]) # no data from the east + self.assertEqual(len(cv.train_index), 1500) + self.assertEqual(len(cv.test_index), 500) + + metrics = cv.metrics().set_index(["set", "stage"]) + self.assertEqual(int(metrics.loc[("test", "raw"), "n"]), 500) + self.assertLess( + metrics.loc[("test", "calibrated"), "rmse"], + metrics.loc[("test", "raw"), "rmse"], + ) + self.assertLess(abs(metrics.loc[("test", "calibrated"), "bias"]), 0.02) + + def test_correct_dataset_and_time_split(self): + """Correction applies to a goal dataset; the time split keeps order.""" + + pairs = _pairs() + cv = CalVal().fit(pairs, split="time", train_fraction=0.5) + self.assertEqual(int(cv.train_index.max()), 999) + ds = xr.Dataset.from_dataframe(pairs.drop(columns="sat_hs")) + out = cv.correct(ds) + west = pairs["pdir0"].to_numpy() == 270.0 + ratio = out["phs0"].to_numpy() / pairs["phs0"].to_numpy() + self.assertAlmostEqual(float(np.mean(ratio[west])), 1.1, delta=0.02) + self.assertTrue(np.allclose(out["pdir0"], ds["pdir0"])) + energy = np.sqrt(out["phs0"] ** 2 + out["phs1"] ** 2) + np.testing.assert_allclose(out["hs"], energy, rtol=1e-6) + + def test_bulk_only(self): + """Without partitions, hs/dir form one family.""" + + pairs = ( + _pairs()[["hs", "phs0", "pdir0", "sat_hs"]] + .rename(columns={"pdir0": "dir"}) + .drop(columns="phs0") + ) + cv = CalVal().fit(pairs) + self.assertEqual(list(cv.families), ["total"]) + self.assertGreater(cv.coefficients.query("sector == 12")["factor"].item(), 1.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/waves/test_hywaves.py b/tests/waves/test_hywaves.py index d58ce41..262584a 100644 --- a/tests/waves/test_hywaves.py +++ b/tests/waves/test_hywaves.py @@ -267,6 +267,34 @@ def test_site_filter(self): bad = metamodel.detect_removed_sites(cases, {"hs_ratio_max": 3.0}) self.assertEqual(bad, ["s2"]) + def test_dry_sites(self): + """Dry site-cases are masked, filtered by fraction and filled.""" + + cases = self.cases.copy(deep=True) + cases["hs"][:3, 1] = 0.0 # SnapWave-style dry: hs = 0, placeholder dir + cases["dir"][:3, 1] = 270.0 + cases["spr"][:3, 1] = np.nan + cases["hs"][:30, 2] = np.nan # SWAN-style dry in half of the cases + + masked = metamodel.mask_dry_sites(cases) + for var in ("hs", "dir", "spr"): + self.assertTrue(bool(masked[var][:3, 1].isnull().all())) + self.assertFalse(bool(masked.hs_forcing.isnull().any())) + + bad = metamodel.detect_removed_sites(masked, {"hs_nan_max": 0.1}) + self.assertEqual(bad, ["s0", "s2"]) # outside the grid, dry 50 % + + filled = metamodel.fill_dry_sites(metamodel.drop_sites(masked, bad)) + self.assertFalse(bool(filled.hs.isnull().any())) + np.testing.assert_array_equal(filled.hs.sel(sites="s1")[:3], 0.0) + np.testing.assert_allclose( + filled.dir.sel(sites="s1")[:3], cases.dir_forcing[:3] + ) + np.testing.assert_allclose( + filled.spr.sel(sites="s1")[:3], cases.spr_forcing[:3] + ) + np.testing.assert_allclose(filled.hs.sel(sites="s3"), cases.hs.sel(sites="s3")) + def test_fit_and_predict(self): """Fit on MDA centroids, predict held-out cases with floors applied.""" diff --git a/tests/waves/test_hywaves_offshore.py b/tests/waves/test_hywaves_offshore.py new file mode 100644 index 0000000..564087a --- /dev/null +++ b/tests/waves/test_hywaves_offshore.py @@ -0,0 +1,71 @@ +import os +import unittest +from unittest import mock + +import numpy as np +import pandas as pd +import xarray as xr + +from bluemath_tk.waves.hywaves.offshore import ( + ERA5_TO_HYWAVES, + EDH_KEY_VARIABLE, + era5_at_goals, + open_era5_waves, +) + + +def _era5_like() -> xr.Dataset: + """Tiny ERA5-like grid: every variable = 1, widths = 0.5 rad.""" + + time = pd.date_range("2013-01-01", periods=6, freq="h") + lat = np.arange(55.0, 50.5, -0.5) + lon = np.arange(0.0, 8.0, 0.5) + shape = (time.size, lat.size, lon.size) + data = { + name: ( + ("valid_time", "latitude", "longitude"), + np.full(shape, 0.5 if name.startswith(("wdw", "dw")) else 1.0), + ) + for name in set(ERA5_TO_HYWAVES.values()) + } + data["swh"] = (data["swh"][0], np.broadcast_to(lon, shape).copy()) # hs = longitude + return xr.Dataset( + data, coords={"valid_time": time, "latitude": lat, "longitude": lon} + ) + + +class TestEra5AtGoals(unittest.TestCase): + """ERA5 waves at goal nodes in HyWaves names.""" + + def test_names_units_and_nodes(self): + """All partitions, widths in degrees, nearest grid point per goal.""" + + out = era5_at_goals( + _era5_like(), + [2.4, 7.1], + [51.6, 54.0], + [1, 2], + time=slice("2013-01-01 01", None), + ) + self.assertEqual(set(out.data_vars), set(ERA5_TO_HYWAVES)) + for i in range(4): + self.assertIn(f"phs{i}", out) + self.assertEqual(out.sizes["time"], 5) + self.assertEqual(list(out.goal.values), [1, 2]) + np.testing.assert_allclose(out["hs"].isel(time=0), [2.5, 7.0]) + np.testing.assert_allclose(out["spr"], np.rad2deg(0.5)) + np.testing.assert_allclose(out["pspr2"], np.rad2deg(0.5)) + np.testing.assert_allclose(out["latitude"], [51.5, 54.0]) + self.assertIn("mean wave period", out["ptp1"].attrs["long_name"].lower()) + + def test_key_required(self): + """Without a key (argument or environment) opening fails early.""" + + with mock.patch.dict(os.environ, {}, clear=True): + self.assertNotIn(EDH_KEY_VARIABLE, os.environ) + with self.assertRaises(ValueError): + open_era5_waves() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/wrappers/test_swan_boundary_wrapper.py b/tests/wrappers/test_swan_boundary_wrapper.py index 0159903..3589ac5 100644 --- a/tests/wrappers/test_swan_boundary_wrapper.py +++ b/tests/wrappers/test_swan_boundary_wrapper.py @@ -13,6 +13,7 @@ from bluemath_tk.wrappers.snapwave import boundary_time_series from bluemath_tk.wrappers.swan.swan_utils import ( + boundary_hs_profile, boundary_par_table, clockwise_triangles, read_adcirc_grid, @@ -91,6 +92,20 @@ def test_par_rows_padded_and_ordered(self): with self.assertRaises(ValueError): boundary_par_table(xy, GOALS[::-1], {"hs": [1.0, 0.0, 0.0]}) + def test_imposed_hs_is_energy_linear(self): + """SWAN interpolates spectra: Hs**2 is linear between PAR rows.""" + + xy, _ = swan_boundary_side(read_adcirc_grid(self.grid)) + rows = boundary_par_table(xy, GOALS, {"hs": [2.0, 0.0, 0.0]}) + table = np.c_[rows["len"], rows["hs"]] + hs = boundary_hs_profile(xy, table) + d = np.r_[0.0, np.cumsum(np.hypot(*np.diff(xy, axis=0).T))] + np.testing.assert_allclose(hs[d <= 0.1], 2.0) + third = np.isclose(d, 0.2) # a third of the way from 0.1 to 0.4 + self.assertTrue(third.any()) + np.testing.assert_allclose(hs[third], 2.0 * np.sqrt(2.0 / 3.0)) + np.testing.assert_allclose(hs[d >= 0.4], 0.0) + def test_clockwise_triangles_are_reoriented(self): """Clockwise triangles are found and a counterclockwise copy written.""" clockwise = op.join(self.test_dir, "clockwise.grd") @@ -144,6 +159,7 @@ def setUp(self): "dir": [270.0, 280.0], "spr": [20.0, 25.0], "wl": [0.0, 1.0], + "save_map": [False, True], }, fixed_parameters={}, output_dir=op.join(self.test_dir, "cases"), @@ -184,6 +200,27 @@ 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( + "BLOCK 'COMPGRID' NOHEAD 'output_map.mat' LAY 3 HSIGN TPS DIR DSPR DEPTH", + inp, + ) + with open(op.join(self.wrapper.cases_dirs[0], "INPUT")) as f: + self.assertNotIn("BLOCK", f.read()) + + def test_read_map(self): + """The BLOCK .mat output becomes a node dataset, exceptions as NaN.""" + + from scipy.io import savemat + + from bluemath_tk.wrappers.swan.swan_utils import read_swan_map + + path = op.join(self.test_dir, "map.mat") + savemat(path, {"Hsig": np.array([[1.0, -9.0, 2.0]]), "Depth": np.ones((1, 3))}) + grid = {"vertices": np.array([[0, 0, 5], [1, 0, 5], [0, 1, 5]], float)} + ds = read_swan_map(path, grid) + self.assertEqual(set(ds.data_vars), {"hs", "depth"}) + np.testing.assert_array_equal(ds.hs.isnull(), [False, True, False]) + np.testing.assert_array_equal(ds.lon, [0, 1, 0]) def test_postprocess_split(self): """Fake tables are joined on case_num and split per goal.""" @@ -245,7 +282,6 @@ def test_hourly_cases(self): finally: shutil.rmtree(test_dir) - def test_wind_field_and_uniform_wind(self): """A wind field is written as u then v blocks, south to north, with its grid.""" test_dir = tempfile.mkdtemp() @@ -293,7 +329,9 @@ def test_wind_field_and_uniform_wind(self): self.assertEqual((context["u10"], context["u10dir"]), (12.0, 300.0)) values = np.loadtxt(op.join(wrapper.cases_dirs[0], "wind.dat")) np.testing.assert_array_equal(values[:2], 5.0) # u block - np.testing.assert_array_equal(values[2:, 0], [2.0, 1.0]) # v at 52N, then 53N + np.testing.assert_array_equal( + values[2:, 0], [2.0, 1.0] + ) # v at 52N, then 53N # the case plot reads INPUT, wind.dat and the TABLE output back import matplotlib @@ -303,7 +341,9 @@ def test_wind_field_and_uniform_wind(self): case_dir = wrapper.cases_dirs[0] with open(op.join(case_dir, "INPUT"), "a") as f: - f.write(f"INPGRID WIND REGULAR {context['wind_grid']}\nREADINP WIND 1 'wind.dat' 3 0 FREE\n") + f.write( + 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") case = load_swan_case(case_dir) @@ -346,5 +386,6 @@ def test_grid_copied_without_symlinks(self): finally: shutil.rmtree(test_dir) + if __name__ == "__main__": unittest.main()