From 5a8f44ed85c98f3bad2af84fe883499578148d10 Mon Sep 17 00:00:00 2001 From: tausiaj Date: Mon, 26 Oct 2026 09:52:51 +0100 Subject: [PATCH 1/2] [JTH] working version for intelligent hyglow reconstruction --- bluemath_tk/interpolation/gps.py | 46 ++++++++------- bluemath_tk/waves/hywaves/metamodel.py | 56 ++++++++++++++++-- bluemath_tk/waves/hywaves/resources.py | 82 ++++++++++++++++++++------ 3 files changed, 143 insertions(+), 41 deletions(-) diff --git a/bluemath_tk/interpolation/gps.py b/bluemath_tk/interpolation/gps.py index 9fdc764..d7ccb38 100644 --- a/bluemath_tk/interpolation/gps.py +++ b/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) diff --git a/bluemath_tk/waves/hywaves/metamodel.py b/bluemath_tk/waves/hywaves/metamodel.py index 00b0d33..23c629b 100644 --- a/bluemath_tk/waves/hywaves/metamodel.py +++ b/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(): diff --git a/bluemath_tk/waves/hywaves/resources.py b/bluemath_tk/waves/hywaves/resources.py index 134b195..ce8d6ba 100644 --- a/bluemath_tk/waves/hywaves/resources.py +++ b/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", ] From 51a10914ae9906f8304a5e6a98f384c88f78a300 Mon Sep 17 00:00:00 2001 From: tausiaj Date: Mon, 26 Oct 2026 09:53:07 +0100 Subject: [PATCH 2/2] [JTH] working version for intelligent hyglow reconstruction tests --- tests/waves/test_hywaves_resources.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/waves/test_hywaves_resources.py b/tests/waves/test_hywaves_resources.py index 4e3ad67..d18c307 100644 --- a/tests/waves/test_hywaves_resources.py +++ b/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)