Usage Guide

This guide demonstrates how to use the torchhydro library to configure and run hydrological models. The core workflow revolves around defining a set of parameters, updating a default configuration, and launching the training and evaluation process with a single function.

Core Concept

The main workflow consists of three steps:

1. Define Parameters: You start by defining all your experiment's parameters (like model choice, dataset, variables, and hyperparameters) using the torchhydro.configs.config.cmd function. This creates a parameter object.
2. Update Configuration: The default configuration is loaded using torchhydro.configs.config.default_config_file(). Then, your custom parameters are merged into it using torchhydro.configs.config.update_cfg().
3. Train and Evaluate: Finally, you pass the consolidated configuration dictionary to the torchhydro.trainers.trainer.train_and_evaluate() function, which handles the entire pipeline: data loading, model building, training, and evaluation.

from torchhydro.configs.config import cmd, default_config_file, update_cfg
from torchhydro.trainers.trainer import train_and_evaluate

# 1. Define parameters for your experiment
args = cmd(...)

# 2. Load default config and update it with your parameters
config_data = default_config_file()
update_cfg(config_data, args)

# 3. Run the training and evaluation pipeline
train_and_evaluate(config_data)

Data configuration

Before running anything, create hydro_setting.yml in your home directory so the data resolver can locate datasets (see Installation):

storage:
  default_source: local
  local:
    root: 'D:/data/hydrodatasets'

In source_cfgs, select a dataset with the dataset_id key:

source_cfgs={"dataset_id": "camels_us"}

Paths are resolved centrally from hydro_setting.yml — there is no per-dataset path to configure. The legacy source_name + source_path form is still accepted for backward compatibility, but source_path is ignored in favor of central resolution.

Below are two practical examples demonstrating this workflow.

Example 1: Training a Standard LSTM on CAMELS Data

This example shows how to train a standard LSTM model for streamflow prediction using the CAMELS-US dataset.

Step 1: Define Parameters

First, we define all the necessary parameters for our experiment. This includes data source, model type, variables, time periods, and training settings.

import os
from hydrodataset.hydro_dataset import StandardVariable
from torchhydro.configs.config import cmd

args = cmd(
    # Experiment name and output directory
    sub=os.path.join("results", "lstm_camels"),

    # Data source — dataset_id auto-resolves the path from ~/hydro_setting.yml
    source_cfgs={"dataset_id": "camels_us"},

    # Use CPU for this example; to use GPU, set ctx=[0], [0, 1], ...
    ctx=[-1],

    # Model selection and hyperparameters
    model_name="CpuLSTM",
    model_hyperparam={
        "n_input_features": 23,
        "n_output_features": 1,
        "n_hidden_states": 256,
    },

    # Basin IDs for training and evaluation
    gage_id=[
        "01013500", "01022500", "01030500", "01031500", "01047000",
        "01052500", "01054200", "01055000", "01057000", "01170100",
    ],

    # Training settings
    batch_size=8,
    train_epoch=2,
    save_epoch=1,

    # Sequence lengths
    hindcast_length=0,
    forecast_length=20,

    # Time settings
    min_time_unit="D",
    min_time_interval="1",

    # Input and output variables
    var_t=[
        StandardVariable.PRECIPITATION, StandardVariable.DAYLIGHT_DURATION,
        StandardVariable.SOLAR_RADIATION, StandardVariable.TEMPERATURE_MAX,
        StandardVariable.TEMPERATURE_MIN, StandardVariable.VAPOR_PRESSURE,
    ],
    var_out=[StandardVariable.STREAMFLOW],

    # Data components
    dataset="StreamflowDataset",
    sampler="KuaiSampler",
    scaler="DapengScaler",

    # Model loading configuration for evaluation
    model_loader={"load_way": "specified", "test_epoch": 2},

    # Date ranges for training, validation, and testing
    train_period=["2000-10-01", "2001-10-01"],
    valid_period=["2001-10-01", "2002-10-01"],
    test_period=["2002-10-01", "2003-10-01"],

    # Loss function and optimizer
    loss_func="RMSESum",
    opt="Adam",
    lr_scheduler={0: 1, 1: 0.5, 2: 0.2},

    # Tensor layout
    which_first_tensor="sequence",
)

Step 2: Run the Pipeline

With the parameters defined, we simply call the update and train functions.

from torchhydro.configs.config import default_config_file, update_cfg
from torchhydro.trainers.trainer import train_and_evaluate

# Load default config and update it
config_data = default_config_file()
update_cfg(config_data, args)

# Run training and evaluation
train_and_evaluate(config_data)

The library will now handle data loading, preprocessing, model training, and finally, evaluation on the test set. Results will be saved in the directory specified by the sub parameter (e.g., results/lstm_camels).

Example 2: Training a Physics-Informed Model (DPL-XAJ)

This example demonstrates a more advanced use case: training a Differentiable Parameter Learning (DPL) model. Here, an LSTM is coupled with the Xinanjiang (XAJ) hydrological model. The network learns to output the parameters of the XAJ model.

Step 1: Define DPL Parameters

The parameter definition is similar, but we specify a different model (DplAttrXaj), dataset (DplDataset), and some additional hyperparameters specific to the physics-informed approach.

import os
from hydrodataset.hydro_dataset import StandardVariable
from torchhydro.configs.config import cmd

dpl_args = cmd(
    sub=os.path.join("results", "dpl_xaj_camels"),
    source_cfgs={"dataset_id": "camels_us"},
    ctx=[0],  # Use GPU 0

    # DPL model selection
    model_name="DplAttrXaj",
    model_hyperparam={
        "n_input_features": 17,
        "n_output_features": 15,  # Number of XAJ model parameters
        "n_hidden_states": 256,
        "kernel_size": 15,
        "warmup_length": 30,  # Warm-up period for the hydrological model
        "param_limit_func": "clamp",
    },

    # DPL models often require a specialized dataset
    dataset="DplDataset",
    constant_only=True,  # DplAttrXaj uses only static attributes as ANN input

    # Use a hybrid loss function for multiple outputs
    loss_func="MultiOutLoss",
    loss_param={
        "loss_funcs": "RMSESum",
        "data_gap": [0, 0],
        "device": [0],
        "item_weight": [1, 0],  # Weight for streamflow vs. other outputs
        "limit_part": [1],
    },

    # Special scaler settings for physical models
    scaler="DapengScaler",
    scaler_params={
        "prcp_norm_cols": ["streamflow"],
        "gamma_norm_cols": [
            StandardVariable.PRECIPITATION,
            StandardVariable.POTENTIAL_EVAPOTRANSPIRATION,
        ],
        "pbm_norm": True,
    },

    gage_id=["01013500", "01022500", "01030500", "01031500", "01047000"],
    train_period=["1985-10-01", "1986-04-01"],
    test_period=["2000-10-01", "2001-10-01"],
    valid_period=None,

    batch_size=50,
    forecast_length=60,
    warmup_length=30,

    # Input variables for the neural network part
    var_t=[
        StandardVariable.PRECIPITATION,
        StandardVariable.POTENTIAL_EVAPOTRANSPIRATION,
    ],
    # Output variable: only streamflow (CAMELS-US has no ET observations).
    # DplAttrXaj internally computes ET via the XAJ physics model, but
    # training loss is only computed against observed streamflow.
    var_out=[StandardVariable.STREAMFLOW],
    n_output=1,

    train_epoch=2,
    model_loader={"load_way": "specified", "test_epoch": 2},
    opt="Adadelta",
    which_first_tensor="sequence",
)

Step 2: Run the DPL Pipeline

The execution step is identical.

from torchhydro.configs.config import default_config_file, update_cfg
from torchhydro.trainers.trainer import train_and_evaluate

# Load default config and update it
config_data = default_config_file()
update_cfg(config_data, dpl_args)

# Run training and evaluation
train_and_evaluate(config_data)

Conventions and Additional Features

Time ranges are left-closed, right-open

Each period ["start", "end"] includes start and excludes end. Adjacent periods therefore do not overlap:

train_period=["2000-10-01", "2001-10-01"]
valid_period=["2001-10-01", "2002-10-01"]  # starts exactly where train ends

A strict overlap between train/valid/test is reported as an advisory warning (results are labelled in-sample); it never blocks execution. Set allow_split_overlap: True to suppress the warning when the in-sample intent is explicit.

Time units

min_time_unit accepts h/H (hourly), d/D (daily), and ME/MS/M (monthly). Use min_time_interval for sub-daily intervals (e.g. "3" with unit "h" for 3-hourly data).

Cloud-Zarr lazy loading

To train on large cloud Zarr stores without loading them fully into memory, use the CloudZarrLazyDataset dataset and CloudZarrChunkBatchSampler sampler. See the Cloud-Zarr Joint Training guide and experiments/train_cloud_zarr_lazy_joint.py for details.

Explore More

torchhydro supports a wide range of models, datasets, and data sources beyond the examples above.

Example Guides

• LSTM on CAMELS-US: Benchmark results on 671 US basins
• Songliao Flood Models: LSTM / LSTM+GNN / WDNE flood forecasting
• Sanxia Multi-Site: LSTM / MTSLSTM multi-site experiments
• MoE / Diffusion / FNO: Mixture-of-Experts, Diffusion, and Fourier Neural Operator models
• Seq2Seq / BALSTM / MTSLSTM: Encoder-decoder, attention-based LSTM, multi-timescale models

Advanced Topics

• Cloud-Zarr Joint Training: Lazy loading from cloud Zarr stores
• Lightning Fabric: Distributed training with Fabric
• ERA5-Land & GPM: ERA5-Land reanalysis and GPM satellite precipitation
• Model Interpretability: SHAP analysis and loss landscape visualization

API Reference

• Models API: All registered model architectures and loss functions
• Datasets API: Dataset classes, data sources, samplers, and scalers
• Trainers API: Training and evaluation pipeline
• Configs API: Configuration system and data resolver