Pretrain a neural network

In this tutorial, we demonstrate the steps to pretrain the Neural Network (NN) component of the hybrid model. The training data is derived from the pre-processed outputs generated after running the conceptual model. The NN is trained to predict key quantities - \(ET\), \(M\), \(Pr\), \(Ps\), and \(Q\) - by matching the outputs of the conceptual model.

IMPORTANT: In this tutorial, we use a file listing 4 basins and demonstrate running a multi-basin model.

Before we start

• This tutorial is rendered from a Jupyter notebook that is hosted on GitHub. If you’d like to run the code yourself, you can access the notebook and configuration files directly from the repository: 02-PretrainNNmodel.

• To run this notebook locally, ensure you have completed the setup steps outlined in Getting started. These steps include setting up the environment, installing the required packages, and preparing the data files necessary for the tutorial.

• Dependency on a Previous Tutorial: Before running this tutorial, you must complete the 01-RunConceptModel Tutorial. After completing it:

  1- Move the generated run folder to src/data.

  2- Update the data_dir field in the config_run_pretrain.yml (here) file to point to this folder.

Import packages

%reload_ext autoreload
%autoreload 2

import sys
from pathlib import Path

# Dynamically set the project directory based on the notebook's location
notebook_dir = Path().resolve()
project_dir = str(notebook_dir.parent.parent)  # Adjust based on your project structure
sys.path.append(project_dir)

from src.thn_run import (
    _load_cfg_and_ds,
    get_basin_interpolators
)

from src.modelzoo_concept import get_concept_model
from src.modelzoo_nn import (
    get_nn_model,
    get_nn_pretrainer,
)

Constants

config_file = 'config_run_pretrain.yml'

Load config file and prepare dataset

cfg, dataset = _load_cfg_and_ds(Path(config_file), model='pretrainer')

-- Loading the config file and the dataset
-- Using device: cpu --
Setting seed for reproducibility: 1111
-- Loading basin dynamics into xarray data set.
100%|██████████| 4/4 [00:00<00:00, 21.40it/s]

cfg._cfg['experiment_name']

'run_pretrain_nn_model'

A folder has been created in the runs directory with the name specified as experiment_name in the configuration file, appended with a YYMMDD_HHMMSS timestamp. This folder will contain the configuration, results, plots, and metrics associated with the run.

# Dataset attributes
dataset.__dict__.keys()

dict_keys(['cfg', 'is_train', '_compute_scaler', 'scaler', 'basins', '_disable_pbar', '_per_basin_target_stds', '_dates', 'start_and_end_dates', 'num_samples', 'period_starts', 'alias_map', 'alias_map_clean', 'ds_train', 'ds_valid', 'ds_static'])

display(
    'basins', dataset.basins,
    'start_and_end_dates', dataset.start_and_end_dates,
    'ds_train', dataset.ds_train,
    'ds_valid', dataset.ds_valid,
)

'basins'

['01013500', '01022500', '01030500', '06431500']

'start_and_end_dates'

{'train': {'start_date': Timestamp('1980-10-01 00:00:00'),
  'end_date': Timestamp('2000-09-30 00:00:00')},
 'valid': {'start_date': Timestamp('2000-10-01 00:00:00'),
  'end_date': Timestamp('2010-09-30 00:00:00')}}

'ds_train'

<xarray.Dataset> Size: 2MB
Dimensions:     (basin: 4, date: 7305)
Coordinates:
  * date        (date) datetime64[ns] 58kB 1980-10-01 1980-10-02 ... 2000-09-30
  * basin       (basin) <U8 128B '01013500' '01022500' '01030500' '06431500'
Data variables: (12/15)
    dayl        (basin, date) float32 117kB 11.33 11.28 11.23 ... 11.52 11.42
    et_bucket   (basin, date) float32 117kB 0.8673 1.153 1.251 ... 1.213 1.155
    m_bucket    (basin, date) float32 117kB 1.192e-07 1.192e-07 ... 0.1266
    obs_runoff  (basin, date) float32 117kB 0.551 0.5607 0.5586 ... 0.441 0.4353
    pr_bucket   (basin, date) float32 117kB 3.1 4.24 ... 1.192e-07 1.192e-07
    prcp        (basin, date) float32 117kB 3.1 4.24 8.02 15.27 ... 0.0 0.0 0.0
    ...          ...
    s_water     (basin, date) float32 117kB 1.303e+03 1.305e+03 ... 984.9 983.5
    srad        (basin, date) float32 117kB 192.6 206.3 165.4 ... 350.7 303.8
    tmax        (basin, date) float32 117kB 10.05 15.82 15.86 ... 22.11 19.83
    tmean       (basin, date) float32 117kB 6.08 10.53 11.84 ... 13.41 12.77
    tmin        (basin, date) float32 117kB 2.11 5.24 7.81 ... 1.95 4.71 5.72
    vp          (basin, date) float32 117kB 711.3 898.6 ... 755.6 868.9

xarray.Dataset

• Dimensions:
  • basin: 4
  • date: 7305
• Coordinates: (2)
  • date
    (date)
    datetime64[ns]
    1980-10-01 ... 2000-09-30

    array(['1980-10-01T00:00:00.000000000', '1980-10-02T00:00:00.000000000',
           '1980-10-03T00:00:00.000000000', ..., '2000-09-28T00:00:00.000000000',
           '2000-09-29T00:00:00.000000000', '2000-09-30T00:00:00.000000000'],
          dtype='datetime64[ns]')

  • basin
    (basin)
    <U8
    '01013500' ... '06431500'

    array(['01013500', '01022500', '01030500', '06431500'], dtype='<U8')

• Data variables: (15)
  • dayl
    (basin, date)
    float32
    11.33 11.28 11.23 ... 11.52 11.42

    array([[11.3279915, 11.283633 , 11.231995 , ..., 11.520014 , 11.424003 ,
            11.424003 ],
           [11.423997 , 11.327997 , 11.295883 , ..., 11.520002 , 11.520002 ,
            11.423997 ],
           [11.328013 , 11.328013 , 11.231973 , ..., 11.520011 , 11.498417 ,
            11.423966 ],
           [11.4240055, 11.328003 , 11.328003 , ..., 11.520002 , 11.520002 ,
            11.4240055]], dtype=float32)

  • et_bucket
    (basin, date)
    float32
    0.8673 1.153 1.251 ... 1.213 1.155

    array([[0.86726254, 1.1531829 , 1.250723  , ..., 1.0023298 , 0.7525559 ,
            1.012144  ],
           [1.4287834 , 1.5383245 , 1.6626121 , ..., 1.2369902 , 0.8758479 ,
            1.0298797 ],
           [1.2459279 , 1.3991284 , 1.6455964 , ..., 1.1769536 , 0.8168555 ,
            1.0002254 ],
           [1.0637116 , 0.6133801 , 0.79956377, ..., 1.1372857 , 1.2130207 ,
            1.1549094 ]], dtype=float32)

  • m_bucket
    (basin, date)
    float32
    1.192e-07 1.192e-07 ... 0.1266

    array([[1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.4116375e-01, 4.7372673e-02, 9.2639700e-02, ..., 1.2195287e-01,
            1.3365604e-01, 1.2664518e-01]], dtype=float32)

  • obs_runoff
    (basin, date)
    float32
    0.551 0.5607 ... 0.441 0.4353

    array([[0.5509981 , 0.56074065, 0.55857563, ..., 0.12881881, 0.12232374,
            0.12015872],
           [0.9325426 , 0.790996  , 0.7118964 , ..., 0.2248094 , 0.22064625,
            0.21231999],
           [0.70644146, 0.67311877, 0.67311877, ..., 0.10996494, 0.11329722,
            0.1106314 ],
           [0.21764375, 0.21191628, 0.21764375, ..., 0.42956004, 0.44101498,
            0.4352875 ]], dtype=float32)

  • pr_bucket
    (basin, date)
    float32
    3.1 4.24 ... 1.192e-07 1.192e-07

    array([[3.0999999e+00, 4.2399998e+00, 8.0200005e+00, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 5.9600000e+00, ..., 5.6000000e-01,
            3.3800001e+00, 1.1920929e-07],
           [2.4200001e+00, 1.8900000e+00, 1.0060000e+01, ..., 3.8599999e+00,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07]], dtype=float32)

  • prcp
    (basin, date)
    float32
    3.1 4.24 8.02 15.27 ... 0.0 0.0 0.0

    array([[ 3.1 ,  4.24,  8.02, ...,  0.  ,  0.  ,  0.  ],
           [ 0.  ,  0.  ,  5.96, ...,  0.56,  3.38,  0.  ],
           [ 2.42,  1.89, 10.06, ...,  3.86,  0.  ,  0.  ],
           [ 0.  ,  0.  ,  0.  , ...,  0.  ,  0.  ,  0.  ]], dtype=float32)

  • ps_bucket
    (basin, date)
    float32
    1.192e-07 1.192e-07 ... 1.192e-07

    array([[1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07]], dtype=float32)

  • q_bucket
    (basin, date)
    float32
    0.02045 0.02127 ... 0.2923 0.2896

    array([[0.02044773, 0.02126511, 0.02301056, ..., 0.626426  , 0.6111313 ,
            0.59631854],
           [0.5994922 , 0.56113166, 0.56499547, ..., 0.35117748, 0.35644755,
            0.35875744],
           [0.79968697, 0.7908769 , 0.8592993 , ..., 0.3722799 , 0.3782645 ,
            0.3686408 ],
           [0.10441243, 0.10383057, 0.10331017, ..., 0.29496604, 0.29227722,
            0.28959465]], dtype=float32)

  • s_snow
    (basin, date)
    float32
    0.0 0.0 0.0 ... 2.735e+03 2.735e+03

    array([[0.0000000e+00, 0.0000000e+00, 0.0000000e+00, ..., 2.0418022e-31,
            9.2899025e-32, 5.5514055e-32],
           [0.0000000e+00, 0.0000000e+00, 0.0000000e+00, ..., 0.0000000e+00,
            0.0000000e+00, 0.0000000e+00],
           [1.1851152e-36, 7.0602143e-37, 3.9647164e-37, ..., 6.5777275e-36,
            3.2139249e-36, 1.8851024e-36],
           [6.5141640e+01, 6.5053032e+01, 6.4984924e+01, ..., 2.7354458e+03,
            2.7353196e+03, 2.7351897e+03]], dtype=float32)

  • s_water
    (basin, date)
    float32
    1.303e+03 1.305e+03 ... 984.9 983.5

    array([[1303.0043 , 1305.3451 , 1310.0562 , ..., 1507.376  , 1505.8998 ,
            1504.4343 ],
           [ 996.22375,  993.64935,  993.9165 , ...,  975.40344,  975.98334,
             976.2348 ],
           [1204.9375 , 1204.4696 , 1207.9741 , ..., 1172.6451 , 1173.3187 ,
            1172.2302 ],
           [ 835.1441 ,  834.33124,  833.6004 , ...,  986.20087,  984.8689 ,
             983.5277 ]], dtype=float32)

  • srad
    (basin, date)
    float32
    192.6 206.3 165.4 ... 350.7 303.8

    array([[192.61, 206.31, 165.35, ..., 334.41, 330.7 , 351.55],
           [272.4 , 252.49, 135.18, ..., 322.8 , 245.17, 348.5 ],
           [207.41, 222.66, 148.39, ..., 280.67, 326.83, 340.99],
           [312.34, 298.73, 318.6 , ..., 388.54, 350.68, 303.83]],
          dtype=float32)

  • tmax
    (basin, date)
    float32
    10.05 15.82 15.86 ... 22.11 19.83

    array([[10.05, 15.82, 15.86, ..., 13.22,  8.84, 14.85],
           [16.65, 17.64, 17.44, ..., 16.04,  9.85, 13.8 ],
           [13.37, 15.42, 17.4 , ..., 16.36,  9.51, 13.52],
           [21.32, 11.91, 16.93, ..., 22.75, 22.11, 19.83]], dtype=float32)

  • tmean
    (basin, date)
    float32
    6.08 10.53 11.84 ... 13.41 12.77

    array([[ 6.08 , 10.53 , 11.835, ...,  5.81 ,  1.675,  6.115],
           [11.56 , 12.92 , 14.22 , ...,  9.485,  4.19 ,  6.765],
           [ 9.11 , 10.94 , 13.63 , ...,  8.385,  2.87 ,  6.02 ],
           [14.09 ,  5.595,  9.695, ..., 12.35 , 13.41 , 12.775]],
          dtype=float32)

  • tmin
    (basin, date)
    float32
    2.11 5.24 7.81 ... 1.95 4.71 5.72

    array([[ 2.11,  5.24,  7.81, ..., -1.6 , -5.49, -2.62],
           [ 6.47,  8.2 , 11.  , ...,  2.93, -1.47, -0.27],
           [ 4.85,  6.46,  9.86, ...,  0.41, -3.77, -1.48],
           [ 6.86, -0.72,  2.46, ...,  1.95,  4.71,  5.72]], dtype=float32)

  • vp
    (basin, date)
    float32
    711.3 898.6 ... 755.6 868.9

    array([[ 711.33,  898.61, 1064.86, ...,  544.96,  400.43,  503.24],
           [ 960.76, 1095.9 , 1320.03, ...,  759.41,  549.3 ,  598.67],
           [ 869.81,  961.86, 1229.14, ...,  626.49,  458.58,  551.06],
           [ 986.25,  583.46,  720.2 , ...,  590.37,  755.56,  868.89]],
          dtype=float32)

• Indexes: (2)
  • date
    PandasIndex

    PandasIndex(DatetimeIndex(['1980-10-01', '1980-10-02', '1980-10-03', '1980-10-04',
                   '1980-10-05', '1980-10-06', '1980-10-07', '1980-10-08',
                   '1980-10-09', '1980-10-10',
                   ...
                   '2000-09-21', '2000-09-22', '2000-09-23', '2000-09-24',
                   '2000-09-25', '2000-09-26', '2000-09-27', '2000-09-28',
                   '2000-09-29', '2000-09-30'],
                  dtype='datetime64[ns]', name='date', length=7305, freq=None))

  • basin
    PandasIndex

    PandasIndex(Index(['01013500', '01022500', '01030500', '06431500'], dtype='object', name='basin'))

• Attributes: (0)

'ds_valid'

<xarray.Dataset> Size: 906kB
Dimensions:     (basin: 4, date: 3652)
Coordinates:
  * date        (date) datetime64[ns] 29kB 2000-10-01 2000-10-02 ... 2010-09-30
  * basin       (basin) <U8 128B '01013500' '01022500' '01030500' '06431500'
Data variables: (12/15)
    dayl        (basin, date) float32 58kB 11.33 11.28 11.23 ... 11.52 11.52
    et_bucket   (basin, date) float32 58kB 1.485 1.449 1.463 ... 1.703 1.2 1.113
    m_bucket    (basin, date) float32 58kB 1.192e-07 1.192e-07 ... 0.1205 0.1076
    obs_runoff  (basin, date) float32 58kB 0.1158 0.1126 ... 0.4525 0.4525
    pr_bucket   (basin, date) float32 58kB 1.192e-07 1.192e-07 ... 1.192e-07
    prcp        (basin, date) float32 58kB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0
    ...          ...
    s_water     (basin, date) float32 58kB 1.504e+03 1.502e+03 ... 1.047e+03
    srad        (basin, date) float32 58kB 327.8 331.3 314.3 ... 352.8 363.1
    tmax        (basin, date) float32 58kB 20.31 20.4 20.0 ... 26.98 20.26 19.66
    tmean       (basin, date) float32 58kB 12.23 11.92 12.17 ... 12.22 11.05
    tmin        (basin, date) float32 58kB 4.16 3.44 4.34 ... 8.61 4.18 2.44
    vp          (basin, date) float32 58kB 822.2 783.9 827.6 ... 810.6 710.9

xarray.Dataset

• Dimensions:
  • basin: 4
  • date: 3652
• Coordinates: (2)
  • date
    (date)
    datetime64[ns]
    2000-10-01 ... 2010-09-30

    array(['2000-10-01T00:00:00.000000000', '2000-10-02T00:00:00.000000000',
           '2000-10-03T00:00:00.000000000', ..., '2010-09-28T00:00:00.000000000',
           '2010-09-29T00:00:00.000000000', '2010-09-30T00:00:00.000000000'],
          dtype='datetime64[ns]')

  • basin
    (basin)
    <U8
    '01013500' ... '06431500'

    array(['01013500', '01022500', '01030500', '06431500'], dtype='<U8')

• Data variables: (15)
  • dayl
    (basin, date)
    float32
    11.33 11.28 11.23 ... 11.52 11.52

    array([[11.3279915, 11.283633 , 11.231995 , ..., 11.560081 , 11.520014 ,
            11.424003 ],
           [11.423997 , 11.327997 , 11.295883 , ..., 11.615997 , 11.520002 ,
            11.520002 ],
           [11.328013 , 11.328013 , 11.231973 , ..., 11.61598  , 11.520011 ,
            11.498417 ],
           [11.4240055, 11.328003 , 11.328003 , ..., 11.616009 , 11.520002 ,
            11.520002 ]], dtype=float32)

  • et_bucket
    (basin, date)
    float32
    1.485 1.449 1.463 ... 1.2 1.113

    array([[1.4852618 , 1.4485817 , 1.4626697 , ..., 1.7580181 , 2.0029075 ,
            1.9035351 ],
           [1.3865494 , 1.5811448 , 1.6381361 , ..., 1.7577513 , 2.252298  ,
            2.118466  ],
           [1.3531102 , 1.5591006 , 1.5008878 , ..., 1.7286984 , 2.201992  ,
            2.062271  ],
           [1.0195109 , 0.88223225, 0.7358465 , ..., 1.7031672 , 1.1998708 ,
            1.1131281 ]], dtype=float32)

  • m_bucket
    (basin, date)
    float32
    1.192e-07 1.192e-07 ... 0.1076

    array([[1.19209290e-07, 1.19209290e-07, 1.19209290e-07, ...,
            1.19209290e-07, 1.19209290e-07, 1.19209290e-07],
           [1.19209290e-07, 1.19209290e-07, 1.19209290e-07, ...,
            1.19209290e-07, 1.19209290e-07, 1.19209290e-07],
           [1.19209290e-07, 1.19209290e-07, 1.19209290e-07, ...,
            1.19209290e-07, 1.19209290e-07, 1.19209290e-07],
           [1.04674108e-01, 8.10469314e-02, 5.02984710e-02, ...,
            1.82069674e-01, 1.20517574e-01, 1.07599907e-01]], dtype=float32)

  • obs_runoff
    (basin, date)
    float32
    0.1158 0.1126 ... 0.4525 0.4525

    array([[0.11582867, 0.11258114, 0.11258114, ..., 0.6917245 , 0.81404823,
            0.93204194],
           [0.20815684, 0.20815684, 0.21231999, ..., 0.26227763, 0.29974586,
            0.31223527],
           [0.1033004 , 0.09930167, 0.09530295, ..., 0.38321117, 0.39920607,
            0.41520098],
           [0.4181051 , 0.44101498, 0.4181051 , ..., 0.44674242, 0.4524699 ,
            0.4524699 ]], dtype=float32)

  • pr_bucket
    (basin, date)
    float32
    1.192e-07 1.192e-07 ... 1.192e-07

    array([[1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.4670000e+01,
            1.0050000e+01, 1.7139999e+01],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 4.1399999e+00,
            3.9100001e+00, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 7.2600002e+00,
            3.5000000e+00, 1.8500000e+00],
           [1.8200001e+00, 1.0000000e+00, 9.8000002e-01, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07]], dtype=float32)

  • prcp
    (basin, date)
    float32
    0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0 0.0

    array([[ 0.  ,  0.  ,  0.  , ..., 14.67, 10.05, 17.14],
           [ 0.  ,  0.  ,  0.  , ...,  4.14,  3.91,  0.  ],
           [ 0.  ,  0.  ,  0.  , ...,  7.26,  3.5 ,  1.85],
           [ 1.82,  1.  ,  0.98, ...,  0.  ,  0.  ,  0.  ]], dtype=float32)

  • ps_bucket
    (basin, date)
    float32
    1.192e-07 1.192e-07 ... 1.192e-07

    array([[1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07],
           [1.1920929e-07, 1.1920929e-07, 1.1920929e-07, ..., 1.1920929e-07,
            1.1920929e-07, 1.1920929e-07]], dtype=float32)

  • q_bucket
    (basin, date)
    float32
    0.5963 0.5758 ... 0.4541 0.4487

    array([[0.59631854, 0.57577556, 0.5563412 , ..., 1.3159372 , 1.5724064 ,
            1.781463  ],
           [0.35875744, 0.34214464, 0.32507095, ..., 0.65764683, 0.6852966 ,
            0.67375565],
           [0.3686408 , 0.35311314, 0.33779895, ..., 0.43548533, 0.45141688,
            0.46435547],
           [0.28959465, 0.29006055, 0.28969634, ..., 0.45984152, 0.45406222,
            0.44873777]], dtype=float32)

  • s_snow
    (basin, date)
    float32
    5.551e-32 3.267e-32 ... 4.136e+03

    array([[ 5.55140549e-32,  3.26676522e-32,  1.79266356e-32, ...,
            -1.17272000e-30, -6.98203631e-31, -4.14229304e-31],
           [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00, ...,
            -1.45789838e+00, -1.45789766e+00, -1.45789695e+00],
           [ 1.88510239e-36,  1.11058074e-36,  6.14546683e-37, ...,
             2.57733954e-32,  1.00508896e-32,  3.12296854e-33],
           [ 2.73518970e+03,  2.73509619e+03,  2.73502954e+03, ...,
             4.13583984e+03,  4.13567529e+03,  4.13552490e+03]], dtype=float32)

  • s_water
    (basin, date)
    float32
    1.504e+03 1.502e+03 ... 1.047e+03

    array([[1504.4343 , 1502.3408 , 1500.2902 , ..., 1551.7046 , 1562.3383 ,
            1569.793  ],
           [ 976.2348 ,  974.389  ,  972.39606, ...,  999.8283 , 1001.43164,
            1000.7704 ],
           [1172.2302 , 1170.4126 , 1168.54   , ..., 1179.2683 , 1180.7859 ,
            1181.9794 ],
           [ 983.5277 ,  983.76154,  983.5788 , ..., 1050.7861 , 1048.9464 ,
            1047.2307 ]], dtype=float32)

  • srad
    (basin, date)
    float32
    327.8 331.3 314.3 ... 352.8 363.1

    array([[327.75, 331.28, 314.32, ..., 200.94, 194.26, 209.82],
           [347.79, 315.86, 305.13, ..., 172.76, 160.72, 274.52],
           [346.19, 330.63, 293.3 , ..., 194.92, 209.87, 298.2 ],
           [207.22, 271.41, 183.92, ..., 371.31, 352.84, 363.08]],
          dtype=float32)

  • tmax
    (basin, date)
    float32
    20.31 20.4 20.0 ... 20.26 19.66

    array([[20.31, 20.4 , 20.  , ..., 18.75, 20.68, 20.37],
           [18.9 , 20.14, 20.68, ..., 18.22, 22.28, 22.44],
           [19.18, 21.04, 19.24, ..., 18.55, 22.88, 23.37],
           [15.98, 15.19,  9.64, ..., 26.98, 20.26, 19.66]], dtype=float32)

  • tmean
    (basin, date)
    float32
    12.23 11.92 12.17 ... 12.22 11.05

    array([[12.235, 11.92 , 12.17 , ..., 14.13 , 16.215, 15.435],
           [11.405, 13.68 , 14.335, ..., 14.575, 18.805, 17.785],
           [10.84 , 13.13 , 12.68 , ..., 14.275, 18.4  , 17.315],
           [10.785,  8.645,  5.86 , ..., 17.795, 12.22 , 11.05 ]],
          dtype=float32)

  • tmin
    (basin, date)
    float32
    4.16 3.44 4.34 ... 8.61 4.18 2.44

    array([[ 4.16,  3.44,  4.34, ...,  9.51, 11.75, 10.5 ],
           [ 3.91,  7.22,  7.99, ..., 10.93, 15.33, 13.13],
           [ 2.5 ,  5.22,  6.12, ..., 10.  , 13.92, 11.26],
           [ 5.59,  2.1 ,  2.08, ...,  8.61,  4.18,  2.44]], dtype=float32)

  • vp
    (basin, date)
    float32
    822.2 783.9 827.6 ... 810.6 710.9

    array([[ 822.23,  783.89,  827.64, ..., 1201.15, 1379.62, 1279.99],
           [ 809.35, 1017.41, 1079.56, ..., 1314.02, 1746.54, 1502.42],
           [ 734.97,  896.45,  931.06, ..., 1240.03, 1588.34, 1340.96],
           [ 908.66,  690.87,  712.85, ..., 1031.12,  810.64,  710.9 ]],
          dtype=float32)

• Indexes: (2)
  • date
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-10-01', '2000-10-02', '2000-10-03', '2000-10-04',
                   '2000-10-05', '2000-10-06', '2000-10-07', '2000-10-08',
                   '2000-10-09', '2000-10-10',
                   ...
                   '2010-09-21', '2010-09-22', '2010-09-23', '2010-09-24',
                   '2010-09-25', '2010-09-26', '2010-09-27', '2010-09-28',
                   '2010-09-29', '2010-09-30'],
                  dtype='datetime64[ns]', name='date', length=3652, freq=None))

  • basin
    PandasIndex

    PandasIndex(Index(['01013500', '01022500', '01030500', '06431500'], dtype='object', name='basin'))

• Attributes: (0)

Feel free to generate plots from the training and validation sets to get familiar with the data

Create interpolators

As the time-series data was loaded on a one-day resolution, we need to run interpolation during the solution of the system of ODEs for adaptative-step methods and fixe-step methods with higher resolution.

# Get the basin interpolators
interpolators = get_basin_interpolators(dataset, cfg, project_dir)

Create the Conceptual model

time_idx0 = 0 # Start from the first time index - 0 for training
model_concept = get_concept_model(cfg, dataset.ds_train, interpolators,
                                  time_idx0, dataset.scaler)

Create the Neural network model

model_nn = get_nn_model(model_concept, dataset.ds_static)

Create the Pretrainer Model

pretrainer = get_nn_pretrainer(model_nn, dataset)

Pretrain the model

pretrain_ok = pretrainer.train(loss=cfg.loss_pretrain, lr=cfg.lr_pretrain,
                               epochs=cfg.epochs_pretrain,
                               disable_pbar=False,
                               any_log=False)

------------------------------------------------------------
-- Pretraining the neural network model -- (cpu)
------------------------------------------------------------
# Epoch 00001:   0%|          | 0/116 [00:00<?, ?it/s, Loss=2.1382e+00]# Epoch 00001: 100%|██████████| 116/116 [00:00<00:00, 175.79it/s, Loss=1.0095e+00]
* Plotting basin 01030500: 100%|██████████| 4/4 [00:02<00:00,  1.90it/s]
# Epoch 00002: 100%|██████████| 116/116 [00:00<00:00, 138.03it/s, Loss=3.6343e-01]
# Epoch 00003: 100%|██████████| 116/116 [00:00<00:00, 161.53it/s, Loss=3.0324e-01]
# Epoch 00004: 100%|██████████| 116/116 [00:00<00:00, 146.59it/s, Loss=2.6700e-01]
# Epoch 00005: 100%|██████████| 116/116 [00:00<00:00, 156.76it/s, Loss=2.3470e-01]
# Epoch 00006: 100%|██████████| 116/116 [00:00<00:00, 163.43it/s, Loss=2.1603e-01]
# Epoch 00007: 100%|██████████| 116/116 [00:00<00:00, 150.90it/s, Loss=2.0421e-01]
# Epoch 00008: 100%|██████████| 116/116 [00:00<00:00, 152.84it/s, Loss=2.0980e-01]
# Epoch 00009: 100%|██████████| 116/116 [00:00<00:00, 173.21it/s, Loss=2.0567e-01]
# Epoch 00010: 100%|██████████| 116/116 [00:00<00:00, 166.42it/s, Loss=2.1032e-01]
# Epoch 00011: 100%|██████████| 116/116 [00:00<00:00, 179.96it/s, Loss=2.1279e-01]
# Epoch 00012: 100%|██████████| 116/116 [00:00<00:00, 149.49it/s, Loss=2.1788e-01]
# Epoch 00013: 100%|██████████| 116/116 [00:00<00:00, 201.65it/s, Loss=2.1850e-01]
# Epoch 00014: 100%|██████████| 116/116 [00:00<00:00, 204.54it/s, Loss=2.1837e-01]
# Epoch 00015: 100%|██████████| 116/116 [00:00<00:00, 205.26it/s, Loss=2.1834e-01]
# Epoch 00016: 100%|██████████| 116/116 [00:00<00:00, 206.15it/s, Loss=2.1539e-01]
# Epoch 00017: 100%|██████████| 116/116 [00:00<00:00, 201.06it/s, Loss=2.1600e-01]
# Epoch 00018: 100%|██████████| 116/116 [00:00<00:00, 188.47it/s, Loss=2.1032e-01]
# Epoch 00019: 100%|██████████| 116/116 [00:00<00:00, 207.30it/s, Loss=2.0384e-01]
# Epoch 00020: 100%|██████████| 116/116 [00:00<00:00, 200.19it/s, Loss=1.9702e-01]
# Epoch 00021: 100%|██████████| 116/116 [00:00<00:00, 138.75it/s, Loss=1.9039e-01]
# Epoch 00022: 100%|██████████| 116/116 [00:01<00:00, 103.36it/s, Loss=1.8611e-01]
# Epoch 00023: 100%|██████████| 116/116 [00:00<00:00, 130.51it/s, Loss=1.8147e-01]
# Epoch 00024: 100%|██████████| 116/116 [00:00<00:00, 130.23it/s, Loss=1.7528e-01]
# Epoch 00025: 100%|██████████| 116/116 [00:00<00:00, 134.12it/s, Loss=1.7129e-01]
# Epoch 00026: 100%|██████████| 116/116 [00:00<00:00, 133.02it/s, Loss=1.6737e-01]
# Epoch 00027: 100%|██████████| 116/116 [00:00<00:00, 129.72it/s, Loss=1.6380e-01]
# Epoch 00028: 100%|██████████| 116/116 [00:00<00:00, 131.77it/s, Loss=1.6077e-01]
# Epoch 00029: 100%|██████████| 116/116 [00:00<00:00, 130.17it/s, Loss=1.5746e-01]
# Epoch 00030: 100%|██████████| 116/116 [00:00<00:00, 133.09it/s, Loss=1.5431e-01]
# Epoch 00031: 100%|██████████| 116/116 [00:00<00:00, 129.34it/s, Loss=1.5180e-01]
# Epoch 00032: 100%|██████████| 116/116 [00:01<00:00, 102.10it/s, Loss=1.5100e-01]
# Epoch 00033: 100%|██████████| 116/116 [00:00<00:00, 132.66it/s, Loss=1.4985e-01]
# Epoch 00034: 100%|██████████| 116/116 [00:00<00:00, 130.04it/s, Loss=1.4877e-01]
# Epoch 00035: 100%|██████████| 116/116 [00:00<00:00, 134.54it/s, Loss=1.4753e-01]
# Epoch 00036: 100%|██████████| 116/116 [00:00<00:00, 133.63it/s, Loss=1.4661e-01]
# Epoch 00037: 100%|██████████| 116/116 [00:00<00:00, 130.98it/s, Loss=1.4545e-01]
# Epoch 00038: 100%|██████████| 116/116 [00:00<00:00, 132.88it/s, Loss=1.4368e-01]
# Epoch 00039: 100%|██████████| 116/116 [00:00<00:00, 135.64it/s, Loss=1.4226e-01]
# Epoch 00040: 100%|██████████| 116/116 [00:00<00:00, 134.35it/s, Loss=1.4064e-01]
# Epoch 00041: 100%|██████████| 116/116 [00:00<00:00, 148.39it/s, Loss=1.3904e-01]
# Epoch 00042: 100%|██████████| 116/116 [00:01<00:00, 111.46it/s, Loss=1.3721e-01]
# Epoch 00043: 100%|██████████| 116/116 [00:00<00:00, 127.96it/s, Loss=1.3627e-01]
# Epoch 00044: 100%|██████████| 116/116 [00:00<00:00, 133.08it/s, Loss=1.3440e-01]
# Epoch 00045: 100%|██████████| 116/116 [00:00<00:00, 136.01it/s, Loss=1.3359e-01]
# Epoch 00046: 100%|██████████| 116/116 [00:00<00:00, 132.79it/s, Loss=1.3252e-01]
# Epoch 00047: 100%|██████████| 116/116 [00:00<00:00, 130.88it/s, Loss=1.3132e-01]
# Epoch 00048: 100%|██████████| 116/116 [00:00<00:00, 132.89it/s, Loss=1.3019e-01]
# Epoch 00049: 100%|██████████| 116/116 [00:00<00:00, 130.17it/s, Loss=1.2775e-01]
# Epoch 00050: 100%|██████████| 116/116 [00:00<00:00, 130.18it/s, Loss=1.2687e-01]
* Plotting basin 01030500: 100%|██████████| 4/4 [00:02<00:00,  1.49it/s]
* Plotting basin 01030500: 100%|██████████| 4/4 [00:09<00:00,  2.30s/it]
* Evaluating basin 06431500 (ds_train): 100%|██████████| 4/4 [00:00<00:00, 24.85it/s]
* Evaluating basin 06431500 (ds_valid): 100%|██████████| 4/4 [00:00<00:00, 46.46it/s]

Feel free to explore the model_plots and model_results folders to evaluate the outcomes of the pretraining stage. These folders provide insights into how well the model fits the target variables: \(ET\), \(M\), \(Pr\), \(Ps\), and \(Q\).

The model_weights folder contains the best-performing version of the pretrained model, which will be loaded by the hybrid model at the start of the training process. This process will be demonstrated in the next tutorial: 03-TrainHybridModel.