trainPUR-Reg.py

import sys
sys.path.append('../')
from hydroDL import master
from hydroDL.master import default
from hydroDL.data import camels
from hydroDL.model import rnn, crit, train


import json
import os
import numpy as np
import torch
import random


# Options for different interface
interfaceOpt = 1
# ==1 default, the improved and more interpretable version. It's easier to see the data flow, model setup and training
# process. Recommended for most users.
# ==0 the original "pro" version we used to run heavy jobs for the paper. It was later improved for clarity to obtain option 1.
# Results are very similar for two options and have little difference in computational performance.

Action = [1, 2]
# Using Action options to control training different models
# 1: Train Base LSTM PUR Models without integrating any soft info
# 2: Train CNN-LSTM to integrate FDCs

# Hyperparameters
EPOCH = 300
BATCH_SIZE=100
RHO=365
HIDDENSIZE=256
saveEPOCH = 10 # save model for every "saveEPOCH" epochs
Ttrain=[19851001, 19951001] # training period
LCrange = [19851001, 19951001]

# Define root directory of database and output
# Modify this based on your own location of CAMELS dataset
# Following the data download instruction in README file, you should organize the folders like
# 'your/path/to/Camels/basin_timeseries_v1p2_metForcing_obsFlow' and 'your/path/to/Camels/camels_attributes_v2.0'
# Then 'rootDatabase' here should be 'your/path/to/Camels'
# You can also define the database directory in hydroDL/__init__.py by modifying pathCamels['DB'] variable

rootDatabase = os.path.join(os.path.sep, 'scratch', 'Camels')  # CAMELS dataset root directory
camels.initcamels(rootDatabase)  # initialize three camels module-scope variables in camels.py: dirDB, gageDict, statDict

rootOut = os.path.join(os.path.sep, 'data', 'rnnStreamflow')  # Model output root directory

# define random seed
# seedid = [159654, 109958, 257886, 142365, 229837, 588859] # six seeds randomly generated using np.random.uniform
seedid = 159654
random.seed(seedid)
torch.manual_seed(seedid)
np.random.seed(seedid)
torch.cuda.manual_seed(seedid)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Fix seed for training, change it to have different runnings with different seeds
# We use the mean discharge of 6 runnings with different seeds to account for randomness


# directory to save training results
exp_name='PUR'
exp_disp='Testrun'
save_path = os.path.join(exp_name, exp_disp, str(seedid))

# Divide CAMELS dataset into 7 PUR regions
gageinfo = camels.gageDict
hucinfo = gageinfo['huc']
gageid = gageinfo['id']
# get the id list of each region
regionID = list()
regionNum = list()
regionDivide = [ [1,2], [3,6], [4,5,7], [9,10], [8,11,12,13], [14,15,16,18], [17] ] # seven regions
for ii in range(len(regionDivide)):
    tempcomb = regionDivide[ii]
    tempregid = list()
    for ih in tempcomb:
        tempid = gageid[hucinfo==ih].tolist()
        tempregid = tempregid + tempid
    regionID.append(tempregid)
    regionNum.append(len(tempregid))

# Only for interfaceOpt=0 using multiple GPUs, not used here
# cid = 0   # starting GPU id
# gnum = 6  # how many GPUs you have

# Region withheld as testing target. Take region 1 as an example.
# Change this to 1,2,..,7 to run models for all 7 PUR regions in CONUS.
testRegion = 1

iexp = testRegion - 1  # index
TestLS = regionID[iexp] # basin ID list for testing, should be withheld for training
TrainLS = list(set(gageid.tolist()) - set(TestLS)) # basin ID for training
gageDic = {'TrainID': TrainLS, 'TestID': TestLS}

# prepare the training dataset
optData = default.optDataCamels
optData = default.update(optData, tRange=Ttrain, subset=TrainLS, lckernel=None, fdcopt=False)
climateList = camels.attrLstSel + ['p_mean','pet_mean','p_seasonality','frac_snow','aridity','high_prec_freq',
                                    'high_prec_dur','low_prec_freq','low_prec_dur']
# climateList = ['slope_mean', 'area_gages2', 'frac_forest', 'soil_porosity', 'max_water_content']
# climateList = []
optData = default.update(optData, varT=camels.forcingLst, varC= climateList)
# varT: forcing used for training   varC: attributes used for training

# The above controls what attributes used for training, change varC for input-selection-ensemble
# for 5 attributes model: climateList = ['slope_mean', 'area_gages2', 'frac_forest', 'soil_porosity', 'max_water_content']
# for no-attribute model: varC = []
# the input-selection ensemble represents using the mean prediction of full, 5-attr and no-attr models,
# in total the mean of 3(different attributes)*6(different random seeds) = 18 models

if interfaceOpt == 1:
# read data from CAMELS dataset
    df = camels.DataframeCamels(
        subset=optData['subset'], tRange=optData['tRange'])
    x = df.getDataTs(
        varLst=optData['varT'],
        doNorm=False,
        rmNan=False)
    y = df.getDataObs(
        doNorm=False,
        rmNan=False,
        basinnorm=True)
    # "basinnorm = True" will call camels.basinNorm() on the original discharge data. This will transform discharge
    # from ft3/s to mm/day and then divided by mean precip to be dimensionless. output = discharge/(area*mean_precip)
    c = df.getDataConst(
        varLst=optData['varC'],
        doNorm=False,
        rmNan=False)

    # process, do normalization and remove nan
    series_data = np.concatenate([x, y], axis=2)
    seriesvarLst = camels.forcingLst + ['runoff']
    # calculate statistics for normalization and save to a dictionary
    statDict = camels.getStatDic(attrLst=climateList, attrdata=c, seriesLst=seriesvarLst, seriesdata=series_data)
    # normalize
    attr_norm = camels.transNormbyDic(c, climateList, statDict, toNorm=True)
    attr_norm[np.isnan(attr_norm)] = 0.0
    series_norm = camels.transNormbyDic(series_data, seriesvarLst, statDict, toNorm=True)

    # prepare the inputs
    xTrain = series_norm[:,:,:-1] # forcing, not include obs
    xTrain[np.isnan(xTrain)]  = 0.0
    yTrain = np.expand_dims(series_norm[:,:,-1], 2)

    if attr_norm.size == 0: # [], no-attribute case
        attrs = None
        Nx = xTrain.shape[-1]
    else:
        # with attributes
        attrs=attr_norm
        Nx = xTrain.shape[-1] + attrs.shape[-1]
    Ny = yTrain.shape[-1]

# define loss function
optLoss = default.optLossRMSE
lossFun = crit.RmseLoss()
# configuration for training
optTrain = default.update(default.optTrainCamels, miniBatch=[BATCH_SIZE, RHO], nEpoch=EPOCH, saveEpoch=saveEPOCH, seed=seedid)
hucdic = 'Reg-'+str(iexp+1)+'-Num'+str(regionNum[iexp])

if 1 in Action:
# Train base LSTM PUR model
    out = os.path.join(rootOut, save_path, hucdic,'Reg-85-95-Sub-Full')
    # out = os.path.join(rootOut, save_path, hucdic,'Reg-85-95-Sub-5attr')
    # out = os.path.join(rootOut, save_path, hucdic,'Reg-85-95-Sub-Noattr')
    if not os.path.isdir(out):
        os.makedirs(out)
    # log training gage information
    gageFile = os.path.join(out, 'gage.json')
    with open(gageFile, 'w') as fp:
        json.dump(gageDic, fp, indent=4)
    # define model config
    optModel = default.update(default.optLstm, name='hydroDL.model.rnn.CudnnLstmModel', hiddenSize=HIDDENSIZE)

    if interfaceOpt == 1:
        # define, load and train model
        optModel = default.update(optModel, nx=Nx, ny=Ny)
        model = rnn.CudnnLstmModel(nx=optModel['nx'], ny=optModel['ny'], hiddenSize=optModel['hiddenSize'])
        # Wrap up all the training configurations to one dictionary in order to save into "out" folder
        masterDict = master.wrapMaster(out, optData, optModel, optLoss, optTrain)
        master.writeMasterFile(masterDict)
        # log statistics
        statFile = os.path.join(out, 'statDict.json')
        with open(statFile, 'w') as fp:
            json.dump(statDict, fp, indent=4)
        # Train the model
        trainedModel = train.trainModel(
            model,
            xTrain,
            yTrain,
            attrs,
            lossFun,
            nEpoch=EPOCH,
            miniBatch=[BATCH_SIZE, RHO],
            saveEpoch=saveEPOCH,
            saveFolder=out)

    if interfaceOpt == 0:
        # Only need to pass the wrapped configuration dict 'masterDict' for training
        # nx, ny will be automatically updated later
        masterDict = master.wrapMaster(out, optData, optModel, optLoss, optTrain)
        master.train(masterDict)

        ## Not used here.
        ## A potential way to run batch jobs simultaneously in background through multiple GPUs and Linux screens.
        ## To use this, must manually set the "pathCamels['DB']" in hydroDL/__init__.py as your own root path of CAMELS data.
        ## Use the following master.runTrain() instead of the above master.train().
        # master.runTrain(masterDict, cudaID=cid % gnum, screen='test-'+str(cid))
        # cid = cid + 1

if 2 in Action:
# Train CNN-LSTM PUR model to integrate FDCs
    # LCrange defines from which period to get synthetic FDC
    LCTstr = str(LCrange[0]) + '-' + str(LCrange[1])
    out = os.path.join(rootOut, save_path, hucdic, 'Reg-85-95-Sub-Full-FDC' + LCTstr)
    # out = os.path.join(rootOut, save_path, hucdic, 'Reg-85-95-Sub-5attr-FDC' + LCTstr)
    # out = os.path.join(rootOut, save_path, hucdic, 'Reg-85-95-Sub-Noattr-FDC' + LCTstr)
    if not os.path.isdir(out):
        os.makedirs(out)
    gageFile = os.path.join(out, 'gage.json')
    with open(gageFile, 'w') as fp:
        json.dump(gageDic, fp, indent=4)

    optData = default.update(default.optDataCamels, tRange=Ttrain, subset=TrainLS,
                             lckernel=LCrange, fdcopt=True)
    # define model
    convNKS = [(10, 5, 1), (5, 3, 3), (1, 1, 1)]
    # CNN parameters for 3 layers: [(Number of kernels 10,5,1), (kernel size 5,3,3), (stride 1,1,1)]
    optModel = default.update(default.optCnn1dLstm, name='hydroDL.model.rnn.CNN1dLCmodel',
                              hiddenSize=HIDDENSIZE, convNKS=convNKS, poolOpt=[2, 2, 1])  # use CNN-LSTM model

    if interfaceOpt == 1:
        # load data and create synthetic FDCs as inputs
        dffdc = camels.DataframeCamels(subset=optData['subset'], tRange=optData['lckernel'])
        datatemp = dffdc.getDataObs(
            doNorm=False, rmNan=False, basinnorm=True)
        # normalize data
        dadata = camels.transNormbyDic(datatemp, 'runoff', statDict, toNorm=True)
        dadata = np.squeeze(dadata)  # dim Nbasin*Nday
        fdcdata = master.master.calFDC(dadata)
        print('FDC was calculated and used!')
        xIn = (xTrain, fdcdata)

        # load model
        Nobs = xIn[1].shape[-1]
        optModel = default.update(optModel, nx=Nx, ny=Ny, nobs=Nobs) # update input dims
        convpara = optModel['convNKS']

        model = rnn.CNN1dLCmodel(
            nx=optModel['nx'],
            ny=optModel['ny'],
            nobs=optModel['nobs'],
            hiddenSize=optModel['hiddenSize'],
            nkernel=convpara[0],
            kernelSize=convpara[1],
            stride=convpara[2],
            poolOpt=optModel['poolOpt'])
        print('CNN1d Local calibartion Kernel is used!')

        # Wrap up all the training configurations to one dictionary in order to save into "out" folder
        masterDict = master.wrapMaster(out, optData, optModel, optLoss, optTrain)
        master.writeMasterFile(masterDict)
        # log statistics
        statFile = os.path.join(out, 'statDict.json')
        with open(statFile, 'w') as fp:
            json.dump(statDict, fp, indent=4)

        # Train the model
        trainedModel = train.trainModel(
            model,
            xIn,  # need to well defined
            yTrain,
            attrs,
            lossFun,
            nEpoch=EPOCH,
            miniBatch=[BATCH_SIZE, RHO],
            saveEpoch=saveEPOCH,
            saveFolder=out)

    if interfaceOpt == 0:
        # Only need to pass the wrapped configuration 'masterDict' for training
        # nx, ny, nobs will be automatically updated later
        masterDict = master.wrapMaster(out, optData, optModel, optLoss, optTrain)
        master.train(masterDict)  # train model

        # master.runTrain(masterDict, cudaID=cid % gnum, screen='test-'+str(cid))
        # cid = cid + 1