Investigate GPU/fast-calc candidates for hazard indices + bias correction (CuPy/xclim-on-GPU)

[peetmate]

Context

Benchmark shared by colleague (their numbers, not ours) — climate downscaling + bias correction via quantile mapping, xarray + CuPy, 4.15 GB / 1000 days daily, 0.25°→0.125°. "Your box" = CG labs, which has no GPUs, so its 40-thread CPU row is the ceiling there; any GPU path must come from Kaggle / SCIO-paid GPUs / the Alliance MPL GPU cluster:

Hardware	Downscale	Quantile map	vs best CPU
Kaggle ~4 vCPU (dask)	~99 s	~805 s	—
your box 40 threads	11.3 s	23.4 s	1.0× / 1.0×
1× T4 (free Kaggle)	2.86 s	16.3 s	4.0× / 1.4×
2× T4	1.44 s	8.48 s	7.9× / 2.8×
1× P100	1.36 s	9.24 s	8.3× / 2.5×

A single free Kaggle GPU beats all CG-labs cores. Paid GPUs for a week (via SCIO) or the Alliance MPL GPU cluster are cheap by comparison. Caveat: GPU path is Python (CuPy), and downscaling needs some hand-written CuPy (no xsdba/xclim GPU path for SDBA quantile-mapping yet); but most xclim indices run directly on GPU.

Task

Investigate which parts of the hazards producer (hazards_upstream/R/04_indices + 03_bias_correction) would actually benefit from a GPU / Python-CuPy fast path, and quantify the win vs porting cost. Ties into converging on Python/xclim (see CDH climate-toolkit).

Candidates to assess (seed analysis)

1. Bias correction + downscaling (03) — quantile mapping. The benchmark IS this; proven 2.5–8× GPU win. Highest-value target if we run downscaling/bias-correction ourselves.
2. fast_calc water-balance trio — NDWS / NDWL0 / NDWL50 (04) — by far the slowest indices (~820–835 s/GCM last bake vs 12–42 s for the rest). Daily EABYEP water balance is per-cell independent, sequential only in time → ideal for GPU vectorisation across cells per day-step. The sequential AVAIL state carry (see hazards#19) constrains naive parallelism but not per-day cell vectorisation.
3. Cheap daily reductions — NDD, NTx, PTOT, TAVG, TMAX, TMIN, TAI, THI, HSH — already 12–42 s/GCM. Embarrassingly parallel; GPU helps but marginal in absolute terms. Most map to native xclim indices that run on GPU, so they'd come "for free" if index calc moves to xclim.

Deliverable

• Per-stage benefit/cost table: current CPU time, expected GPU speedup, porting effort (R→Python/CuPy, hand-written vs xclim-native), state/sequencing constraints.
• Recommendation on what (if anything) to move to GPU first, and whether it justifies the SCIO / Alliance-MPL GPU resource ask.
• Decision input for the broader hazards-engine convergence on Python/xclim.

Filed alongside the methods review (#14) and the CDH/xclim convergence direction.

[bjyberg]

I ran a few more tests yesterday on the hazards indices and it is worth noting that some of them also would require us to do our own cupy/numpy versions. The benefits were also not quite as noticeable as with the downscaling and bias correction. However the tests were only for a few simple temperature indices, so it would be worth testing a few more of the complex and slow ones (ndlw0 I believe) to see the benefit

[bjyberg]

I will share my code for all of these benchmarks in the CDH org so thy can be modified and re-tested