train.py

"""
Implementation for a temporal graph representation learning based solution by taking different 
input features as NDVI, SAR coefficients, Weather values to predict (both estimating missing values and forecasting) 
soil moisture which can leverage both spatial and temporal dependency of soil moisture values.

"""

from __future__ import division
from __future__ import print_function

import time
import argparse
import glob
import numpy as np
import os
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from scipy import stats


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

#Importing function to load data
from utils_usa import load_data_usa
from utils_spain import load_data_spain
from utils_mississippi import load_data_mississippi
from utils_alabama import load_data_alabama

#Importing GAT based models
from models.dglr.dglr_share import DGLR_SHARE
from models.dglr.dglr import DGLR

#Importing Regression models
from models.regression.Regression_ind import FC_IND
from models.regression.Regression_shared import FC_SHARE


# Training settings
parser = argparse.ArgumentParser()

parser.add_argument('--dataset', type=str, default='USA', help='Dataset to be used.')

parser.add_argument('--model', type=str, default='DGLR', help='Model to be used.')

parser.add_argument('--nofstations', type=int, default=68,help='Number of nodes in the graph.')

parser.add_argument('--noftimes', type=int, default=10,help='Number of times to run the experiment.')

parser.add_argument('--no-cuda', action='store_true', default=False, help='Disables CUDA training.')

parser.add_argument('--epochs', type=int, default=20,help='Number of epochs to train.')   

parser.add_argument('--lr', type=float, default=0.01,help='Initial learning rate.')

parser.add_argument('--hidden', type=int, default=20,help='Number of hidden units.')

parser.add_argument('--out', type=int, default=20,help='Number of output units.')

parser.add_argument('--window', type=int, default=4,help='Window Length.')

parser.add_argument('--stride', type=int, default=1,help='Stride for the Window.')

parser.add_argument('--nheads', type=int, default=3, help='Number of head attentions.')

parser.add_argument('--alpha', type=float, default=0.1, help='Alpha for the leaky_relu.')

parser.add_argument('--negslope', type=float, default=0.1, help='Negative slope for the leaky_relu of regression.')

args = parser.parse_args()
args.cuda = False

# Load data
if args.dataset == "USA":
    train_set ,test_set ,  features, SM, AList = load_data_usa()

if args.dataset == "Spain":
    train_set ,test_set ,  features, SM, AList = load_data_spain()

if args.dataset == "Mississippi":
    train_set ,test_set ,  features, SM, AList,date = load_data_mississippi()

if args.dataset == "Alabama":
    train_set ,test_set ,  features, SM, AList,date = load_data_alabama()

#RMSE
def error(actual, predicted):
    return actual - predicted

def mse(actual, predicted):
    return np.mean(np.square(error(actual, predicted)))

def rmse(actual, predicted):
    return np.sqrt(mse(actual, predicted))

#SMAPE
def smape(A, F):
    return (100 * np.sum(np.abs(F - A) / (np.abs(A) + np.abs(F))))/len(A)

#Correlation
def correlation(a,b):
    corr=stats.pearsonr(a,b)[0]
    return corr

#Average over iterations
rmse_std=np.zeros((args.noftimes,args.nofstations))
corr_std=np.zeros((args.noftimes,args.nofstations))
smape_std=np.zeros((args.noftimes,args.nofstations))

for itr in range(args.noftimes):

    # Model and optimizer
    if args.model=='DGLR':
        model_egnn = DGLR(nodes=args.nofstations,infeat=features[0].shape[1],outfeat=args.out,hiddenfeat=args.hidden,device="cuda" if args.cuda else "cpu",alpha=args.alpha,nheads=args.nheads)
        model_regression = FC_IND(nodes=args.nofstations,ntime=args.window,nfeatures=args.out,negslope=args.negslope,device="cuda" if args.cuda else "cpu")

    if args.model=='DGLR_SHARE':
        model_egnn = DGLR_SHARE(nodes=args.nofstations, infeat=features[0].shape[1],outfeat=args.out,hiddenfeat=args.hidden,device="cuda" if args.cuda else "cpu",alpha=args.alpha,nheads=args.nheads)
        model_regression = FC_SHARE(nodes=args.nofstations,ntime=args.window,nfeatures=args.out,negslope=args.negslope,device="cuda" if args.cuda else "cpu")

    optimizer_egnn = optim.Adam(model_egnn.parameters(),lr=args.lr)

    optimizer_regression = optim.Adam(model_regression.parameters(),lr=args.lr)

    #Results arrays
    groundtruth=np.zeros((args.nofstations,train_set+test_set-args.window))
    prediction=np.zeros((args.nofstations,train_set+test_set-args.window))
    trainloss=[]

    # Train model
    t_total = time.time()
    mse_loss=nn.MSELoss()

    for epoch in range(args.epochs):
        loss_train=0
        
        for i in range(args.nofstations):
            model_egnn.train()
            model_regression.train()

            optimizer_egnn.zero_grad()
            optimizer_regression.zero_grad()

            print('Epoch ',epoch,' ID ', i)

            loss_wd=0

            for wd in range(0,train_set-args.window,args.stride):
                
                train_AList=AList[wd:wd+args.window]
                train_data=features[wd:wd+args.window]
                target_SM=torch.FloatTensor([SM[wd+args.window][i]])

                nodeemb = model_egnn(train_AList, train_data)
                predict_SM=model_regression(nodeemb,i)

                loss_wd=loss_wd + mse_loss(predict_SM, target_SM)

                print("Target: ",target_SM)
                print("Prediction: ",predict_SM)
                print('------')

                if epoch==args.epochs-1:
                    prediction[i][wd]=predict_SM
                    groundtruth[i][wd]=target_SM

            print('ID Loss ',loss_wd)

            loss_train=loss_train+loss_wd.item()
        
            loss_wd.backward()

            optimizer_egnn.step()
            optimizer_regression.step()   
            print('------------------')
            
        trainloss.append(loss_train)

        print('Training Loss ',loss_train)

        print('================')

    print("Optimization Finished!")
    print("Total time elapsed: {:.4f}s".format(time.time() - t_total))
    print('=========')

    # Test
    for i in range(args.nofstations):
        for wd in range(train_set-args.window,train_set+test_set-args.window,args.stride):

            test_AList=AList[wd:wd+args.window]
            test_data=features[wd:wd+args.window]
            target_SM=torch.FloatTensor([SM[wd+args.window][i]])

            model_egnn.eval()
            model_regression.eval()

            nodeemb = model_egnn(test_AList, test_data)
            predict_SM=model_regression(nodeemb,i)

            loss_test = mse_loss(predict_SM, target_SM)

            print("Test set results:","loss= {:.4f}".format(loss_test.item()))

            print('For ID ',i)
            print("Target: ",target_SM)
            print("Prediction: ",predict_SM)
            print('------')
            
            prediction[i][wd]=predict_SM
            groundtruth[i][wd]=target_SM


        rmse_std[itr][i]=rmse(groundtruth[i][train_set-args.window:train_set+test_set-args.window], prediction[i][train_set-args.window:train_set+test_set-args.window])
        corr_std[itr][i]=correlation(groundtruth[i][train_set-args.window:train_set+test_set-args.window], prediction[i][train_set-args.window:train_set+test_set-args.window])
        smape_std[itr][i]=smape(groundtruth[i][train_set-args.window:train_set+test_set-args.window], prediction[i][train_set-args.window:train_set+test_set-args.window])


mean_rmse_itr=np.mean(rmse_std,axis=0)
mean_smape_itr=np.mean(smape_std,axis=0)
mean_corr_itr=np.mean(corr_std,axis=0)
std_rmse_itr=np.std(rmse_std,axis=0)
std_smape_itr=np.std(smape_std,axis=0)
std_corr_itr=np.std(corr_std,axis=0)

print('Average Test RMSE: ',np.mean(mean_rmse_itr),'+/-',np.mean(std_rmse_itr))
print('Average Test Correlation: ',np.mean(mean_corr_itr),'+/-',np.mean(std_corr_itr))
print('Average Test SMAPE: ',np.mean(mean_smape_itr),'+/-',np.mean(std_smape_itr))