OpenHydroNet: Riverine Flood Forecasting

🌊 This repository implements the state-of-the-art models that power Google FloodHub.

This is not an officially supported Google product. This project is not eligible for the Google Open Source Software Vulnerability Rewards Program.

The repository provides open-source replication of Google’s global flood-forecasting models. By open-sourcing these models, we aim to foster transparency, enable in-house integration in production systems, and accelerate academic research.

This repository is a fork of NeuralHydrology, which has been heavily modified and extended to support forecast sequences using the specific model architectures that are used operationally in the Google FloodHub.

📖 Documentation

Detailed instructions on how to configure, train, and evaluate OpenHydroNet models can be found on our documentation page: 👉 openhydronet.readthedocs.io

Watch our high-level video introduction to the interactive tutorial on YouTube: OpenHydroNet Tutorial Video

Models

This repository contains implementations of the core models used in Google's production forecasting systems.

Mean-Embedding-Forecast-LSTM

The Mean Embedding Forecast LSTM is a forecasting model that uses separate embedding networks for hindcast and forecast inputs. It aggregates these inputs using masked means before passing them into respective LSTMs for the hindcast and forecast periods.

• Status: Current production model (as of December 2025) for Google FloodHub.
• Reference: Gauch, Martin, et al. "How to deal with missing input data." Hydrology and Earth System Sciences (2025).

Handoff-Forecast-LSTM

The State Handoff Forecast LSTM is a forecasting model that uses a state-handoff to transition from a hindcast sequence (LSTM) model to a forecast sequence (LSTM) model. The hindcast model runs from the past up to the present (the issue time of the forecast) and then passes the cell state and hidden state of the LSTM into a (nonlinear) handoff network, which is used to initialize a new LSTM that rolls out over the forecast period.

• Status: Former production model for Google FloodHub.
• Reference: Nearing, Grey, et al. "Global prediction of extreme floods in ungauged watersheds." Nature (2024).

Data Assimilation

Data assimilation (DA) is an optional step during evaluate / infer that corrects a trained model with recent streamflow observations. For every forecast issue date, selected latent components of the model (the static and dynamic embeddings) are treated as free variables and adjusted by gradient descent so that the model's output matches the observations in a window of recent time steps ending at the issue date. A background (regularization) term keeps the components close to their unassimilated values. The model weights are never changed, and no new model is created.

• Supported model: mean_embedding_forecast_lstm only (components static_embedding, hindcast_embedding, forecast_embedding). The Handoff-Forecast-LSTM does not support DA.
• Configuration: the assimilation_config block, validated by AssimilationConfig; see Configuration below.

Installation

We recommend using Conda to manage dependencies like PyTorch and CUDA.

1. Create and Activate the Environment:

   # Create the environment from the file in the repo  
   conda env create -f environments/conda.yml
   
   # Activate the environment (MANDATORY)  
   conda activate googlehydrology  
   

2. Install the Package:
   Install in editable mode so that changes to the source code are reflected immediately:

   # Run from the root of the repository  
   pip install -e .
   

🚀 Tutorial Notebook

The most direct way to explore this repository is through our interactive tutorial: OpenHydroNet Tutorial Notebook.

What you will learn:

• Model Evaluation: Load pre-trained Google Hydrology models and calculate performance metrics (NSE, KGE) on real-world basin data.
• Fine-Tuning for Performance: Learn how to fine-tune the static_embedding_fc layer. This is a powerful technique for improving predictions on "outlier" basins (e.g., basins with unusual sizes or geology) without retraining the entire model.
• Visualizing Results: Compare model hydrographs against observed discharge data.

Run it now:

Data Setup

OpenHydroNet standardizes entirely on the high-performance, cloud-native Zarr format for streamflow targets, static attributes, meteorological forcings, and normalizer scalers.

1. Caravan Data (Zarr Format)

Caravan data is structured into modular Zarr stores:

• attributes.zarr: Catchment static attributes (area, elevation, soil, geology).
• streamflow.zarr: Gauge streamflow observations.

If you have downloaded the legacy Caravan NetCDF/CSV dataset from Zenodo, convert it to Zarr in a single step using the built-in CLI:

run convert-caravan --caravan-dir ~/data/Caravan-nc --output-dir ~/data/Caravan-zarr

2. MultiMet Dynamics Data

The MultiMet meteorological forcing data extension is accessed directly from Google Cloud Storage or local disk. Point your configuration to: gs://caravan-multimet/v1.1 (or your local dynamics directory).

3. Catchment Delineation (Creating Polygons for New Gauges)

If you have latitude and longitude coordinates for streamflow gauges and need their upstream watershed boundary polygons and drainage areas ($\text{km}^2$), use the delineate-catchment command-line tool included in this repository. It traces upstream drainage areas across 90-meter flow-direction map tiles and writes polygons in Caravan-compatible GeoParquet, GeoJSON, or Shapefile format for downstream MultiMet and static attribute extraction.

• Package Guide & CLI Reference: catchment_delineation/README.md
• Official Documentation: docs/source/usage/catchment_delineation.rst
• MultiMet Forcing Data: See MultiMet Data above (gs://caravan-multimet/v1.1).

Usage

The package installs the run command as the primary entry point.

Training a Model

run train --config-file /path/to/your/training_config_file.yml

Evaluation

Calculate performance metrics (NSE, KGE) on the test set:

run evaluate --run-dir /path/to/your/model_run/

Inference

Generate predictions (without skipping NaN observations):

run infer --run-dir /path/to/your/model_run/

Evaluation / inference with data assimilation

Run evaluate or infer with the --assimilate flag. The run's config must contain an assimilation_config block (see Configuration); the flag overrides assimilate: false in the config:

run evaluate --run-dir /path/to/your/model_run/ --assimilate
run infer --run-dir /path/to/your/model_run/ --assimilate

If the run's config.yml has no assimilation_config block, add the block to a copy of the config and pass it together with --run-dir. For evaluate / infer, --config-file replaces the run's config.yml entirely (it is not merged), so the copy must keep all other settings:

run infer --run-dir /path/to/your/model_run/ --config-file /path/to/config_with_da.yml --assimilate

DA outputs never overwrite the regular results: all output file stems get the suffix _data_assimilation, i.e. test_metrics_data_assimilation.csv and test_results_data_assimilation.zarr (the Zarr store is always written when DA is on, also in evaluate mode). Figures and hot-start state files are suffixed in the same way.

Configuration

Experiments are defined by YAML files. Update the following paths in your config (e.g., tutorial/configs/train-config.yml):

• run_dir: Where weights and logs are saved.
• train_basin_file: Path to the list of basin IDs.
• data_dir: Path to your root directory containing attributes.zarr, streamflow.zarr, and dynamic meteorological data (e.g., ~/data/Caravan-zarr). Used as the default for the three keys below.
• statics_data_dir / targets_data_dir: Optional paths to individual component Zarr stores (statics_data_path / targets_data_path are accepted as aliases).
• dynamics_data_dir: Path to forcing data (e.g., gs://caravan-multimet/v1.1 or local directory; alias dynamics_data_path).

To enable data assimilation, add an assimilation_config block (only mean_embedding_forecast_lstm). A minimal block:

assimilate: false  # or true; the CLI flag --assimilate also turns it on
assimilation_config:
  assimilation_components: [hindcast_embedding, static_embedding]
  assimilation_window: 30        # observed steps before the issue date
  initial_learning_rate: 0.01    # mandatory, not inherited from training
  regularization_weight: 0.5     # background term weight, default 0.0
  epochs: 50                     # default 100
  loss: MSE                      # MSE, NSE or CMAL; default MSE

The window must satisfy 1 <= assimilation_window <= seq_length - lead_time. The training keys epochs, initial_learning_rate, optimizer, loss and clip_gradient_norm are not inherited from the run config; the DA block sets them or the DA defaults apply. All keys are documented in googlehydrology/utils/assimilationconfig.py and in the configuration docs.

Example Configurations

The ~/flood-forecasting/example-configs directory contains reference YAML files that define the experimental setups for different model architectures and datasets.

• floodhub-settings-config.yml
  • Model Architecture: mean_embedding_forecast_lstm
  • Dataset: MultiMet (Global Caravan dataset)
  • Description: This configuration is designed to replicate the training settings of the current (2025) operational FloodHub model as closely as possible within this open-source framework.
• handoff-forecast-lstm-config.yml
  • Model Architecture: handoff_forecast_lstm
  • Dataset: MultiMet (Global Caravan dataset)
  • Description: Provides the settings used for the former operational model. This configuration aligns with the methodology described in the Nature (2024) paper for global ungauged flood prediction.
• camels-multimet-mean-embedding-forecast-lstm-config.yml
  • Model Architecture: mean_embedding_forecast_lstm
  • Dataset: CAMELS-US (531 basins)
  • Description: A benchmarking configuration for the Mean-Embedding model tailored for the CAMELS-US dataset. It is optimized for evaluating model stability and performance on a standard hydrological benchmark. Our team uses this as a reference point during model development, and it is included in this repository because this is what we use to ensure that any changes to the repository work as expected.
• camels-multimet-handoff-forecast-lstm-config.yml
  • Model Architecture: handoff_forecast_lstm
  • Dataset: CAMELS-US (531 basins)
  • Description: A benchmarking configuration for the State Handoff model tailored for the CAMELS-US dataset, used to compare the handoff approach against other architectures on US-based basin data.
• camels-multimet-mean-embedding-forecast-lstm-assimilation-config.yml
  • Model Architecture: mean_embedding_forecast_lstm
  • Dataset: CAMELS-US (531 basins)
  • Description: The CAMELS-US Mean-Embedding benchmark configuration extended with a fully commented assimilation_config block (all three embeddings, 90-step window, NSE loss). DA stays off (assimilate: false) so that a plain run evaluate reproduces the regular forecasts; run run evaluate --run-dir ... --assimilate to assimilate.

Extracting Static Attributes for Your Own Watersheds

To run OpenHydroNet models on a watershed, the model needs a table of static watershed characteristics (such as area, elevation, slope, soil type, land cover, and long-term average climate). For basins in the published Caravan dataset, these tables are already included.

If you want to run models on your own watersheds, this repository includes a static data workflow (multimet/static_extractor) that takes a map file of your watershed boundaries (.geojson, .shp, or .gpkg) and builds a Caravan-compatible CSV table of static attributes using the community HydroATLAS and ERA5-Land datasets.

extract-caravan-static \
    --input /path/to/watershed_polygons.geojson \
    --output /path/to/extracted_caravan_attributes.csv \
    --gdb-path /path/to/BasinATLAS_v10.gdb \
    --era5-source hybas \
    --era5-cache-dir /path/to/era5_climate

👉 Full Usage Guide & Command-Line Flags: See multimet/README.md.

Issue Reporting

If you encounter bugs, please use the GitHub Issue Tracker. Provide a clear description, steps to reproduce, and the expected behavior.