[JTH] add more core functions to keep hywaves clean, and also re-use …

bluemath_tk/core/metrics.py

@@ -28,6 +28,7 @@
 from collections.abc import Callable, Iterable
 
 import numpy as np
+import pandas as pd
 
 ArrayLike = np.ndarray | Iterable[float]
 
@@ -58,7 +59,9 @@
 
 
 def paired_finite(
-    reference: ArrayLike, estimate: ArrayLike
+    reference: ArrayLike,
+    estimate: ArrayLike,
+    bounds: tuple[float, float] | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Flatten two paired arrays and keep only the pairs finite on both sides.
@@ -69,6 +72,9 @@ def paired_finite(
         Reference values.
     estimate : array-like
         Estimated values, paired element-wise with *reference*.
+    bounds : tuple of float, optional
+        ``(low, high)``: also drop pairs with either side outside it, e.g.
+        corrupted-but-finite model output. Default is None (finite only).
 
     Returns
     -------
@@ -89,6 +95,9 @@ def paired_finite(
             f"({ref.size} != {est.size})"
         )
     mask = np.isfinite(ref) & np.isfinite(est)
+    if bounds is not None:
+        low, high = bounds
+        mask &= (ref >= low) & (ref <= high) & (est >= low) & (est <= high)
 
     return ref[mask], est[mask]
 
@@ -639,6 +648,7 @@ def compute_metrics(
     *,
     circular: bool = False,
     rrmse_normalization: str = "mean",
+    bounds: tuple[float, float] | None = None,
 ) -> dict[str, float]:
     """
     Compute several metrics at once on the finite pairs of two arrays.
@@ -658,6 +668,8 @@ def compute_metrics(
         ``NaN`` so tables keep the same columns. Default is False.
     rrmse_normalization : {"mean", "rms"}, optional
         Passed to :func:`rrmse`. Default is "mean".
+    bounds : tuple of float, optional
+        Passed to :func:`paired_finite`. Default is None.
 
     Returns
     -------
@@ -682,7 +694,7 @@ def compute_metrics(
             f"Unknown metric(s) {unknown}; available: {list(AVAILABLE_METRICS)}"
         )
 
-    ref, est = paired_finite(reference, estimate)
+    ref, est = paired_finite(reference, estimate, bounds)
     out: dict[str, float] = {"n": int(ref.size)}
     for name in names:
         if ref.size == 0:
@@ -696,3 +708,49 @@ def compute_metrics(
             out[name] = _LINEAR[name](ref, est)
 
     return out
+
+
+def pool_metrics(
+    table: pd.DataFrame,
+    by: Iterable[str],
+    metrics: Iterable[str] = ("bias", "rmse", "rrmse", "r2"),
+) -> pd.DataFrame:
+    """
+    Pool per-group metrics into one row per *by* group, weighted by ``n``.
+
+    ``bias`` (n-weighted mean) and ``rmse`` (root of the n-weighted mean
+    square) are the exact pooled values. Ratio and correlation metrics
+    (``rrmse``, ``r2``, ...) need the raw pairs to be pooled exactly, so they
+    are n-weighted means, an approximation.
+
+    Parameters
+    ----------
+    table : pd.DataFrame
+        One row per group, with an ``n`` column and the *metrics* columns.
+    by : iterable of str
+        Columns that define the pooled rows.
+    metrics : iterable of str, optional
+        Metric columns to pool.
+
+    Returns
+    -------
+    pd.DataFrame
+        The *by* columns, the pooled *metrics* and the summed ``n``.
+    """
+
+    by, metrics = list(by), list(metrics)
+
+    def pool(group: pd.DataFrame) -> pd.Series:
+        w = group["n"].to_numpy(dtype=float)
+        if w.sum() <= 0:
+            return pd.Series({m: np.nan for m in metrics} | {"n": 0})
+        out = {
+            m: float(np.sqrt(np.average(group[m] ** 2, weights=w)))
+            if m == "rmse"
+            else float(np.average(group[m], weights=w))
+            for m in metrics
+        }
+        return pd.Series(out | {"n": int(w.sum())})
+
+    out = table.groupby(by, sort=False)[["n", *metrics]].apply(pool).reset_index()
+    return out.astype({"n": int})

bluemath_tk/core/models.py

@@ -20,7 +20,7 @@
 from sklearn.preprocessing import StandardScaler
 
 from .constants import EARTH_RADIUS
-from .logging import get_file_logger
+from .logging import _coerce_level, _console_handlers, get_file_logger
 from .operations import (
     denormalize,
     destandarize,
@@ -31,6 +31,36 @@
 )
 
 
+#: Console level of every model logger (PCA, MDA, GP, ...); see :func:`set_models_console_level`.
+MODELS_CONSOLE_LEVEL: Union[int, str] = "WARNING"
+
+
+def set_models_console_level(level: Union[int, str]) -> None:
+    """
+    Console level of the model loggers, existing and future (files keep their level).
+
+    Models log routine messages on every fit; ``"ERROR"`` keeps a notebook
+    that fits many models readable while ``logs/*.log`` keep the detail.
+
+    Parameters
+    ----------
+    level : int or str
+        Logging level, e.g. ``"ERROR"``.
+    """
+
+    global MODELS_CONSOLE_LEVEL
+    MODELS_CONSOLE_LEVEL = level
+    names = set()
+    stack = [BlueMathModel]
+    while stack:
+        cls = stack.pop()
+        names.add(cls.__name__)
+        stack.extend(cls.__subclasses__())
+    for name in names & set(logging.Logger.manager.loggerDict):
+        for handler in _console_handlers(logging.getLogger(name)):
+            handler.setLevel(_coerce_level(level))
+
+
 class BlueMathModel(ABC):
     """
     Abstract base class for handling default functionalities across the project.
@@ -148,8 +178,13 @@ def logger(self) -> logging.Logger:
         - Thread-safe logger creation
         """
 
-        if self._logger is None:
-            self._logger = get_file_logger(name=self.__class__.__name__)
+        # an unpickled model refers to its logger by name: in a new process
+        # that logger has no handlers yet
+        if self._logger is None or not self._logger.handlers:
+            self._logger = get_file_logger(
+                name=self._logger.name if self._logger else self.__class__.__name__,
+                console_level=MODELS_CONSOLE_LEVEL,
+            )
         return self._logger
 
     @logger.setter
@@ -200,7 +235,9 @@ def set_logger_name(
         >>> model.set_logger_name("my_model", level="DEBUG", console=False)
         """
 
-        self.logger = get_file_logger(name=name, level=level, console=console)
+        self.logger = get_file_logger(
+            name=name, level=level, console=console, console_level=MODELS_CONSOLE_LEVEL
+        )
 
     def save_model(self, model_path: str, exclude_attributes: List[str] = None) -> None:
         """

bluemath_tk/core/plotting/labels.py

@@ -0,0 +1,89 @@
+"""
+Point labels that do not overlap each other.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from matplotlib.axes import Axes
+
+#: Candidate offsets (points, ha, va): 8 compass directions per ring, near to far.
+_OFFSETS: list[tuple[tuple[float, float], str, str]] = [
+    ((radius * ux / ((ux**2 + uy**2) ** 0.5), radius * uy / ((ux**2 + uy**2) ** 0.5)), ha, va)
+    for radius in (6, 11, 16, 22, 28, 34, 42, 55)
+    for ux, uy, ha, va in (
+        (1, 1, "left", "bottom"),
+        (1, -1, "left", "top"),
+        (-1, 1, "right", "bottom"),
+        (-1, -1, "right", "top"),
+        (0, 1, "center", "bottom"),
+        (0, -1, "center", "top"),
+        (1, 0, "left", "center"),
+        (-1, 0, "right", "center"),
+    )
+]
+
+
+def _overlap(a: Sequence[float], b: Sequence[float], margin: float) -> bool:
+    return not (
+        a[2] + margin < b[0] or b[2] + margin < a[0] or a[3] + margin < b[1] or b[3] + margin < a[1]
+    )
+
+
+def place_labels(
+    ax: Axes,
+    x: Sequence[float],
+    y: Sequence[float],
+    texts: Sequence[str],
+    *,
+    leader_colors: Sequence[str] | None = None,
+    **text_kwargs,
+) -> None:
+    """
+    Annotate points, each at the nearest offset that clears the labels already placed.
+
+    The axes limits must be final (the label sizes are measured on the
+    rendered figure). A label pushed far from its point gets a leader line.
+
+    Parameters
+    ----------
+    ax : Axes
+        Axes with its final limits.
+    x, y : sequence of float
+        Points in data coordinates.
+    texts : sequence of str
+        One label per point.
+    leader_colors : sequence of str, optional
+        Leader-line colour per point. Default is grey.
+    **text_kwargs
+        Passed to ``ax.annotate`` (``fontsize``, ``path_effects``, ...).
+    """
+
+    ax.figure.canvas.draw()
+    renderer = ax.figure.canvas.get_renderer()
+    inv = ax.transData.inverted()
+    pad = abs(inv.transform((3.5 * ax.figure.dpi / 72.0, 0))[0] - inv.transform((0, 0))[0])
+    placed: list[tuple[float, float, float, float]] = []
+    for i, (px, py, text) in enumerate(zip(x, y, texts)):
+        for offset, ha, va in _OFFSETS:
+            probe = ax.annotate(text, (px, py), xytext=offset, textcoords="offset points",
+                                ha=ha, va=va, **text_kwargs)
+            bbox = probe.get_window_extent(renderer)
+            probe.remove()
+            (x0, y0), (x1, y1) = inv.transform((bbox.x0, bbox.y0)), inv.transform((bbox.x1, bbox.y1))
+            box = (min(x0, x1), min(y0, y1), max(x0, x1), max(y0, y1))
+            if not any(_overlap(box, other, pad) for other in placed):
+                break
+        placed.append(box)
+        leader = (offset[0] ** 2 + offset[1] ** 2) ** 0.5 > 9
+        color = leader_colors[i] if leader_colors is not None else "0.4"
+        ax.annotate(
+            text, (px, py), xytext=offset, textcoords="offset points", ha=ha, va=va,
+            arrowprops=dict(arrowstyle="-", color=color, lw=0.7, alpha=0.75, shrinkA=0, shrinkB=3)
+            if leader else None,
+            **text_kwargs,
+        )
+
+
+__all__ = ["place_labels"]

bluemath_tk/core/plotting/scatter.py

@@ -126,6 +126,7 @@ def validation_scatter(
     seed: int = 0,
     fmt: str = ".3f",
     rrmse_normalization: str = "mean",
+    bounds: tuple[float, float] | None = None,
 ) -> dict[str, float]:
     """
     Density scatter of an estimate against its reference, with metrics.
@@ -177,6 +178,9 @@ def validation_scatter(
         Seed for the drawing subsample. Default is 0.
     fmt : str, optional
         Number format for the annotation. Default is ".3f".
+    bounds : tuple of float, optional
+        Drop pairs outside ``(low, high)`` before plotting and scoring (see
+        :func:`bluemath_tk.core.metrics.paired_finite`).
     rrmse_normalization : {"mean", "rms"}, optional
         Passed to :func:`bluemath_tk.core.metrics.rrmse`. Default is "mean".
 
@@ -186,7 +190,7 @@ def validation_scatter(
         The metrics computed on all finite pairs (always includes ``n``).
     """
 
-    ref, est = paired_finite(reference, estimate)
+    ref, est = paired_finite(reference, estimate, bounds)
     if circular:
         ref, est = np.mod(ref, 360.0), np.mod(est, 360.0)
 

bluemath_tk/interpolation/gps.py

@@ -683,6 +683,29 @@ def fit(
         self.is_fitted = True
         self.logger.info("GP models fitted successfully")
 
+    def to(self, device: str | torch.device) -> "ExactGPInterpolation":
+        """
+        Move the fitted models and likelihoods to *device* (e.g. ``"cpu"``).
+
+        Parameters
+        ----------
+        device : str or torch.device
+            Target device for prediction.
+
+        Returns
+        -------
+        ExactGPInterpolation
+            This instance.
+        """
+
+        self.device = torch.device(device)
+        for model in self._models.values():
+            model.to(self.device)
+        for likelihood in self._likelihoods.values():
+            likelihood.to(self.device)
+
+        return self
+
     def predict(
         self,
         dataset: pd.DataFrame,