Hywaves/fast reconstruct

bluemath_tk/interpolation/gps.py

@@ -711,6 +711,7 @@ def predict(
         dataset: pd.DataFrame,
         return_std: bool = False,
         verbose: int = 1,
+        batch_size: int | None = None,
     ) -> pd.DataFrame:
         """
         Predict using the fitted GP model.
@@ -723,6 +724,12 @@ def predict(
             If True, returns standard deviations. Default is False.
         verbose : int, optional
             Verbosity level. Default is 1.
+        batch_size : int, optional
+            Rows predicted per kernel evaluation. Every row is independent,
+            so the result is the same (to float32 rounding); a kernel of a few
+            hundred rows x the training points stays in the CPU cache, which
+            was ~8x faster than all rows at once for 6,648 rows x 550 training
+            points (512 rows; 256 was slower). Default: all rows at once.
 
         Returns
         -------
@@ -760,27 +767,26 @@ def predict(
             model.eval()
             likelihood.eval()
 
+            step = batch_size or max(1, len(X_tensor))
+            means, lowers, uppers = [], [], []
             with torch.no_grad(), gpytorch.settings.fast_pred_var():
-                pred_dist = likelihood(model(X_tensor))
-                predictions_dict[target_var] = pred_dist.mean.cpu().numpy()
-                if return_std:
-                    # stds_dict[f"{target_var}_std"] = pred_dist.stddev.cpu().numpy()
-                    # self._target_scale_factor[f"{target_var}_std"] = (
-                    #     self._target_scale_factor[target_var]
-                    # )
-                    (
-                        stds_dict[f"{target_var}_lower_ci"],
-                        stds_dict[f"{target_var}_upper_ci"],
-                    ) = (
-                        pred_dist.confidence_region()[0].cpu().numpy(),
-                        pred_dist.confidence_region()[1].cpu().numpy(),
-                    )
-                    self._target_scale_factor[f"{target_var}_lower_ci"] = (
-                        self._target_scale_factor[target_var]
-                    )
-                    self._target_scale_factor[f"{target_var}_upper_ci"] = (
-                        self._target_scale_factor[target_var]
-                    )
+                for start in range(0, len(X_tensor), step):
+                    pred_dist = likelihood(model(X_tensor[start : start + step]))
+                    means.append(pred_dist.mean.cpu().numpy())
+                    if return_std:
+                        lower, upper = pred_dist.confidence_region()
+                        lowers.append(lower.cpu().numpy())
+                        uppers.append(upper.cpu().numpy())
+            predictions_dict[target_var] = np.concatenate(means)
+            if return_std:
+                stds_dict[f"{target_var}_lower_ci"] = np.concatenate(lowers)
+                stds_dict[f"{target_var}_upper_ci"] = np.concatenate(uppers)
+                self._target_scale_factor[f"{target_var}_lower_ci"] = (
+                    self._target_scale_factor[target_var]
+                )
+                self._target_scale_factor[f"{target_var}_upper_ci"] = (
+                    self._target_scale_factor[target_var]
+                )
 
         # Convert to DataFrame
         result = pd.DataFrame(predictions_dict)

bluemath_tk/waves/hywaves/metamodel.py

@@ -13,7 +13,9 @@
    (``dir`` as ``dir_u``/``dir_v``).
 4. :func:`fit_gp` - an exact GP from the forcing to all stacked PCs.
 5. :func:`predict_fields` - GP -> PCA inverse -> fields, with physical floors
-   (:data:`MIN_HS`, :data:`MIN_DIRECTIONAL_SPREAD_DEG`).
+   (:data:`MIN_HS`, :data:`MIN_DIRECTIONAL_SPREAD_DEG`); the same in two
+   steps, :func:`predict_pcs` (GP, the costly part, run in parallel
+   processes by a reconstruction) and :func:`fields_from_pcs`.
 
 Offshore hindcasts enter through :func:`goal_forcing` (bulk or one spectral
 partition, Hs included) and :func:`valid_timesteps`; the predicted fields of
@@ -56,6 +58,11 @@
 # summed.
 MIN_HS = 0.0
 
+# Rows (time steps) per GP kernel evaluation in predict_fields: a kernel of
+# 512 rows x the training points stays in the CPU cache; one year (6,648
+# valid hours, MDA 550) was ~8x faster than in one evaluation.
+PREDICT_BATCH_ROWS = 512
+
 #: 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_zero_energy` the
 #: training fields have no NaN, so this is only a safety net.
@@ -522,6 +529,30 @@ def fit_gp(
 # ---------------------------------------------------------------------------
 
 
+def predict_pcs(gp: ExactGPInterpolation, forcing: pd.DataFrame) -> np.ndarray:
+    """
+    Stacked PCs for each forcing row (the GP step of :func:`predict_fields`).
+
+    Rows are predicted :data:`PREDICT_BATCH_ROWS` at a time.
+
+    Parameters
+    ----------
+    gp : ExactGPInterpolation
+        Fitted GP (see :func:`fit_gp`).
+    forcing : pd.DataFrame
+        Forcing rows, columns as at fit time.
+
+    Returns
+    -------
+    np.ndarray
+        ``(rows, stacked PCs)``, in :func:`stack_pcs` order.
+    """
+
+    return gp.predict(
+        dataset=forcing.rename_axis("case_num"), batch_size=PREDICT_BATCH_ROWS, verbose=0
+    ).to_numpy()
+
+
 def predict_fields(
     gp: ExactGPInterpolation,
     pcas: Mapping[str, PCA],
@@ -532,14 +563,31 @@ def predict_fields(
     """
     Nearshore fields for each forcing row: GP -> PCA inverse.
 
+    :func:`predict_pcs` then :func:`fields_from_pcs`; see the latter for the
+    fields returned.
+    """
+
+    return fields_from_pcs(predict_pcs(gp, forcing), pcas, forcing, vars_to_predict, sites)
+
+
+def fields_from_pcs(
+    pcs: np.ndarray,
+    pcas: Mapping[str, PCA],
+    forcing: pd.DataFrame,
+    vars_to_predict: Mapping[str, Any],
+    sites: list[str],
+) -> xr.Dataset:
+    """
+    Nearshore fields from predicted PCs (the PCA step of :func:`predict_fields`).
+
     ``dir`` is rebuilt from ``dir_u``/``dir_v``; ``raw`` variables are copied
     from the forcing to every site; ``hs`` and ``spr`` are floored at
     :data:`MIN_HS` and :data:`MIN_DIRECTIONAL_SPREAD_DEG`.
 
     Parameters
     ----------
-    gp : ExactGPInterpolation
-        Fitted GP (see :func:`fit_gp`).
+    pcs : np.ndarray
+        :func:`predict_pcs` of *forcing*.
     pcas : mapping
         Fitted PCAs by target name.
     forcing : pd.DataFrame
@@ -560,7 +608,7 @@ def predict_fields(
     row_dim = forcing.index.name or "case_num"
     forcing = forcing.rename_axis("case_num")
     rows = forcing.index.values
-    blocks = split_pcs(gp.predict(dataset=forcing).values, pcas)
+    blocks = split_pcs(pcs, pcas)
 
     out: dict[str, xr.DataArray] = {}
     for name, spec in vars_to_predict.items():

bluemath_tk/waves/hywaves/resources.py

@@ -29,8 +29,10 @@
   processes (:class:`FitResources`); a fit needs <= ~1.2 GiB.
 - **Reconstruction**, one call = *months* x *sites*:
 
-  1. goal contributions: GP + PCA inverse per goal and partition, threaded,
-     in the calling process (``device``, ``threads``);
+  1. goal contributions: the GP of every goal x partition in its own
+     process (``threads // 2`` processes of 2 threads, 512-row kernels;
+     ~5x faster than one threaded process), then the PCA inverse to the
+     sites, threaded, in the calling process (``device``, ``threads``);
   2. spectral sum: JONSWAP x Cartwright spectra summed per site and hour in
      ``workers`` single-thread processes. Every site and hour is
      independent, so the sum is split into site x time blocks
@@ -76,16 +78,24 @@
 #: Memory of a sum worker process itself (imports, its share of the inputs).
 WORKER_BYTES = 3 * GIB // 4
 
-# Xeon Gold 5418Y (2026-10-15). One core: building and summing the spectra per
-# step x site x contribution, plus the bulk stats per step x site. Workers
-# slow each other down (memory bandwidth): speedup = w / (1 + _CONTENTION w),
-# measured 7.1 / 10.9 / 14.1 with 8 / 18 / 36 workers. GP + PCA (36 threads):
-# a fixed part per month x contribution (models, kernel) plus one per site.
+# Xeon Gold 5418Y. Sum (2026-10-15), one core: building and summing the
+# spectra per step x site x contribution, plus the bulk stats per step x site;
+# workers slow each other down (memory bandwidth): speedup = w / (1 +
+# _CONTENTION w), measured 7.1 / 10.9 / 14.1 with 8 / 18 / 36 workers.
+# GP step (2026-10-16, MDA 550, 18 processes x 2 threads, 512-row kernels):
+# ~6 s per call to start the processes, then per month x contribution (GP)
+# and per site x month x contribution (PCA inverse, 36 threads).
 _SUM_SECONDS_PER_PART = 2.2e-5
 _SUM_SECONDS_PER_SITE = 1.0e-4
 _CONTENTION = 0.04
-_GP_SECONDS_PER_MONTH_PART = 0.125
-_GP_SECONDS_PER_SITE_MONTH_PART = 1.9e-4
+_GP_SECONDS_PER_CALL = 6.0
+_GP_SECONDS_PER_MONTH_PART = 0.05
+_GP_SECONDS_PER_SITE_MONTH_PART = 2.9e-4
+
+#: Cost of a sum worker pool (s): starting the processes (~6 s) and sending
+#: them the contributions; one site-year sum took 11-18 s in a pool for
+#: ~0.5 s of work.
+POOL_SECONDS = 10.0
 
 
 @dataclass
@@ -234,13 +244,49 @@ def sum_budget(memory: int, held: int, workers: int) -> int:
     return max(0, min(workers * SUM_BYTES_PER_WORKER, memory - held - workers * WORKER_BYTES))
 
 
+def _single_core_seconds(sites: int, steps: int, parts: int) -> float:
+    return sites * steps * (_SUM_SECONDS_PER_PART * parts + _SUM_SECONDS_PER_SITE)
+
+
+def _speedup(workers: int) -> float:
+    return workers / (1 + _CONTENTION * workers)
+
+
+def sum_workers(sites: int, steps: int, parts: int, workers: int) -> int:
+    """
+    Workers worth using for one spectral sum.
+
+    *workers* when the time they save beats :data:`POOL_SECONDS`, else 1
+    (the sum runs in the calling process; e.g. one site over a year).
+
+    Parameters
+    ----------
+    sites, steps : int
+        Sites and time steps summed.
+    parts : int
+        Contributions per site.
+    workers : int
+        Workers available (``ReconstructionResources.workers``).
+
+    Returns
+    -------
+    int
+        *workers* or 1.
+    """
+
+    single = _single_core_seconds(sites, steps, parts)
+    saved = single - single / _speedup(max(1, workers))
+
+    return workers if saved > POOL_SECONDS else 1
+
+
 def hours(site_months: int, months: int, parts: int, workers: int) -> float:
     """
     Rough wall time (hours) of a reconstruction.
 
-    From rates measured on one 48-core Xeon Gold 5418Y node (GP + PCA with
-    36 threads, the sum slowed by memory bandwidth as workers grow); other
-    machines differ, so take it as an order of magnitude.
+    From rates measured on one 48-core Xeon Gold 5418Y node (GP processes and
+    PCA, the sum slowed by memory bandwidth as workers grow); other machines
+    differ, so take it as an order of magnitude.
 
     Parameters
     ----------
@@ -259,12 +305,12 @@ def hours(site_months: int, months: int, parts: int, workers: int) -> float:
         Hours.
     """
 
-    workers = max(1, workers)
-    speedup = workers / (1 + _CONTENTION * workers)
-    per_site_month = STEPS_PER_MONTH * (_SUM_SECONDS_PER_PART * parts + _SUM_SECONDS_PER_SITE)
-    seconds = site_months * (per_site_month / speedup + _GP_SECONDS_PER_SITE_MONTH_PART * parts)
+    per_site_month = _single_core_seconds(1, STEPS_PER_MONTH, parts)
+    seconds = site_months * (per_site_month / _speedup(max(1, workers)) + _GP_SECONDS_PER_SITE_MONTH_PART * parts)
+    calls = -(-months // MONTHS_PER_CALL["cpu"])
+    seconds += calls * _GP_SECONDS_PER_CALL + months * parts * _GP_SECONDS_PER_MONTH_PART
 
-    return (seconds + months * parts * _GP_SECONDS_PER_MONTH_PART) / 3600
+    return seconds / 3600
 
 
 def _halvings(n: int) -> list[int]:
@@ -377,6 +423,7 @@ def plan_fit(
     "STEPS_PER_MONTH",
     "FitResources",
     "MIN_SUM_BYTES_PER_WORKER",
+    "POOL_SECONDS",
     "ReconstructionResources",
     "SUM_BYTES_PER_WORKER",
     "WORKER_BYTES",
@@ -385,4 +432,5 @@ def plan_fit(
     "plan_fit",
     "plan_reconstruction",
     "sum_budget",
+    "sum_workers",
 ]

tests/waves/test_hywaves_resources.py

@@ -35,6 +35,10 @@ def test_sum_budget_and_hours(self):
         self.assertEqual(resources.sum_budget(GIB, 2 * GIB, 4), 0)
         self.assertLess(resources.hours(100, 1, 10, 36), resources.hours(100, 1, 10, 1))
 
+    def test_small_sums_skip_the_pool(self):
+        self.assertEqual(resources.sum_workers(1, 8760, 10, 36), 1)
+        self.assertEqual(resources.sum_workers(4971, 744, 10, 36), 36)
+
     def test_gpu_months_and_fits(self):
         gpu = Machine(16, 100 * GIB, 128 * GIB, "A100", 40 * GIB)
         self.assertEqual(resources.plan_reconstruction(38, 24, 10, host=gpu).months, 4)