model.utils.training

Training loss utilities for ADELM.

Module Contents

Functions

sample_r2_loss	Compute the sample-level 1 - R² loss by pooling all finite site-day values.
basin_nse_loss	Basin-averaged NSE loss (Kratzert et al., 2019).
kge_loss	Kling–Gupta-style composite loss over the full series, per basin then averaged.
sample_r2_score	Compute sample-level R² over all finite positions.
site_r2_loss	Compute the site-level 1 - R² loss on per-site multi-year means.
site_r2_score	Compute R² on per-site means over valid samples.

API

model.utils.training.sample_r2_loss(pred: Tensor, target: Tensor)

Compute the sample-level 1 - R² loss by pooling all finite site-day values.

This is the primary training loss. It constrains day-to-day and seasonal variability within sites. Use together with site_r2_loss() when cross-site climatological differences also matter.

model.utils.training.basin_nse_loss(pred: Tensor, target: Tensor, target_std: Tensor, eps: float = 0.1)

Basin-averaged NSE loss (Kratzert et al., 2019).

Weighted mean squared error with a per-basin weight 1/(std + eps)**2, where std is the precomputed observed standard deviation of each basin. The weight normalises every basin’s errors to a comparable scale (so large rivers do not dominate), giving an NSE-like objective that – unlike NSE itself – is a plain additive mean and therefore composes cleanly across time chunks and basins (suitable for chunked backpropagation).

Parameters:

• pred (torch.Tensor) – [n_basin, T] modelled and observed series. Non-finite modelled or observed entries are ignored.

• target (torch.Tensor) – [n_basin, T] modelled and observed series. Non-finite modelled or observed entries are ignored.

• target_std (torch.Tensor) – [n_basin] per-basin observed std (precomputed; treated as constant).

• eps (float) – Stabiliser for near-constant basins.

Returns:

Scalar loss, or None if target has no finite values.

Return type:

torch.Tensor or None

model.utils.training.kge_loss(pred: Tensor, target: Tensor, weights=(1.0, 1.0, 1.0), min_days: int = 30, eps: float = 1e-06)

Kling–Gupta-style composite loss over the full series, per basin then averaged.

For each basin the modelled and observed series are compared through three scale-separated terms – the linear correlation r (timing and seasonal phase), the ratio of standard deviations alpha (variability), and the ratio of means beta (volume bias):

L_i = w_r (1 - r)^2 + w_a (1 - alpha)^2 + w_b (1 - beta)^2

The loss is the mean of L_i over basins with at least min_days finite observations. Because alpha penalises under-dispersion, a flattened series is no longer the cheapest fit (unlike a plain squared error), and the correlation term gives a gradient on seasonal timing – but only when the series passed in spans long enough to resolve the seasonal cycle.

Parameters:

• pred (torch.Tensor) – [n_basin, T] modelled and observed series. Non-finite modelled or observed entries are dropped per basin.

• target (torch.Tensor) – [n_basin, T] modelled and observed series. Non-finite modelled or observed entries are dropped per basin.

• weights (tuple of float) – (w_r, w_a, w_b) weights on the correlation, variability, and bias terms.

• min_days (int) – Minimum finite observations a basin needs to contribute a stable term.

• eps (float) – Stabiliser for the standard deviations and ratios.

Returns:

Scalar loss, or None if no basin has enough finite observations.

Return type:

torch.Tensor or None

model.utils.training.sample_r2_score(pred: Tensor, target: Tensor)

Compute sample-level R² over all finite positions.

model.utils.training.site_r2_loss(pred: Tensor, target: Tensor, min_valid_sites: int = 1, min_samples_per_site: int = 1)

Compute the site-level 1 - R² loss on per-site multi-year means.

Reduces pred and target from [site, ...] to one mean per site using only finite positions, then applies sample_r2_loss() to the reduced vectors. This constrains the cross-site climatological gradient, which is the signal most relevant to learning attribute-to-parameter mappings that transfer in space.

Sites with fewer than min_samples_per_site valid samples are excluded. Returns None when fewer than min_valid_sites sites remain, so the caller can skip the gradient step.

model.utils.training.site_r2_score(pred: Tensor, target: Tensor, min_valid_sites: int = 1, min_samples_per_site: int = 1)

Compute R² on per-site means over valid samples.