This repository presents a cleaned, refactored, documented, and reorganized version of the research code originally released on Zenodo, with assistance from OpenAI's ChatGPT (GPT-5.6 Sol, High reasoning mode).
The original code was developed for the experiments reported in Wang & Gupta (2024), while this repository restructures those implementations into a more usable, reproducible, and maintainable model zoo.
It contains single-node Mass-Conserving Perceptron (MCP) model variants, together with cleaned training and evaluation scripts, historical checkpoints, standardized model notation, and documented execution conventions for model reproduction, continuation, fine-tuning, and further development.
Readers are encouraged to consult the original paper for the formal model names, notation, and mathematical definitions associated with each MCP variant.
- Full article: Wang & Gupta (2024)
- DOI: 10.1029/2023WR036461
Wang, Y.-H. and Gupta, H.V. (2024). A mass-conserving-perceptron for machine-learning-based modeling of geoscientific systems. Water Resources Research, 60(4), e2023WR036461.
The figure below illustrates the functional components and modeling utilities that can be represented within a single-node MCP, including:
- input bias correction
- input bypass
- gate functional complexity
- output-gate context dependence
- unobserved mass exchange with the environment
- mass conservation.
Conceptual illustration of the functional components and modeling utilities represented in a single-node MCP.
The model notation used throughout the repository is:
| Symbol | Meaning |
|---|---|
| Output gate | |
| Loss gate | |
| Remember gate | |
| Constant gate | |
| Sigmoid-based time-variable gate | |
| PET-constrained loss gate | |
| Piecewise-linear precipitation bias-correction gate with dimension |
|
| Piecewise-quadratic precipitation bias-correction gate with dimension |
|
| Single hidden ANN layer with |
|
| Multi-information (MI): the gate uses additional state information | |
| MR | Mass-relaxation mechanism |
The basic family contains constant (
For variable gates, the output gate depends on the current storage state and the loss gate depends on current PET.
| Model | Notation | Description | MCP_Zoo class |
|---|---|---|---|
| M1 | Constant output + constant loss | MCPBRNN_constant_OutLoss |
|
| M2 | Constant output + PET-variable loss | MCPBRNN_Generic_constant_Out_variableLoss |
|
| M3 | Storage-variable output + constant loss | MCPBRNN_Generic_variable_Out_constantLoss |
|
| M4 | Storage-variable output + PET-variable loss | MCPBRNN_Generic_Scaling |
|
| M5 | M4 + PET-constrained loss | MCPBRNN_Generic_PETconstraint_Scaling |
Historical checkpoint references:
| Model | Historical checkpoint |
|---|---|
| M1 | model_epoch42.pt |
| M2 | model_epoch468.pt |
| M3 | model_epoch4963.pt |
| M4 | model_epoch48.pt |
| M5 | model_epoch27.pt |
These models extend M5 with a precipitation-dependent bias-correction function.
For both model families below, M1–M5 indicate the dimensionality of the bias-correction function, with dimensions ranging from 1 to 5. They do not refer to the five base MCP models described in Section 2.
The model family is denoted as
where
Accordingly:
-
M1-IBCorrPLcorresponds to$B_{L1}$ -
M2-IBCorrPLcorresponds to$B_{L2}$ -
M3-IBCorrPLcorresponds to$B_{L3}$ -
M4-IBCorrPLcorresponds to$B_{L4}$ -
M5-IBCorrPLcorresponds to$B_{L5}$
The model family is denoted as
where
Accordingly:
-
M1-IBCorrPQcorresponds to$B_{Q1}$ -
M2-IBCorrPQcorresponds to$B_{Q2}$ -
M3-IBCorrPQcorresponds to$B_{Q3}$ -
M4-IBCorrPQcorresponds to$B_{Q4}$ -
M5-IBCorrPQcorresponds to$B_{Q5}$
Folders:
M-IBCorrPL/
├── M1-IBCorrPL/
├── M2-IBCorrPL/
├── M3-IBCorrPL/
├── M4-IBCorrPL/
└── M5-IBCorrPL/
M-IBCorrPQ/
├── M1-IBCorrPQ/
├── M2-IBCorrPQ/
├── M3-IBCorrPQ/
├── M4-IBCorrPQ/
└── M5-IBCorrPQ/
These models extend M5 by incorporating a mass-relaxation (MR) mechanism.
The four MR variants are:
| Model | Notation | Description |
|---|---|---|
| MR-1 | Regular | |
| MR-2 | Independent | |
| MR-3 | Regular-Relaxed | |
| MR-4 | Independent-Relaxed |
Folders:
MR-1/
MR-2/
MR-3/
MR-4/
The model family is denoted as
Folder structure:
M-ComplexGate/
└── LossGateOnly/
├── loss-dim-1/
├── loss-dim-2/
├── loss-dim-3/
├── loss-dim-4/
└── loss-dim-5/
Class:
MCPBRNN_Generic_LossANNGate_PETconstraintThe output gate is storage-dependent sigmoid. The loss gate is a PET-driven ANN.
The model family is denoted as
Folder structure:
M-ComplexGate/
└── OutputGateOnly/
├── output-dim-1/
├── output-dim-2/
├── output-dim-3/
├── output-dim-4/
└── output-dim-5/
Class:
MCPBRNN_Generic_OutputANNGate_PETconstraintThe output gate is a storage-driven ANN. The loss gate remains PET-dependent sigmoid.
The model family is denoted as
There are 25 combinations:
loss-dim-1-out-dim-1
...
loss-dim-5-out-dim-5
Here,
Class:
MCPBRNN_Generic_ANNGate_PETconstraint_GenericThe output and loss gates each use a one-hidden-layer ANN.
MI means that a gate uses additional information beyond its standard current-time input.
For the current implementation:
- the output gate can use current storage plus lagged storage (
$t-1$ ); - the loss gate can additionally use current storage information;
- the superscript
$^{+}$ denotes the MI extension.
The model family is denoted as
Class:
MCPBRNN_Generic_PETconstraint_MIlossThe model family is denoted as
Class:
MCPBRNN_Generic_PETconstraint_MIoutputThe model family is denoted as
There are 25 combinations.
Here,
Class:
MCPBRNN_Generic_PETconstraint_MIoutputlossThese cases use direct sigmoid gates instead of ANN layers.
| Folder | Notation | Class |
|---|---|---|
loss-gate-only/ |
MCPBRNN_Generic_PETconstraint_MIloss_Sigmoid |
|
output-gate-only/ |
MCPBRNN_Generic_PETconstraint_MIoutput_Sigmoid |
|
loss-output-gate/ |
MCPBRNN_Generic_PETconstraint_MIoutputloss_Sigmoid |
project/
├── 20220527-MDUPLEX-LeafRiver/
├── MCPBRNN_lib_tools/
│ ├── MCP_Zoo.py
│ ├── Eval_Metric.py
│ └── Loss_Function.py
│
├── M1/
├── M2/
├── M3/
├── M4/
├── M5/
│
├── M-IBCorrPL/
│ ├── M1-IBCorrPL/
│ └── ...
│
├── M-IBCorrPQ/
│ ├── M1-IBCorrPQ/
│ └── ...
│
├── M-MassRelaxation/
│ ├── MR1/
│ ├── MR2/
│ ├── MR3/
│ └── MR4/
│
├── M-ComplexGate/
│ ├── LossGateOnly/
│ ├── OutputGateOnly/
│ └── BothLossOutputGate/
│
└── M-MI-ComplexGate/
├── LossGateOnly/
├── OutputGateOnly/
├── BothLossOutputGate/
└── MI-Only/
├── loss-gate-only/
├── output-gate-only/
└── loss-output-gate/
The 20220527-MDUPLEX-LeafRiver folder contains the daily hydrometeorological data for the Leaf River catchment in Mississippi, USA.
The current training and evaluation scripts use two files:
-
LeafRiverDaily_43YR.txt— 43-year daily data containing three columns: precipitation ($P$ ), potential evapotranspiration ($PET$ ), and streamflow ($Q$ ). -
LeafRiverDaily_43YR_Flag.txt— data-partition flags used to identify spinup, training, selection, and testing periods.
The partition flags are:
| Flag | Period |
|---|---|
-99999 |
Spinup |
-1 |
Training |
0 |
Selection |
1 |
Testing |
Other files in this directory are retained from earlier experiments and alternative data-partition schemes, but are not used by the current cleaned Model Zoo scripts.
All model classes should be consolidated in:
MCPBRNN_lib_tools/MCP_Zoo.py
All model scripts can be executed from the repository root.
The scripts use argparse for optional runtime settings. Fixed-structure models can be run directly, whereas model families with multiple functional dimensions require the corresponding dimension argument to select the intended model case.
The included checkpoint is used by default unless --checkpoint is explicitly provided.
The five base models have fixed structures and can be run directly.
Training:
python M1/mcpbrnn_Main_constantO_constantL_clean.pyEvaluation:
python M1/mcpbrnn_Main_constantO_constantL_EVAL_clean.pyTraining:
python M2/mcpbrnn_Main_constantO_variableL_clean.pyEvaluation:
python M2/mcpbrnn_Main_constantO_variantL_EVAL_clean.pyTraining:
python M3/mcpbrnn_Main_variableO_constantL_clean.pyEvaluation:
python M3/mcpbrnn_Main_variantO_constantL_EVAL_clean.pyTraining:
python M4/mcpbrnn_Main_M4_Scaling_clean.pyEvaluation:
python M4/mcpbrnn_Main_M4_Scaling_EVAL_clean.pyTraining:
python M5/mcpbrnn_Main_PETconstraint_clean.pyEvaluation:
python M5/evaluate_Main_PETconstraint_clean.pyThe precipitation bias-correction families share one training script and one evaluation script across dimensions 1–5.
The desired bias-correction dimension is specified using:
--dim x
where x = 1,...,5.
For example, to run dimension 3:
python M-IBCorrPL/mcpbrnn_Main_PETconstraint_IBcorrPL_clean.py --dim 3Evaluation:
python M-IBCorrPL/mcpbrnn_Main_PETconstraint_IBcorrPL_EVAL_clean.py --dim 3The script automatically associates --dim 3 with:
M-IBCorrPL/M3-IBCorrPL/
For example, to run dimension 3:
python M-IBCorrPQ/mcpbrnn_Main_PETconstraint_IBcorrPQ_clean.py --dim 3Evaluation:
python M-IBCorrPQ/mcpbrnn_Main_PETconstraint_IBcorrPQ_EVAL_clean.py --dim 3The script automatically associates --dim 3 with:
M-IBCorrPQ/M3-IBCorrPQ/
The four mass-relaxation cases have fixed structures and can be run directly.
python MR-1/mcpbrnn_Main_MR1_Regular_clean.py
python MR-1/mcpbrnn_Main_MR1_Regular_EVAL_clean.pypython MR-2/mcpbrnn_Main_MR2_Independent_clean.py
python MR-2/mcpbrnn_Main_MR2_Independent_EVAL_clean.pypython MR-3/mcpbrnn_Main_MR3_Regular_Relaxed_clean.py
python MR-3/mcpbrnn_Main_MR3_Regular_Relaxed_EVAL_clean.pypython MR-4/mcpbrnn_Main_MR4_Independent_Relaxed_clean.py
python MR-4/mcpbrnn_Main_MR4_Independent_Relaxed_EVAL_clean.pySpecify the loss-gate ANN dimension using --dim.
For example:
python M-ComplexGate/LossGateOnly/mcpbrnn_Main_lossANNGate_clean.py --dim 3
python M-ComplexGate/LossGateOnly/mcpbrnn_Main_PETconstraint_lossANNGate_EVAL_clean.py --dim 3This selects:
M-ComplexGate/LossGateOnly/loss-dim-3/
Specify the output-gate ANN dimension using --dim.
For example:
python M-ComplexGate/OutputGateOnly/mcpbrnn_Main_outputANNGate_clean.py --dim 3
python M-ComplexGate/OutputGateOnly/mcpbrnn_Main_PETconstraint_outputANNGate_EVAL_clean.py --dim 3This selects:
M-ComplexGate/OutputGateOnly/out-dim-3/
The loss- and output-gate ANN dimensions are specified independently using:
--loss_dim y --out_dim x
For example, for a loss-gate dimension of 3 and an output-gate dimension of 5:
python M-ComplexGate/BothLossOutputGate/mcpbrnn_Main_ANNGate_clean.py --loss_dim 3 --out_dim 5Evaluation:
python M-ComplexGate/BothLossOutputGate/mcpbrnn_Main_ANNGate_PETconstraint_EVAL_clean.py --loss_dim 3 --out_dim 5This selects:
M-ComplexGate/BothLossOutputGate/loss-dim-3-out-dim-5/
Specify the loss-gate ANN dimension using --dim.
python M-MI-ComplexGate/LossGateOnly/mcpbrnn_Main_MI_Loss_PETconstraint_clean.py --dim 3
python M-MI-ComplexGate/LossGateOnly/mcpbrnn_Main_MI_Loss_PETconstraint_EVAL_clean.py --dim 3Specify the output-gate ANN dimension using --dim.
python M-MI-ComplexGate/OutputGateOnly/mcpbrnn_Main_MI_Output_PETconstraint_clean.py --dim 3
python M-MI-ComplexGate/OutputGateOnly/mcpbrnn_Main_MI_Output_PETconstraint_EVAL_clean.py --dim 3Specify both ANN dimensions using --loss_dim and --out_dim.
For example:
python M-MI-ComplexGate/BothLossOutputGate/mcpbrnn_Main_MI_LossOutput_PETconstraint_clean.py --loss_dim 3 --out_dim 5Evaluation:
python M-MI-ComplexGate/BothLossOutputGate/mcpbrnn_Main_MI_LossOutput_PETconstraint_EVAL_clean.py --loss_dim 3 --out_dim 5The MI-only sigmoid models have fixed structures and can be run directly.
python M-MI-ComplexGate/MI-Only/loss-gate-only/mcpbrnn_Main_MI_Loss_PETconstraint_Sigmoid_clean.py
python M-MI-ComplexGate/MI-Only/loss-gate-only/mcpbrnn_Main_MI_Loss_PETconstraint_Sigmoid_EVAL_clean.pypython M-MI-ComplexGate/MI-Only/output-gate-only/mcpbrnn_Main_MI_Output_PETconstraint_Sigmoid_clean.py
python M-MI-ComplexGate/MI-Only/output-gate-only/mcpbrnn_Main_MI_Output_PETconstraint_Sigmoid_EVAL_clean.pypython M-MI-ComplexGate/MI-Only/loss-output-gate/mcpbrnn_Main_MI_LossOutput_PETconstraint_Sigmoid_clean.py
python M-MI-ComplexGate/MI-Only/loss-output-gate/mcpbrnn_Main_MI_LossOutput_PETconstraint_Sigmoid_EVAL_clean.pyTraining scripts support optional arguments such as:
--case_no
--epoch_no
--time_lag
--seed_no
--checkpoint
--data_dir
--output_dir
--device
Dimension-dependent model families additionally use:
--dim
or:
--loss_dim
--out_dim
The available arguments and defaults for any script can be inspected using:
python <script>.py --helpFor cleaned scripts, the preferred convention is:
<model-folder>/
├── <one checkpoint>.pt
├── training script
└── EVAL script
If --checkpoint is not supplied, the script searches the relevant model folder for exactly one .pt file.
Training workflow:
existing checkpoint
↓
strict load
↓
continue training
↓
evaluate every epoch
↓
select best epoch by selection KGE
↓
best_model_epochX.pt
EVAL workflow:
checkpoint
↓
strict load
↓
forward-only evaluation
↓
metrics + diagnostic time series
Unless a case explicitly differs:
time_lag = 0
learning_rate = 0.025
learning_rates = {300: 0.0125, 600: 0.0125}Data split flags:
| Flag | Split |
|---|---|
-99999 |
spinup |
-1 |
training |
0 |
selection |
1 |
testing |
Reported metrics:
NSE
KGE
KGE-A
KGE-B
Corr
MSE / mse
KGEss
Training diagnostics also include KGE at lags 1, 2, and 3.
Cleaned output files use:
time, <result variable>, phase
with phases:
spinup1
spinup2
spinup3
simulation
The time axis contains:
- three repetitions of WY1949 (
1948-10-01to1949-09-30) for spinup; - formal simulation from
1948-10-01through1988-09-30; - total length: 15,705 rows.
Repeated calendar dates across spinup phases are intentional; phase distinguishes them.
Cleaned model classes follow these rules:
- preserve historical checkpoint parameter names and tensor shapes;
- preserve pre-update versus post-update hydrologic timing;
- use remember gate terminology for
gR; - use
x.new_zeros(...)for device-safe output tensors; - keep the
seq_len = 1assumption explicit; - remove unused imports, dead prototype code, and obsolete initialization logic;
- keep equations compact and readable;
- do not rename
state_dictparameter keys; - Python class names may be renamed only when necessary to avoid duplicate class names in
MCP_Zoo.py.
For MR3 and MR4, only the Python class names were disambiguated; the checkpoint parameter keys remain unchanged.
The repository was organized for the following Python environment.
numpy==2.2.2
pandas==2.2.3
torch==2.9.0
torchvision==0.24.0
tqdm==4.67.1
scikit-learn==1.6.1
Install with:
pip install -r requirement.txtname: MCP-SingleNode-ModelZoo
channels:
- defaults
dependencies:
- python=3.10
- pip
- notebook
- pip:
- numpy==2.2.2
- pandas==2.2.3
- torch==2.9.0
- torchvision==0.24.0
- tqdm==4.67.1
- scikit-learn==1.6.1Create and activate the Conda environment with:
conda env create -f environment.yml
conda activate MCP-SingleNode-ModelZooThis repository provides a consolidated and reproducible collection of single-node MCP model variants, including cleaned training and evaluation scripts, checkpoints, and data conventions.
For questions or suggestions, please contact Yuan-Heng Wang, Ph.D. at yhwang0730@gmail.com or yhwang0730@arizona.edu.
