train_graphrefine.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. Graph Reconstruction module is added
along with the soil moisture prediction module.

"""

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 model
from models.dglr.dglr import DGLR

#Importing Regression model
from models.regression.Regression_ind import FC_IND

# Training settings
parser = argparse.ArgumentParser()

parser.add_argument('--dataset', type=str, default='USA', help='Dataset 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
    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")

    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=[]
    predictionloss=[]
    graphrefineloss=[]

    # Train model
    gamma=1
    delta=1000

    t_total = time.time()
    A=AList

    mse_loss=nn.MSELoss()
    crossentropyloss=nn.BCEWithLogitsLoss()

    for epoch in range(args.epochs):
        loss_train=0
        loss_pred=0
        loss_graph_refine=0
        Ait = torch.zeros(size=(train_set,args.nofstations,args.out))
        
        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
            loss_refine_wd=0

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

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

                recons_adj_id=torch.cat(([F.relu(nodeemb[j:j+args.nofstations] @ nodeemb[j:j+args.nofstations].T)[i]/torch.sum(F.relu(nodeemb[j:j+args.nofstations] @ nodeemb[j:j+args.nofstations].T)[i])  for j in range(0,args.nofstations*(args.window+1),args.nofstations)]),dim=0)
                target_adj_id=torch.cat(([AList[j][i] for j in range(wd,wd+args.window+1)]),dim=0)

                loss_wd=loss_wd + mse_loss(predict_SM, target_SM)
                loss_refine_wd=loss_refine_wd+crossentropyloss(recons_adj_id,target_adj_id)

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

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

                Ait[wd:wd+args.window+1,i,:]=torch.stack(([nodeemb[j:j+args.nofstations][i].clone().detach() for j in range(0,args.nofstations*(args.window+1),args.nofstations)]),dim=0)

            print('ID Loss ',loss_wd)
            print('Graph Refine ID Loss ',loss_refine_wd)

            loss_train=loss_train+loss_wd.item()+loss_refine_wd.item()
            loss_pred=loss_pred+loss_wd.item()
            loss_graph_refine=loss_graph_refine+loss_refine_wd.item()
            
            loss_total=gamma * loss_wd+ delta* loss_refine_wd 
            loss_total.backward()

            optimizer_egnn.step()
            optimizer_regression.step()   
            print('------------------')

        A=[ F.relu(Ait[j] @ Ait[j].T)/torch.sum(F.relu(Ait[j] @ Ait[j].T),dim=1) for j in range(train_set)] 

        trainloss.append(loss_train)
        predictionloss.append(loss_pred)
        graphrefineloss.append(loss_graph_refine)

        print('Training Loss ',loss_train)
        print('Prediction Training Loss ',loss_pred)
        print('Graph Refine Training Loss ',loss_graph_refine)

        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=[A[train_set-1]]* 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))