model.runtime.runners.gauge_learning_runner

Gauge-learning runner for ADELM: differentiable calibration of the coupled land + river-routing model against gauge discharge.

The land domain is the union of the selected gauges’ upstream basins (a :func:model.environment.river_network.build_training_domain dict); the loss is the basin-averaged NSE between each gauge cell’s streamflow_mm_day and the observed discharge (both normalised by the same model upstream area; see set_gauge_targets).

Unlike the site-learning runner, the full forcing record is never held in memory. Training streams one time chunk of drivers from disk at a time and truncates backpropagation at the chunk boundary (TBPTT), so the peak memory is O(chunk x n_cells) independent of the total record length. This is what makes a continental domain over a multi-decade record tractable.

Module Contents

Classes

GaugeEvaluation	Result of a no-grad streaming forward over the calibration window.
GaugeLearningRunner	Calibrate the coupled land+routing model on a basin domain against gauge Q.

Data

__all__

API

class model.runtime.runners.gauge_learning_runner.GaugeEvaluation

Result of a no-grad streaming forward over the calibration window.

loss: torch.Tensor

None

predictions: torch.Tensor

None

gauge_ids: list

‘field(…)’

class model.runtime.runners.gauge_learning_runner.GaugeLearningRunner(config, model, site_ids=None, time_values=None, data=None)

Bases: model.runtime.runners.grid_simulation_runner.GridSimulationRunner

Calibrate the coupled land+routing model on a basin domain against gauge Q.

Each gauge is an entity the spatial cross-validation machinery (:class:model.utils.CrossValidation) can fold over: a fold masks out its held-out gauges by writing NaN into the target, which :func:basin_nse_loss ignores.

Initialization

set_gauge_targets(domain, obs, months)

Bind the basin domain and the gauge discharge targets.

Parameters

domain : dict Output of :func:build_training_domain (sets the routing geometry and the forcing cells via :meth:set_training_domain). obs : dict Output of :func:load_gauge_discharge; discharge [n_g, T] and std [n_g] are aligned to obs['gauge_ids']. months : sequence[(int, int)] (year, month) pairs whose concatenated days equal the target time axis T. Drivers are streamed one chunk of these at a time.

train(optimizer, targets=None, epochs=1, epoch_offset=0, train_chunk_months=1, max_grad_norm=1.0, skip_nan_grads=True, show_progress=True, mem_probe=False, **_ignored)

Run epochs streaming-TBPTT passes over the calibration window.

Each pass walks the record in chunks of train_chunk_months: load that chunk’s drivers from disk, forward from the carried (detached) state, compute the configured gauge loss over the chunk, backprop, step, and detach. With a season-spanning chunk the loss sees the full seasonal cycle; peak memory stays O(chunk x n_cells).

evaluate(targets=None, stage=None)

Streaming no-grad forward over the window; returns loss + predictions.

preload_and_freeze()

Pre-load land weights and freeze the configured parameters.

Returns (trainable_tensors, open_param_names, frozen_param_names). The open parameters are the nn_global/nn_feature_based params that are NOT frozen – these are what the optimiser trains.

manages_own_data

True

prepare_run(workflow, run_dir)

describe_run(workflow, train_targets, val_targets)

fit(workflow, train_targets, val_targets, *, epochs, lr, train_chunk_size, max_grad_norm, show_progress, debug)

evaluate_run(workflow, results, eval_targets=None)

save_run_outputs(workflow, run_dir, evaluation, results, test_targets=None)

property entity_ids

property cv_config

fold_metadata_lines(fold)

evaluate_only(workflow, run_dir)

Global per-gauge evaluation of the saved checkpoint (no training).

evaluate_global sets up its own basin window and loads the saved checkpoint, so no separate prepare_run is needed here.

evaluate_global(checkpoint_path=None)

Per-gauge NSE/KGE eval over the calibration window for all gauges; writes a CSV.

model.runtime.runners.gauge_learning_runner.__all__

[‘GaugeLearningRunner’, ‘GaugeEvaluation’]