Pywr-DRB Part 2: Running Ensemble Simulations with the DRB Reconstruction Ensemble

Motivation

We recently posted “Introducing Pywr-DRB - Part 1” where we highlighted the latest 2.0 version of our Pywr-DRB model.

Pywr-DRB is designed to simulate water resource systems operations in the Delaware River Basin (DRB) under different flow conditions. The current Pywr-DRB package includes 9 different streamflow datasets which can be used to run different simulations.

However, all of the pre-packaged datasets are single-scenario, deterministic datasets. We do not include any ensemble datasets in the Pywr-DRB installation because of memory constraints, however it is possible to run large, multi-realization ensembles of streamflows using Pywr-DRB.

This post is intended to:

1. Introduce the 74-year probabilistic DRB streamflow reconstruction (1945-2023) dataset available on Zenodo
2. Demonstrate how to prepare and run ensemble data using Pywr-DRB

This post should be considered an “Advanced” Pywr-DRB tutorial and it assumes that you are already familiar with the basic Pywr-DRB API functionality for single-scenario simulations.

If you would like to get up-to-speed on the Pywr-DRB basics, I’d recommend you begin by exploring the links below to understand the basics.

Links:

• Pywr-DRB Repository
  • Pywr-DRB v2.0.0 Release Notes
• Pywr-DRB Documentation
• Blog post repo: TrevorJA/Pywr-DRB-Reconstruction-Ensemble-Tutorial

DRB Streamflow Reconstruction Ensemble Dataset

The DRB streamflow reconstruction dataset was created because we wanted a dataset which:

• includes the 1960s drought of record
• accounts for uncertainty in historical streamflows via probabilistic predictions
• maximizes the use of historical data when predicting historic streamflows

In short, the reconstruction ensemble includes 1,000 realizations of daily streamflow at all Pywr-DRB model nodes from 1945-2023. These 1,000 realizations are each slightly different and reflect the uncertainty in (1) the streamflow timing and (2) streamflow magnitude for ungauged catchments.

The methods used to generate the probabilistic streamflow reconstruction are beyond the scope of this post, however you can review the details of the methodology in the paper cited below:

Amestoy, T. J., Hamilton, A. L., & Reed, P. M. (2026). Integrated river basin assessment framework combining probabilistic streamflow reconstruction, Bayesian bias correction, and drought storyline analysis. Environmental Modelling & Software, 195, 106756. https://doi.org/10.1016/j.envsoft.2025.106756

The dataset is currently available on Zenodo as a separate data release:

Amestoy, T., & Reed, P. (2025). Delaware River Basin Probabilistic Daily Streamflow Reconstruction Ensemble and Water Systems Model Data 1945-2023 (1.0.0) [Data set]. Zenodo. https://doi.org/10.5281/zenodo.15101164

We have included the median realization of the reconstruction ensemble in the latest Pywr-DRB v2.0.0 release. That dataset (named "pub_nhmv10_BC_withObsScaled") is a single timeseries reflecting the median streamflow conditions across the ensemble realizations and includes the full 1945-2023 period.

Here, however, we want to run simulations using the full 1,000 member reconstruction ensemble.

Ensemble Simulation in Pywr-DRB

I’ll start by assuming that you have already installed the pywrdrb package. If not, see Tutorial 01 - Installation Guide.

The full workflow, can be performed using the run_reconstruction_parallel.sh via SLURM.

However, this script assumes that you are running on an HPC system with SLURM (I am running on Hopper at Cornell) and may require modifications based on your computing resources.

I will describe the workflow in 5 steps, which correspond to the script 01_, … , 05_ in the blog post repo: TrevorJA/Pywr-DRB-Reconstruction-Ensemble-Tutorial:

1. Download the reconstruction ensemble from Zenodo
2. Calculate marginal catchment inflow
3. Predicting inflows at downstream nodes
4. Running the ensemble simulation in parallel
5. Post-processing results

In the descriptions that follow, I share only the snippets of code that I believe to be most helpful when understanding the key workflow components. If you are trying to run this yourself, you will need to visit the repository linked above, which has the full code.

But first, a note on the necessary input files that are expected when running a custom input dataset.

Necessary input files for custom dataset simulation

On the Pywr-DRB docs site, we have a more detailed Tutorial 03 ‘Using Customized Data to Run Pywr-DRB’ however that tutorial is focused on single-realization datasets.

As described in Pywr-DRB Tutorial 03, when running Pywr-DRB with a custom input dataset, you must originally provide full natural streamflow estimates at all Pywr-DRB node locations. This natural streamflow file must be called gage_flow_mgd.<filetype>.

For single-realization datasets, this file must be csv type.

For many-realization, ensemble datasets, this file must be an HDF5 (.hdf5) file.

Before we can run the Pywr-DRB simulation, we need to generate some additional input files which are derived from the gage_flow_mgd file. Ultimately, before running an ensemble simulation, we will need:

• gage_flow_mgd.hdf5
  • Contains full-natural streamflow timeseries for each node.
• catchment_inflow_mgd.hdf5
  • Contains marginal catchment inflow timeseries for each node.
• predicted_inflows_mgd.hdf5
  • Contains 1-4 day ahead predictions of inflows at multiple nodes, which are used to inform the NYC reservoir operations for downstream flow targets.

In Step 1, we will download the gage_flow_mgd.hdf5 file from Zenodo.

In Steps 2 and 3, we will use the gage_flow_mgd.hdf5 file to generate the catchment_inflow_mgd.hdf5 and predicted_inflows_mgd.hdf5 respectively.

Step 1: Download the reconstruction from Zenodo

First, we need to download the reconstruction ensemble dataset from Zenodo.

The 01_download_reconstruction.py script is setup to only download the drb_historic_streamflow_ensemble_data.zip from Zenodo , which is 10.69GB. This download took ~40 minutes on Hopper, but may be different depending on the internet speed.

python 01_download_reconstruction.py

After the download, we need to unzip the downloaded file:

# Unzip the file
unzip drb_historic_streamflow_ensemble_data.zip

The drb_historic_streamflow_ensemble_data.zip file contains two HDF5 files, each containing 1,000 realizations of daily full natural streamflow at 33 Pywr-DRB node locations across the basin:

• outputs/ensembles/gage_flow_obs_pub_nhmv10_BC_ObsScaled_ensemble.hdf5
• outputs/ensembles/gage_flow_obs_pub_nhmv10_ObsScaled_ensemble.hdf5

For the purposes of this demonstration, I am going to focus solely on the bias-corrected (BC) version of the streamflow reconstruction.

When setting-up Pywr-DRB simulations with custom datasets, it is recommended to organize the data such that you have a single folder with the dataset name, where we will store all and access all of the input data later on:

./<dataset_name>/*

In this case, I create a pywrdrb_inputs/<dataset_name> directory:

mkdir -p pywrdrb_inputs/obs_pub_nhmv10_BC_ObsScaled_ensemble

Then we need to move and rename the reconstruction streamflow file to the new directory:

mv outputs/ensembles/gage_flow_obs_pub_nhmv10_BC_ObsScaled_ensemble.hdf5 pywrdrb_inputs/obs_pub_nhmv10_BC_ObsScaled_ensemble/gage_flow_mgd.hdf5

Now we are ready to move forward.

At this point, we have our gage_flow_mgd.hdf5 file, and pywrdrb is going to help us prepare the remaining input data.

Step 2: Calculating marginal catchment inflows

Pywr-DRB requires “catchment inflows” rather than total streamflows as input data.

The script 02_calculate_catchment_inflow.py performs this conversion for each realization individually:

from pywrdrb.pre.flows import _subtract_upstream_catchment_inflows

inflow_ensemble = {}
for realization in gage_flow_ensemble:
    flows_i = gage_flow_ensemble[realization].copy()
    inflow_ensemble[realization] = _subtract_upstream_catchment_inflows(flows_i)

Where the _subtract_upstream_catchment_inflows function iteratively subtracts upstream gage flows from downstream totals, working through the node network to separate the local catchment contribution at each location. The result is a set of catchment inflows that can drive the Pywr-DRB simulation without double-counting upstream flows.

Step 3: Predicting inflows at downstream nodes

Pywr-DRB uses autoregressive models to generate the 1-4 day ahead inflow predictions.

The script 03_predict_inflows.py uses uses the pywrdrb.pre.PredictedInflowEnsemblePreprocessor() to generate the ensemble of predicted inflows in parallel using MPI.

It is possible to run this same code with use_mpi=False, however it may be slow depending on the number of realizations you are running.

from pywrdrb.pre import PredictedInflowEnsemblePreprocessor

# Generate predictions using AR models
preprocessor = PredictedInflowEnsemblePreprocessor(
    flow_type=inflow_type,
    ensemble_hdf5_file=catchment_inflow_filename,
    realization_ids=realization_ids,
    use_mpi=True   # Uses MPI to parallelize
)

preprocessor.load()
preprocessor.process()        
preprocessor.save()

After running this script, you will notice that a new file (predicted_inflows_mgd.hdf5) appears in the pywrdrb_inputs/ folder.

Step 4: Running the simulation

At this point, we have created all of the necessary input data and are ready to execute the simulation.

Given that we are focused on ensembles in this post, I will demonstrate how to run the ensemble simulations in parallel using MPI.

4.1 Distributing realizations across parallel simulations

Pywr-DRB supports two-levels of parallelization:

1. Inter-core parallelization using MPI
2. Intra-core parallelization using pywr

In the first case, a single core is setup to run a single simulation. This is the basic parallelization.

In the second case, the pywr core model runs multiple realizations during a single model run. In this case, pywr is keeping track of a matrix of system variables, where the matrix contains all realizations at once. This can be performed on a single core, however it may be memory constrained.

For example, assume that we have 100 realizations and only 10 core.

I can run all 100 realizations at once, by assigning:

• 1 simulation per core
• 10 realizations per simulation (referred to as a “batch” in the code below)

This is the approach used in 04_run_pywrdrb_simulations.py.

When distributing the realizations, I prefer MPI-level parallelization assuming I have more core than realizations. I.e., I want to max out my use of available core, while minimizing the number of realizations per simulation.

Given a list of realization_ids (e.g., [1, … ,500]), I split these realization ID numbers into “batches” of <= 10 realizations (to avoid memory bottlenecks) for each MPI rank using:

# Broadcast realization IDs
realization_ids = comm.bcast(realization_ids, root=0)

# Split realizations into batches across ranks
rank_realization_ids = np.array_split(realization_ids, size)[rank]
rank_realization_ids = list(rank_realization_ids)
n_rank_realizations = len(rank_realization_ids)

# Split rank realizations into batches
n_batches = math.ceil(n_rank_realizations / N_REALIZATIONS_PER_PYWRDRB_BATCH)
batched_indices = {}
for i in range(n_batches):
    batch_start = i * N_REALIZATIONS_PER_PYWRDRB_BATCH
    batch_end = min((i + 1) * N_REALIZATIONS_PER_PYWRDRB_BATCH, n_rank_realizations)
    batched_indices[i] = rank_realization_ids[batch_start:batch_end]

The result is that, for each individual MPI rank, I end up with a dictionary of:

{batch_number : realization_index_list}

# Example:
{
0 : [1,2,3],
1 : [4,5,6]
}

Then, when it comes time for that rank to prepare and run the pywrdrb simulation, it will focus on the specific realization_index_list for each batched simulation run.

Important: When using the pywrdrb.ModelBuilder() we must provide the specific list of realization IDs to focus on. When it sets up the simulation, it will only load those specific realizations from the input files.

Also important: Given that we will have multiple model instances being created at once, we will have multiple .json files being generated as well, where each JSON contains the settings and specifications for a specific simulation realization. With this in mind, we need to make sure each MPI rank is creating a unique JSON file for each simulation batch, as shown below.

for batch, indices in batched_indices.items():
    
    # Model options for this rank, batch
    model_options = {
        "inflow_ensemble_indices": indices,
    }
    
    # Build model
    mb = pywrdrb.ModelBuilder(
        inflow_type=f'{inflow_type}',
        start_date=start_date,
        end_date=end_date,
        options=model_options,
    )
    
    # Save model with unique filename
    model_fname = f"{output_dir}/{inflow_type}_rank{rank}_batch{batch}.json"
    mb.make_model()
    mb.write_model(model_fname)

    # rest of the workflow
    ...

Now, since I am running in parallel, I will have multiple output files generated.

So, in order to avoid overwriting each core’s output, I provide a unique output filename for each MPI rank and each simulation batch:

# unqiue name for output for this rank and batch
batch_output_filename = f"{output_dir}/{inflow_type}_rank{rank}_batch{batch}.hdf5"

recorder = pywrdrb.OutputRecorder(
    model=model,
    output_filename=batch_output_filename,
)

# Run simulation
model.run()

And let it run!

4.2 Combining batched output files

Lastly, it is nice to re-combine the simulation outputs from the distributed parallel simulations after they are all done.

This is done at the end of 04_run_pywrdrb_simulations.py, as shown in the snippet below.

WARNING(S):

• Combining multiple HDF5 files may overflow your machine memory if you have a large number of realizations.
  • This works fine on Hopper for 100 realizations
  • However, it may cause memory overflow if you try this with 1,000
• Also, the combine_batched_hdf5_outputs is not guaranteed to preserve the specific ordering of the realizations in the final output. In this case, it does not matter since we are not tracking individual realizations, but rather compared with the spread of the ensemble.

import glob
from pywrdrb.utils.hdf5 import combine_batched_hdf5_outputs

# Find all batch files for this set
batch_pattern = f"{output_dir}/{inflow_type}_rank*_batch*.hdf5"
all_batch_files = glob.glob(batch_pattern)

# Combine and save a single as output
combine_batched_hdf5_outputs(all_batch_files, output_filename)

In the 04_run_pywrdrb_simulations.py script I also identify and delete the individual rank, batch files after combining them.

Step 5: Post-processing data

By default, Pywr-DRB will save all output data within an hdf5 file, regardless of whether you are running a single realization or an ensemble.

We designed the pywrdrb.Data object to handle ensemble data by default - when loading an output file, it will automatically retrieve all of the ensemble realizations within the output file.

The 05_pot_example_results.py script shows how to load the simulation output.

import pywrdrb
from config import inflow_type, output_filename

### Load output data
results_sets = ['res_storage', 'major_flow']
data = pywrdrb.Data(print_status=True)

data.load_output(
    results_sets=results_sets,
    output_filenames=[output_filename])

After running data = pywrdrb.Data().load_output(), the full ensemble will be stored within the data object in a hierarchical structure according to:

data.<results_set>[<output_name>][<scenario_id>] -> pd.DataFrame

For example: data.res_storage[inflow_type][0] will return a DataFrame for the first realization of the res_storage results set, containing a column for each reservoir and the date index.

Conclusions

With the input files prepared and the MPI batching scheme above, Pywr-DRB can simulate the full 1,000-member reconstruction ensemble in a single HPC job, and the pywrdrb.Data object makes the resulting ensemble output straightforward to load and analyze. The complete, runnable workflow is in the Pywr-DRB-Reconstruction-Ensemble-Tutorial repository, and the original version of this post is on the WaterProgramming blog.