 



# High level summary:
# The goal of the script is to train one of many neural networks
# on previously process ASTER imagery.

# Input data:
# This input data is expected to be in the form of numpy arrays
# representing sequences of cropped thermal infrared images from
# 1999-2012. More specifically, this is anticipated to be temperature
# above the environmental background (kelvin/celsius) for volcanoes.
# This input data is gathered from the previously numbered scripts in
# the directory.

# Script arguments: (model selection and training data)
# Model selection:
# This training code has been designed with several different
# neural network architectures in mind. Namely,
#     1) "LSTM"
#     2 LSTM variants that handle irregular time intervals
#         * "Time-LSTM"
#         * "TimeAwareLSTM"
#     3) Convolutional LSTM (that allows images sequences)
#         * "ConvLSTM"
#     4) New/proposed models that allow image sequences that are
#        irregular through time
#         * Convolutional LSTM + Time-LSTM
#             - "ConvTimeLSTM"
#         * Convolutional LSTM + TimeAwareLSTM
#             - "ConvTimeAwareLSTM"
# The quoted keyword-names are the required input for this argument and will
# select which neural network is trained; this in turn affects preprocessing
# the data into different forms, the training batch size, etc...
#
# Training data:
# This training code has been designed with two sets of training data in mind.
# These two options are:
#     1) Single volcano; in this case, the argument should be the directory
#        name containing the volcano-specific data. All of these directories
#        are expected to be within the "data" directory which is located in
#        current working directory.
#     2) All volcanoes; in this case the argument should be "all".

# Hardware:
# This code was developed for Penn State's Advanced CyberInfrastructure's
# GPU instances which utilize dual NVIDIA Tesla K80 GPU cards resulting
# in 4 GPUs with 11 GB of RAM each
# Later adapted for more powerful hardware...
# 10 GPUs (Tesla V100s-PCI), 32 GB of RAM each


import argparse

parser = argparse.ArgumentParser(description='level2 training')
parser.add_argument('--model_selection', type=str, help='model')
parser.add_argument('--training_data_set', type=str, help='training_data_set')
parser.add_argument('--n_layers', type=int, help='n_layers')
parser.add_argument('--hidden_dim_ls', nargs='+', type=int, help='hidden_dim_ls')
args = parser.parse_args()

model_selection = args.model_selection
training_data_set = args.training_data_set
n_layers = args.n_layers
hidden_dim_ls = args.hidden_dim_ls


# Gathering function inputs
### model_selection = input("Which model do you select?")
print('Model selected:', model_selection)
assert(model_selection in ['Identity', 'AR', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM', 'ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet'])
### training_data_set = input("Which set of training data do you want to use?")
print('Training data set:', training_data_set)
### n_layers = int(input("Enter number of layers: ")) 
### hidden_dim_ls = []
# iterating till the range 
### for i in range(n_layers): 
###     layer_dim = int(input()) 
###     hidden_dim_ls.append(layer_dim)
print('Hidden layers:', hidden_dim_ls)


# Libraries and imports
import os
import numpy as np
import pandas as pd
from datetime import datetime
import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.utils import data
if model_selection == 'LSTM':
	from models.LSTM import StackedLSTM as LSTM_Model
elif model_selection == 'Identity':
	from models.Identity import Identity2 as LSTM_Model
elif model_selection == 'AR':
	from models.AR import AR as LSTM_Model
elif model_selection == 'TimeLSTM':
	from models.TimeLSTM import StackedTimeLSTM as LSTM_Model
elif model_selection == 'TimeAwareLSTM':
	from models.TimeAwareLSTM import StackedTimeAwareLSTM as LSTM_Model
elif model_selection == 'ConvLSTM':
	from models.ConvLSTM import ConvLSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTM':
	from models.ConvTimeLSTM2 import ConvTime_LSTM2 as LSTM_Model
elif model_selection == 'ConvTimeAwareLSTM':
	from models.ConvTimeAwareLSTM2 import ConvTimeAware_LSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTMUnet':
	from models.ConvTimeLSTMUnet import ConvTime_LSTM2_Unet as LSTM_Model
from helper_fns.processing import scale_and_remove_na
from helper_fns.efcnt_data_AST08 import efficient_Dataset
from optimization import ssim
import numpy.ma as ma
from scipy.spatial import KDTree


# Setting seed for reproducibility
torch.manual_seed(0)


# Basic data import, step 0
print("Importing and formatting data")
volcanoes = os.listdir("data")
assert((training_data_set in volcanoes) or (training_data_set == "all"))


# Training parameters
# This needs to actually be variable, will do with later exploration
assert(hidden_dim_ls[0] in [2000, 1028, 128, 64])
if hidden_dim_ls[0] == 2000: # [2000, 2000, 1]
	assert(model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM'])
	# batch_size_dict = {'AR':84, 'Identity':84, 'LSTM':4, 'TimeLSTM':2, 'TimeAwareLSTM':2}
	batch_size_dict = {'AR':84, 'Identity':84, 'LSTM':24, 'TimeLSTM':15, 'TimeAwareLSTM':15}
	batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 1028: # [1028, 512, 1]
	assert(model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM'])
	# batch_size_dict = {'LSTM':12, 'TimeLSTM':6, 'TimeAwareLSTM':10}
	batch_size_dict = {'LSTM':36, 'TimeLSTM':18, 'TimeAwareLSTM':30}
	batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 128: 
	if hidden_dim_ls[1] == 64: # [128, 64, 64, 1] for ConvLSTMs
		assert(model_selection in ['ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM'])
		batch_size_dict = {'ConvLSTM':320, 'ConvTimeLSTM':216, 'ConvTimeAwareLSTM':320}
		batch_size = batch_size_dict[model_selection]
	else: # [128, 1] for LSTMs
		assert(model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM'])
		# batch_size_dict = {'LSTM':128, 'TimeLSTM':80, 'TimeAwareLSTM':114}
		batch_size_dict = {'LSTM':360, 'TimeLSTM':224, 'TimeAwareLSTM':320}
		batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 64: # [64, 64, 1]; [64, 128, 64, 1] for Unet
	assert(model_selection in ['ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet'])
	# batch_size_dict = {'ConvLSTM':108, 'ConvTimeLSTM':76, 'ConvTimeAwareLSTM':124, 'ConvTimeLSTMUnet':36}
	batch_size_dict = {'ConvLSTM':320, 'ConvTimeLSTM':216, 'ConvTimeAwareLSTM':352, 'ConvTimeLSTMUnet':80}
	batch_size = batch_size_dict[model_selection]
lag_dict = {"all":6, "Erebus":3, "ErtaAle":9, "Etna":3, "Kilauea":10, "Masaya":3, "Nyamuragira":3, "Nyiragongo":3, "Pacaya":3, "Puuoo":4}

num_input_scenes = lag_dict[training_data_set]
train_percent = 0.85
out_samp_perc = 0.15


# Removing possible directory clutter
try:
	volcanoes.remove(".ipynb_checkpoints")
except ValueError as e:
	do = 'nothing'

count = 0
vol_cutoff_indices = []
vol_cutoff_indices_test = []
vol_name_ls = []
for vol in os.listdir('data'):
	### Basic data import ###
	numpy_data_location = "AST_08_data/" + vol + "/background_subtracted"
	table_data_location = "data/" + vol + "/good_df.csv"
	volcano_scenes = np.load(numpy_data_location, allow_pickle = True)
	tabular_metadata = pd.read_csv(table_data_location)
	### Taking care of the unretrievable scenes for Erebus ###
	### This should always be done when referencing processed data (AST_08_data) with metadata (data) ###
	if vol == 'Erebus':
		tabular_metadata['acquisition_datetimes'] = [file.split('_')[2] for file in tabular_metadata['nighttime_volcano_files']]
		tabular_metadata = tabular_metadata[tabular_metadata.acquisition_datetimes != '00304102013143221'] 
		tabular_metadata = tabular_metadata[tabular_metadata.acquisition_datetimes != '00308162018142109']
		tabular_metadata = tabular_metadata.reset_index()
	### Separate model inputs and outputs
	# Determine number in each partition
	train_n = int(np.floor((len(volcano_scenes) - num_input_scenes)*train_percent))
	out_n = int(np.floor((len(volcano_scenes) - num_input_scenes)*out_samp_perc))
	# For every data partition
	# Array for the prior scenes
	#   "train_n - 1" is to remove the first scene that wont have an associated TimeAwareLSTM time interval
	x_scenes_train = ma.zeros([train_n - 1, num_input_scenes, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	x_scenes_test = ma.zeros([out_n, num_input_scenes, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	# Array for the time differences between scenes
	time_differences_train = np.ones(x_scenes_train.shape)
	time_differences_test = np.ones(x_scenes_test.shape)
	# Array for the target scenes
	y_scenes_train = ma.zeros([train_n - 1, 1, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	y_scenes_test = ma.zeros([out_n, 1, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	# Formatting the string dates as datetime objects
	formatted_dates = [datetime.strptime(date, '%Y-%m-%d') for date in tabular_metadata['dates']]
	# For all observations - acknowledging that the first (n-1) wont have n prior observations
	#     Also, the first data point wont have a TimeAwareLSTM time value, so it is omitted
	for i in range(num_input_scenes + 1, x_scenes_train.shape[0] + x_scenes_test.shape[0] + num_input_scenes+1):
		if i < (train_n + num_input_scenes):
			# Store the image data
			x_scenes_train[i - num_input_scenes - 1, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :]
			y_scenes_train[i - num_input_scenes - 1, 0, :, :] = volcano_scenes[i, :, :]
			# Compute the time differences and store
			# Time LSTM uses forward-time interval
			if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
				dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]
				dates_i = formatted_dates[(i - num_input_scenes):i]
				for j in range(len(dates_i_plus_1)):
					time_differences_train[i - num_input_scenes - 1, j] = (dates_i_plus_1[j] - dates_i[j]).days
			# While TimeAwareLSTM uses backwards-time interval
			else:
				dates_i = formatted_dates[(i - num_input_scenes):i]
				dates_i_minus_1 = formatted_dates[(i - num_input_scenes - 1):(i - 1)]
				for j in range(len(dates_i)):
					time_differences_train[i - num_input_scenes - 1, j] = (dates_i[j] - dates_i_minus_1[j]).days
		else:
			# Store the image data
			x_scenes_test[i - train_n - num_input_scenes, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :]
			y_scenes_test[i - train_n - num_input_scenes, 0, :, :] = volcano_scenes[i, :, :]
			# Compute the time differences and store
			# Time LSTM uses forward-time interval
			if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
				dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]
				dates_i = formatted_dates[(i - num_input_scenes):i]
				for j in range(len(dates_i_plus_1)):
					time_differences_test[i - num_input_scenes - train_n - 1, j] = (dates_i_plus_1[j] - dates_i[j]).days
			# While TimeAwareLSTM uses backwards-time interval
			else:
				dates_i = formatted_dates[(i - num_input_scenes):i]
				dates_i_minus_1 = formatted_dates[(i - num_input_scenes - 1):(i - 1)]
				for j in range(len(dates_i)):
					time_differences_test[i - num_input_scenes - train_n - 1, j] = (dates_i[j] - dates_i_minus_1[j]).days
	if count == 0:
		x_train = x_scenes_train
		t_train = time_differences_train
		y_train = y_scenes_train
		x_test = x_scenes_test
		t_test = time_differences_test
		y_test = y_scenes_test
	else:
		x_train = ma.append(x_train, x_scenes_train, axis = 0)
		t_train = np.append(t_train, time_differences_train, axis = 0)
		y_train = ma.append(y_train, y_scenes_train, axis = 0)
		x_test = ma.append(x_test, x_scenes_test, axis = 0)
		t_test = np.append(t_test, time_differences_test, axis = 0)
		y_test = ma.append(y_test, y_scenes_test, axis = 0)
	count += 1
	vol_cutoff_indices.append(y_train.shape[0])
	vol_cutoff_indices_test.append(y_test.shape[0])
	vol_name_ls.append(vol)
	print('\timported ' + str(x_scenes_train.shape[0]) + ' training scenes from ' + vol)
	print('\t\timported ' + str(x_scenes_test.shape[0]) + ' test scenes from ' + vol)


print("Processing data")
# Gap fill missing values with previous ones or nearest neighbor
for i in range(len(x_train)):
	for j in range(x_train.shape[1]):
		# Identifying missing values
		ma = np.ma.masked_invalid(x_train[i, j, :, :])
		# If the mask found NA values
		if not np.all(x_train[i, j, :, :].mask == False):
			# Using previous value to fill
			if j == 0:
				# Unless there's no previous, then nearest neighbor interpolate
				if i == 0 or i in vol_cutoff_indices:
					print('\t\tNo prior scene; nearest neighbor interpolation')
					scene = x_train[i, j, :, :]
					# Explicitly retrieving good and bad locations from the mask
					# note that these differ from (x, y) which earlier defined bad pixel locations
					X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
					xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
					xybad = np.array((X[scene.mask], Y[scene.mask])).T
					# Performing the nearest neighbor gap-filling
					x_train[i, j, :, :][ma.mask == True] = scene[~scene.mask][KDTree(xygood).query(xybad)[1]]
				else:
					x_train[i, j, :, :][ma.mask == True] = x_train[i-1, j, :, :][ma.mask == True]
					t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][ma.mask == True] + t_train[i-1, j, :, :][ma.mask == True]
			else:
				x_train[i, j, :, :][ma.mask == True] = x_train[i, j-1, :, :][ma.mask == True]
				t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][ma.mask == True] + t_train[i, j-1, :, :][ma.mask == True]
for i in range(len(x_test)):
	for j in range(x_test.shape[1]):
		# Identifying missing values
		ma = np.ma.masked_invalid(x_test[i, j, :, :])
		# If the mask found NA values
		if not np.all(ma.mask == False):
			# Using previous value to fill
			if j == 0:
				# Unless there's no previous, then nearest neighbor interpolate
				if i == 0 or i in vol_cutoff_indices_test:
					scene = x_test[i, j, :, :]
					# Explicitly retrieving good and bad locations from the mask
					# note that these differ from (x, y) which earlier defined bad pixel locations
					X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
					xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
					xybad = np.array((X[scene.mask], Y[scene.mask])).T
					# Performing the nearest neighbor gap-filling
					x_test[i, j, :, :][ma.mask == True] = scene[~scene.mask][KDTree(xygood).query(xybad)[1]]
				else:
					x_test[i, j, :, :][ma.mask == True] = x_test[i-1, j, :, :][ma.mask == True]
					t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma.mask == True] + t_test[i-1, j, :, :][ma.mask == True]
			else:
				x_test[i, j, :, :][ma.mask == True] = x_test[i, j-1, :, :][ma.mask == True]
				t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma.mask == True] + t_test[i, j-1, :, :][ma.mask == True]
for i in range(len(y_train)):
	ma = np.ma.masked_invalid(y_train[i, :, :])
	# If the mask found NA values
	if not np.all(ma.mask == False):
		if i == 0 or i in vol_cutoff_indices: 
			y_train[i, :, :][ma.mask == True] = x_train[i, [-1], :, :][ma.mask == True]
		else:
			y_train[i, :, :][ma.mask == True] = y_train[i-1, :, :, :][ma.mask == True]
for i in range(len(y_test)):
	ma = np.ma.masked_invalid(y_test[i, :, :, :])
	# If the mask found NA values
	if not np.all(ma.mask == False):
		if i == 0 or i in vol_cutoff_indices_test:
			y_test[i, :, :, :][ma.mask == True] = x_test[i, [-1], :, :][ma.mask == True]
		else:
			y_test[i, :, :, :][ma.mask == True] = y_test[i-1, :, :, :][ma.mask == True]


# Scale 0-1 using min and max from training set
# Also attempts to find NAs and replace with 0s, but those shouldnt exist anymore
stored_parameters = np.zeros([2, 3])
x_train, stored_parameters = scale_and_remove_na(x_train, stored_parameters, 0)
x_test = (x_test - stored_parameters[0, 0]) / (stored_parameters[1, 0] - stored_parameters[0, 0])
t_train, stored_parameters = scale_and_remove_na(t_train, stored_parameters, 1)
t_test = (t_test - stored_parameters[0, 1]) / (stored_parameters[1, 1] - stored_parameters[0, 1])
y_train, stored_parameters = scale_and_remove_na(y_train, stored_parameters, 2)
y_test = (y_test - stored_parameters[0, 2]) / (stored_parameters[1, 2] - stored_parameters[0, 2])
np.save("outputs/transformation_parameters.npy", stored_parameters)


# Convert to torch tensors
x_train = torch.from_numpy(x_train).type(torch.FloatTensor)
x_test = torch.from_numpy(x_test).type(torch.FloatTensor)
t_train = torch.from_numpy(t_train).type(torch.FloatTensor)
t_test = torch.from_numpy(t_test).type(torch.FloatTensor)
y_train = torch.from_numpy(y_train).type(torch.FloatTensor)
y_test = torch.from_numpy(y_test).type(torch.FloatTensor)


# Defining model parameters
# Picking one of the like-sequence tensors within the list to set parameters
print("Setting up methods")
channels = 1
height = x_train.shape[2]
width = x_train.shape[3]
lstm_model = LSTM_Model(input_dim=channels,hidden_dim=hidden_dim_ls,GPU=True,input_size=(height,width),num_layers=n_layers)
# Print number of model parameters
total_params = sum(p.numel() for p in lstm_model.parameters() if p.requires_grad)
print('\tTotal number of model parameters:', total_params)



# Passing to GPU
if model_selection != 'AR':
	lstm_model.cuda()


# Defining data sets and loaders for parallelization option
if training_data_set == 'all':
	training_set = efficient_Dataset(data_indices=range(y_train.shape[0]), x = x_train, t=t_train, y = y_train)
else:
	print('\tNote: imported all data, but only using training data for', training_data_set)
	vol_ID = vol_name_ls.index(training_data_set)
	index_max = vol_cutoff_indices[vol_ID]
	if vol_ID == 0:
		index_min = 0
	else:
		index_min = vol_cutoff_indices[vol_ID - 1]
	curr_data_indices = range(index_min, index_max)
	training_set = efficient_Dataset(data_indices=curr_data_indices, x = x_train, t=t_train, y = y_train)
test_set = efficient_Dataset(data_indices=range(y_test.shape[0]), x = x_test, t = t_test, y = y_test)
train_loader = torch.utils.data.DataLoader(dataset = training_set, batch_size = batch_size, shuffle = True)
test_loader = torch.utils.data.DataLoader(dataset = test_set, batch_size = batch_size, shuffle = True)


# Determining compute options (GPU? Parallel?)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if model_selection != 'AR':
	lstm_model = torch.nn.DataParallel(lstm_model)


# Setting optimization methods
loss = torch.nn.MSELoss()
# Defining many regularization strengths to iterate over to combat overfitting
assert(model_selection in ['AR', 'LSTM', "ConvTimeLSTMUnet"])
if model_selection == 'AR':
	penalization = 1
elif model_selection == 'LSTM':
	penalization = 0.0001
elif model_selection == 'ConvTimeLSTMUnet':
	penalization = 0.1
optimizer = torch.optim.Adam(lstm_model.parameters(), weight_decay = penalization)
# Training loop
print("Beginning training for regularization strength:", penalization)
loss_list = []
#epochs = int(np.ceil((7*10**5) / x_train.shape[0]))
epochs = 100
loop_begin_time = datetime.now()
for i in range(epochs):
	# Marking the beginning time of epoch
	begin_time = datetime.now()
	for data in train_loader:
		
		# data loader
		batch_x, batch_t, batch_y = data
		batch_x.unsqueeze_(2) # code was originally wrote for 5-band level-1 data
		batch_t.unsqueeze_(2) # this unsqueeze is just allowing the explicit acknowledgement
		batch_y.unsqueeze_(2) # that the level-2 data is 1-band deep
		
		# reshaping data if needed for non-spatial LSTMs
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_x = batch_x.permute(0, 3, 4, 1, 2)
			x_sh = batch_x.shape
			batch_x = batch_x.reshape(x_sh[0]*x_sh[1]*x_sh[2], x_sh[3], x_sh[4])
		# Only further processing time in a time-conscious, non-spatial LSTM
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
			batch_t = batch_t.permute(0, 3, 4, 1, 2)
			t_sh = batch_t.shape
			batch_t = batch_t.reshape(t_sh[0]*t_sh[1]*t_sh[2], t_sh[3], t_sh[4])
			# This next line is fragile to the assumption that
			# bands have the same sampling time difference
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_t = batch_t[:,:,[0]]
		
		# move to GPU
		if model_selection != 'AR':
			batch_x = batch_x.to(device)
			# Only move time tensors to GPU if time-conscious LSTM
			if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
				batch_t = batch_t.to(device)
			batch_y = batch_y.to(device)
		
		# Run the model, determining forward pass based on model selected
		if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
			batch_y_hat = lstm_model(batch_x)
		elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_y_hat = lstm_model(batch_x, batch_t)
		elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet']:
			batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
		
		# Extracting the target prediction based on model output
		if model_selection not in ['AR', 'Identity']:
			batch_y_hat = batch_y_hat[0][0]
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_y_hat = batch_y_hat.reshape(x_sh)
			batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
		batch_y_hat = batch_y_hat[:, [-1], :, :, :]
		
		# calculate and store the loss
		batch_loss = loss(batch_y, batch_y_hat)
		loss_list.append(batch_loss.item())
		
		# update parameters
		if model_selection != 'Identity':
			optimizer.zero_grad()
			batch_loss.backward()
			optimizer.step()
		
	# Marking the end time and computing time difference, also printing epoch information
	end_time = datetime.now()
	time_diff = (end_time - loop_begin_time).total_seconds()
	# With l2 regularization iterations, too many prints; only printing every 10 epochs
	if (i + 1) % 10 == 0:
		print('\tEpoch:', i, '\n\t\tMost recent batch loss:', batch_loss.item(), '\n\t\t' + str(time_diff) + ' seconds elapsed')
loop_end_time = datetime.now()
loop_time_diff = (loop_end_time - loop_begin_time).total_seconds()
print('\tTotal training-loop time:', loop_time_diff)


# Converting loss values into array and saving
loss_array = np.asarray(loss_list)
np.save("outputs/loss_over_iterations" + str(penalization) + ".npy", loss_array)


# Trying to free GPU memory
del batch_x
del batch_t
del batch_y
del batch_y_hat
del batch_loss
torch.cuda.empty_cache()


print("Beginning final evaluation")

	# Determine train set performance by volcano
with torch.no_grad():
    count = 0
    for i in range(len(y_train)):
        batch_x = x_train[[i], :, :, :].unsqueeze(2)
        batch_t = t_train[[i], :, :, :].unsqueeze(2)
        batch_y = y_train[[i], :, :, :].unsqueeze(2)
        
        # reshaping data if needed for non-spatial LSTMs
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0]*x_sh[1]*x_sh[2], x_sh[3], x_sh[4])
        # Only further processing time in a time-conscious, non-spatial LSTM
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0]*t_sh[1]*t_sh[2], t_sh[3], t_sh[4])
            # This next line is fragile to the assumption that
            # bands have the same sampling time difference
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_t = batch_t[:,:,[0]]
            
        # move to GPU
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            # Only move time tensors to GPU if time-conscious LSTM
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
            
        # Run the model, determining forward pass based on model selected
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet']:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        
        # Extracting the target prediction based on model output
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        
        # Moving data off GPU now that model has ran
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        
        # Transformating the data to temperature values
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = (batch_y * (train_y_max - train_y_min)) + train_y_min
        batch_y_hat = (batch_y_hat * (train_y_max - train_y_min)) + train_y_min
        
        # Storing all temperature-valued truths and predictions for
        # one root mean squared error calculation to get error in
        # terms of temperature
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim = 0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim = 0)
        count = count + 1
    # calculate and store the loss
    train_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    train_set_loss = torch.sqrt(train_set_loss)
    train_set_loss = train_set_loss.item()
    # print traination set loss
    print("\tTraining set loss:", train_set_loss)
    # Saving the predictions and corresponding truths
    np.save("outputs/train_prediction" + str(penalization) + ".npy", cpu_y_hat_temps.numpy())
    np.save("outputs/train_truth" + str(penalization) + ".npy", cpu_y_temps.numpy())
    
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices[vol_ID - 1]
        index_max = vol_cutoff_indices[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTraining set loss for', vol, ':', vol_loss)
        vol_ID += 1

# Saving the train set loss
np.save("outputs/final_train_loss" + str(penalization) + ".npy", np.asarray(train_set_loss))

# Determine test set performance by volcano
with torch.no_grad():
	count = 0
	for i in range(len(y_test)):
		batch_x = x_test[[i], :, :, :].unsqueeze(2)
		batch_t = t_test[[i], :, :, :].unsqueeze(2)
		batch_y = y_test[[i], :, :, :].unsqueeze(2)
		
		# reshaping data if needed for non-spatial LSTMs
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_x = batch_x.permute(0, 3, 4, 1, 2)
			x_sh = batch_x.shape
			batch_x = batch_x.reshape(x_sh[0]*x_sh[1]*x_sh[2], x_sh[3], x_sh[4])
		# Only further processing time in a time-conscious, non-spatial LSTM
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
			batch_t = batch_t.permute(0, 3, 4, 1, 2)
			t_sh = batch_t.shape
			batch_t = batch_t.reshape(t_sh[0]*t_sh[1]*t_sh[2], t_sh[3], t_sh[4])
			# This next line is fragile to the assumption that
			# bands have the same sampling time difference
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_t = batch_t[:,:,[0]]
			
		# move to GPU
		if model_selection != 'AR':
			batch_x = batch_x.to(device)
			# Only move time tensors to GPU if time-conscious LSTM
			if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
				batch_t = batch_t.to(device)
			batch_y = batch_y.to(device)
			
		# Run the model, determining forward pass based on model selected
		if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
			batch_y_hat = lstm_model(batch_x)
		elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_y_hat = lstm_model(batch_x, batch_t)
		elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet']:
			batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
		
		# Extracting the target prediction based on model output
		if model_selection not in ['AR', 'Identity']:
			batch_y_hat = batch_y_hat[0][0]
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_y_hat = batch_y_hat.reshape(x_sh)
			batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
		batch_y_hat = batch_y_hat[:, [-1], :, :, :]
		
		# Moving data off GPU now that model has ran
		batch_y = batch_y.cpu()
		batch_y_hat = batch_y_hat.cpu()
		
		# Transformating the data to temperature values
		train_y_min = torch.tensor(stored_parameters[0, 0])
		train_y_max = torch.tensor(stored_parameters[1, 0])
		batch_y = (batch_y * (train_y_max - train_y_min)) + train_y_min
		batch_y_hat = (batch_y_hat * (train_y_max - train_y_min)) + train_y_min
		
		# Storing all temperature-valued truths and predictions for
		# one root mean squared error calculation to get error in
		# terms of temperature
		if count == 0:
			cpu_y_temps = batch_y
			cpu_y_hat_temps = batch_y_hat
		else:
			cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim = 0)
			cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim = 0)
		count = count + 1
	# calculate and store the loss
	test_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
	test_set_loss = torch.sqrt(test_set_loss)
	test_set_loss = test_set_loss.item()
	# print test set loss
	print("\tTest set loss:", test_set_loss)
	# Saving the predictions and corresponding truths
	np.save("outputs/test_prediction" + str(penalization) + ".npy", cpu_y_hat_temps.numpy())
	np.save("outputs/test_truth" + str(penalization) + ".npy", cpu_y_temps.numpy())
	
	vol_ID = 0
	for vol in vol_name_ls:
		if vol_ID == 0:
			index_min = 0
		else:
			index_min = vol_cutoff_indices_test[vol_ID - 1]
		index_max = vol_cutoff_indices_test[vol_ID]
		pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
		true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
		vol_loss = loss(pred_vol, true_vol)
		vol_loss = torch.sqrt(vol_loss)
		vol_loss = vol_loss.item()
		print('\t\tTest set loss for', vol, ':', vol_loss)
		vol_ID += 1

# Saving the testation set loss
np.save("outputs/final_test_loss" + str(penalization) + ".npy", np.asarray(test_set_loss))

# Saving the model
torch.save(lstm_model, "outputs/model" + str(penalization) + ".pt")

# Print breakpoint for human reading between regularization strengths
print('\n')

01bed6e6add94c3188e6a02546f7761a57ed269b
<mask token>

<mask token>
parser.add_argument('--model_selection', type=str, help='model')
parser.add_argument('--training_data_set', type=str, help='training_data_set')
parser.add_argument('--n_layers', type=int, help='n_layers')
parser.add_argument('--hidden_dim_ls', nargs='+', type=int, help=
    'hidden_dim_ls')
<mask token>
print('Model selected:', model_selection)
assert model_selection in ['Identity', 'AR', 'LSTM', 'TimeLSTM',
    'TimeAwareLSTM', 'ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM',
    'ConvTimeLSTMUnet']
print('Training data set:', training_data_set)
print('Hidden layers:', hidden_dim_ls)
<mask token>
if model_selection == 'LSTM':
    from models.LSTM import StackedLSTM as LSTM_Model
elif model_selection == 'Identity':
    from models.Identity import Identity2 as LSTM_Model
elif model_selection == 'AR':
    from models.AR import AR as LSTM_Model
elif model_selection == 'TimeLSTM':
    from models.TimeLSTM import StackedTimeLSTM as LSTM_Model
elif model_selection == 'TimeAwareLSTM':
    from models.TimeAwareLSTM import StackedTimeAwareLSTM as LSTM_Model
elif model_selection == 'ConvLSTM':
    from models.ConvLSTM import ConvLSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTM':
    from models.ConvTimeLSTM2 import ConvTime_LSTM2 as LSTM_Model
elif model_selection == 'ConvTimeAwareLSTM':
    from models.ConvTimeAwareLSTM2 import ConvTimeAware_LSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTMUnet':
    from models.ConvTimeLSTMUnet import ConvTime_LSTM2_Unet as LSTM_Model
<mask token>
torch.manual_seed(0)
print('Importing and formatting data')
<mask token>
assert training_data_set in volcanoes or training_data_set == 'all'
assert hidden_dim_ls[0] in [2000, 1028, 128, 64]
if hidden_dim_ls[0] == 2000:
    assert model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
        'TimeAwareLSTM']
    batch_size_dict = {'AR': 84, 'Identity': 84, 'LSTM': 24, 'TimeLSTM': 15,
        'TimeAwareLSTM': 15}
    batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 1028:
    assert model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM']
    batch_size_dict = {'LSTM': 36, 'TimeLSTM': 18, 'TimeAwareLSTM': 30}
    batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 128:
    if hidden_dim_ls[1] == 64:
        assert model_selection in ['ConvLSTM', 'ConvTimeLSTM',
            'ConvTimeAwareLSTM']
        batch_size_dict = {'ConvLSTM': 320, 'ConvTimeLSTM': 216,
            'ConvTimeAwareLSTM': 320}
        batch_size = batch_size_dict[model_selection]
    else:
        assert model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM']
        batch_size_dict = {'LSTM': 360, 'TimeLSTM': 224, 'TimeAwareLSTM': 320}
        batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 64:
    assert model_selection in ['ConvLSTM', 'ConvTimeLSTM',
        'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet']
    batch_size_dict = {'ConvLSTM': 320, 'ConvTimeLSTM': 216,
        'ConvTimeAwareLSTM': 352, 'ConvTimeLSTMUnet': 80}
    batch_size = batch_size_dict[model_selection]
<mask token>
try:
    volcanoes.remove('.ipynb_checkpoints')
except ValueError as e:
    do = 'nothing'
<mask token>
for vol in os.listdir('data'):
    numpy_data_location = 'AST_08_data/' + vol + '/background_subtracted'
    table_data_location = 'data/' + vol + '/good_df.csv'
    volcano_scenes = np.load(numpy_data_location, allow_pickle=True)
    tabular_metadata = pd.read_csv(table_data_location)
    if vol == 'Erebus':
        tabular_metadata['acquisition_datetimes'] = [file.split('_')[2] for
            file in tabular_metadata['nighttime_volcano_files']]
        tabular_metadata = tabular_metadata[tabular_metadata.
            acquisition_datetimes != '00304102013143221']
        tabular_metadata = tabular_metadata[tabular_metadata.
            acquisition_datetimes != '00308162018142109']
        tabular_metadata = tabular_metadata.reset_index()
    train_n = int(np.floor((len(volcano_scenes) - num_input_scenes) *
        train_percent))
    out_n = int(np.floor((len(volcano_scenes) - num_input_scenes) *
        out_samp_perc))
    x_scenes_train = ma.zeros([train_n - 1, num_input_scenes,
        volcano_scenes.shape[1], volcano_scenes.shape[2]])
    x_scenes_test = ma.zeros([out_n, num_input_scenes, volcano_scenes.shape
        [1], volcano_scenes.shape[2]])
    time_differences_train = np.ones(x_scenes_train.shape)
    time_differences_test = np.ones(x_scenes_test.shape)
    y_scenes_train = ma.zeros([train_n - 1, 1, volcano_scenes.shape[1],
        volcano_scenes.shape[2]])
    y_scenes_test = ma.zeros([out_n, 1, volcano_scenes.shape[1],
        volcano_scenes.shape[2]])
    formatted_dates = [datetime.strptime(date, '%Y-%m-%d') for date in
        tabular_metadata['dates']]
    for i in range(num_input_scenes + 1, x_scenes_train.shape[0] +
        x_scenes_test.shape[0] + num_input_scenes + 1):
        if i < train_n + num_input_scenes:
            x_scenes_train[i - num_input_scenes - 1, :, :, :] = volcano_scenes[
                i - num_input_scenes:i, :, :]
            y_scenes_train[i - num_input_scenes - 1, 0, :, :] = volcano_scenes[
                i, :, :]
            if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
                dates_i_plus_1 = formatted_dates[i - num_input_scenes + 1:i + 1
                    ]
                dates_i = formatted_dates[i - num_input_scenes:i]
                for j in range(len(dates_i_plus_1)):
                    time_differences_train[i - num_input_scenes - 1, j] = (
                        dates_i_plus_1[j] - dates_i[j]).days
            else:
                dates_i = formatted_dates[i - num_input_scenes:i]
                dates_i_minus_1 = formatted_dates[i - num_input_scenes - 1:
                    i - 1]
                for j in range(len(dates_i)):
                    time_differences_train[i - num_input_scenes - 1, j] = (
                        dates_i[j] - dates_i_minus_1[j]).days
        else:
            x_scenes_test[i - train_n - num_input_scenes, :, :, :
                ] = volcano_scenes[i - num_input_scenes:i, :, :]
            y_scenes_test[i - train_n - num_input_scenes, 0, :, :
                ] = volcano_scenes[i, :, :]
            if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
                dates_i_plus_1 = formatted_dates[i - num_input_scenes + 1:i + 1
                    ]
                dates_i = formatted_dates[i - num_input_scenes:i]
                for j in range(len(dates_i_plus_1)):
                    time_differences_test[i - num_input_scenes - train_n - 1, j
                        ] = (dates_i_plus_1[j] - dates_i[j]).days
            else:
                dates_i = formatted_dates[i - num_input_scenes:i]
                dates_i_minus_1 = formatted_dates[i - num_input_scenes - 1:
                    i - 1]
                for j in range(len(dates_i)):
                    time_differences_test[i - num_input_scenes - train_n - 1, j
                        ] = (dates_i[j] - dates_i_minus_1[j]).days
    if count == 0:
        x_train = x_scenes_train
        t_train = time_differences_train
        y_train = y_scenes_train
        x_test = x_scenes_test
        t_test = time_differences_test
        y_test = y_scenes_test
    else:
        x_train = ma.append(x_train, x_scenes_train, axis=0)
        t_train = np.append(t_train, time_differences_train, axis=0)
        y_train = ma.append(y_train, y_scenes_train, axis=0)
        x_test = ma.append(x_test, x_scenes_test, axis=0)
        t_test = np.append(t_test, time_differences_test, axis=0)
        y_test = ma.append(y_test, y_scenes_test, axis=0)
    count += 1
    vol_cutoff_indices.append(y_train.shape[0])
    vol_cutoff_indices_test.append(y_test.shape[0])
    vol_name_ls.append(vol)
    print('\timported ' + str(x_scenes_train.shape[0]) +
        ' training scenes from ' + vol)
    print('\t\timported ' + str(x_scenes_test.shape[0]) +
        ' test scenes from ' + vol)
print('Processing data')
for i in range(len(x_train)):
    for j in range(x_train.shape[1]):
        ma = np.ma.masked_invalid(x_train[i, j, :, :])
        if not np.all(x_train[i, j, :, :].mask == False):
            if j == 0:
                if i == 0 or i in vol_cutoff_indices:
                    print('\t\tNo prior scene; nearest neighbor interpolation')
                    scene = x_train[i, j, :, :]
                    X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
                    xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
                    xybad = np.array((X[scene.mask], Y[scene.mask])).T
                    x_train[i, j, :, :][ma.mask == True] = scene[~scene.mask][
                        KDTree(xygood).query(xybad)[1]]
                else:
                    x_train[i, j, :, :][ma.mask == True] = x_train[i - 1, j,
                        :, :][ma.mask == True]
                    t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][
                        ma.mask == True] + t_train[i - 1, j, :, :][ma.mask ==
                        True]
            else:
                x_train[i, j, :, :][ma.mask == True] = x_train[i, j - 1, :, :][
                    ma.mask == True]
                t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][
                    ma.mask == True] + t_train[i, j - 1, :, :][ma.mask == True]
for i in range(len(x_test)):
    for j in range(x_test.shape[1]):
        ma = np.ma.masked_invalid(x_test[i, j, :, :])
        if not np.all(ma.mask == False):
            if j == 0:
                if i == 0 or i in vol_cutoff_indices_test:
                    scene = x_test[i, j, :, :]
                    X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
                    xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
                    xybad = np.array((X[scene.mask], Y[scene.mask])).T
                    x_test[i, j, :, :][ma.mask == True] = scene[~scene.mask][
                        KDTree(xygood).query(xybad)[1]]
                else:
                    x_test[i, j, :, :][ma.mask == True] = x_test[i - 1, j, :, :
                        ][ma.mask == True]
                    t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][
                        ma.mask == True] + t_test[i - 1, j, :, :][ma.mask ==
                        True]
            else:
                x_test[i, j, :, :][ma.mask == True] = x_test[i, j - 1, :, :][
                    ma.mask == True]
                t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma
                    .mask == True] + t_test[i, j - 1, :, :][ma.mask == True]
for i in range(len(y_train)):
    ma = np.ma.masked_invalid(y_train[i, :, :])
    if not np.all(ma.mask == False):
        if i == 0 or i in vol_cutoff_indices:
            y_train[i, :, :][ma.mask == True] = x_train[i, [-1], :, :][ma.
                mask == True]
        else:
            y_train[i, :, :][ma.mask == True] = y_train[i - 1, :, :, :][ma.
                mask == True]
for i in range(len(y_test)):
    ma = np.ma.masked_invalid(y_test[i, :, :, :])
    if not np.all(ma.mask == False):
        if i == 0 or i in vol_cutoff_indices_test:
            y_test[i, :, :, :][ma.mask == True] = x_test[i, [-1], :, :][ma.
                mask == True]
        else:
            y_test[i, :, :, :][ma.mask == True] = y_test[i - 1, :, :, :][ma
                .mask == True]
<mask token>
np.save('outputs/transformation_parameters.npy', stored_parameters)
<mask token>
print('Setting up methods')
<mask token>
print('\tTotal number of model parameters:', total_params)
if model_selection != 'AR':
    lstm_model.cuda()
if training_data_set == 'all':
    training_set = efficient_Dataset(data_indices=range(y_train.shape[0]),
        x=x_train, t=t_train, y=y_train)
else:
    print('\tNote: imported all data, but only using training data for',
        training_data_set)
    vol_ID = vol_name_ls.index(training_data_set)
    index_max = vol_cutoff_indices[vol_ID]
    if vol_ID == 0:
        index_min = 0
    else:
        index_min = vol_cutoff_indices[vol_ID - 1]
    curr_data_indices = range(index_min, index_max)
    training_set = efficient_Dataset(data_indices=curr_data_indices, x=
        x_train, t=t_train, y=y_train)
<mask token>
if model_selection != 'AR':
    lstm_model = torch.nn.DataParallel(lstm_model)
<mask token>
assert model_selection in ['AR', 'LSTM', 'ConvTimeLSTMUnet']
if model_selection == 'AR':
    penalization = 1
elif model_selection == 'LSTM':
    penalization = 0.0001
elif model_selection == 'ConvTimeLSTMUnet':
    penalization = 0.1
<mask token>
print('Beginning training for regularization strength:', penalization)
<mask token>
for i in range(epochs):
    begin_time = datetime.now()
    for data in train_loader:
        batch_x, batch_t, batch_y = data
        batch_x.unsqueeze_(2)
        batch_t.unsqueeze_(2)
        batch_y.unsqueeze_(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_loss = loss(batch_y, batch_y_hat)
        loss_list.append(batch_loss.item())
        if model_selection != 'Identity':
            optimizer.zero_grad()
            batch_loss.backward()
            optimizer.step()
    end_time = datetime.now()
    time_diff = (end_time - loop_begin_time).total_seconds()
    if (i + 1) % 10 == 0:
        print('\tEpoch:', i, '\n\t\tMost recent batch loss:', batch_loss.
            item(), '\n\t\t' + str(time_diff) + ' seconds elapsed')
<mask token>
print('\tTotal training-loop time:', loop_time_diff)
<mask token>
np.save('outputs/loss_over_iterations' + str(penalization) + '.npy', loss_array
    )
del batch_x
del batch_t
del batch_y
del batch_y_hat
del batch_loss
torch.cuda.empty_cache()
print('Beginning final evaluation')
with torch.no_grad():
    count = 0
    for i in range(len(y_train)):
        batch_x = x_train[[i], :, :, :].unsqueeze(2)
        batch_t = t_train[[i], :, :, :].unsqueeze(2)
        batch_y = y_train[[i], :, :, :].unsqueeze(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = batch_y * (train_y_max - train_y_min) + train_y_min
        batch_y_hat = batch_y_hat * (train_y_max - train_y_min) + train_y_min
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim=0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim=0)
        count = count + 1
    train_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    train_set_loss = torch.sqrt(train_set_loss)
    train_set_loss = train_set_loss.item()
    print('\tTraining set loss:', train_set_loss)
    np.save('outputs/train_prediction' + str(penalization) + '.npy',
        cpu_y_hat_temps.numpy())
    np.save('outputs/train_truth' + str(penalization) + '.npy', cpu_y_temps
        .numpy())
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices[vol_ID - 1]
        index_max = vol_cutoff_indices[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTraining set loss for', vol, ':', vol_loss)
        vol_ID += 1
np.save('outputs/final_train_loss' + str(penalization) + '.npy', np.asarray
    (train_set_loss))
with torch.no_grad():
    count = 0
    for i in range(len(y_test)):
        batch_x = x_test[[i], :, :, :].unsqueeze(2)
        batch_t = t_test[[i], :, :, :].unsqueeze(2)
        batch_y = y_test[[i], :, :, :].unsqueeze(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = batch_y * (train_y_max - train_y_min) + train_y_min
        batch_y_hat = batch_y_hat * (train_y_max - train_y_min) + train_y_min
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim=0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim=0)
        count = count + 1
    test_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    test_set_loss = torch.sqrt(test_set_loss)
    test_set_loss = test_set_loss.item()
    print('\tTest set loss:', test_set_loss)
    np.save('outputs/test_prediction' + str(penalization) + '.npy',
        cpu_y_hat_temps.numpy())
    np.save('outputs/test_truth' + str(penalization) + '.npy', cpu_y_temps.
        numpy())
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices_test[vol_ID - 1]
        index_max = vol_cutoff_indices_test[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTest set loss for', vol, ':', vol_loss)
        vol_ID += 1
np.save('outputs/final_test_loss' + str(penalization) + '.npy', np.asarray(
    test_set_loss))
torch.save(lstm_model, 'outputs/model' + str(penalization) + '.pt')
print('\n')

<mask token>
parser = argparse.ArgumentParser(description='level2 training')
parser.add_argument('--model_selection', type=str, help='model')
parser.add_argument('--training_data_set', type=str, help='training_data_set')
parser.add_argument('--n_layers', type=int, help='n_layers')
parser.add_argument('--hidden_dim_ls', nargs='+', type=int, help=
    'hidden_dim_ls')
args = parser.parse_args()
model_selection = args.model_selection
training_data_set = args.training_data_set
n_layers = args.n_layers
hidden_dim_ls = args.hidden_dim_ls
print('Model selected:', model_selection)
assert model_selection in ['Identity', 'AR', 'LSTM', 'TimeLSTM',
    'TimeAwareLSTM', 'ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM',
    'ConvTimeLSTMUnet']
print('Training data set:', training_data_set)
print('Hidden layers:', hidden_dim_ls)
<mask token>
if model_selection == 'LSTM':
    from models.LSTM import StackedLSTM as LSTM_Model
elif model_selection == 'Identity':
    from models.Identity import Identity2 as LSTM_Model
elif model_selection == 'AR':
    from models.AR import AR as LSTM_Model
elif model_selection == 'TimeLSTM':
    from models.TimeLSTM import StackedTimeLSTM as LSTM_Model
elif model_selection == 'TimeAwareLSTM':
    from models.TimeAwareLSTM import StackedTimeAwareLSTM as LSTM_Model
elif model_selection == 'ConvLSTM':
    from models.ConvLSTM import ConvLSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTM':
    from models.ConvTimeLSTM2 import ConvTime_LSTM2 as LSTM_Model
elif model_selection == 'ConvTimeAwareLSTM':
    from models.ConvTimeAwareLSTM2 import ConvTimeAware_LSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTMUnet':
    from models.ConvTimeLSTMUnet import ConvTime_LSTM2_Unet as LSTM_Model
<mask token>
torch.manual_seed(0)
print('Importing and formatting data')
volcanoes = os.listdir('data')
assert training_data_set in volcanoes or training_data_set == 'all'
assert hidden_dim_ls[0] in [2000, 1028, 128, 64]
if hidden_dim_ls[0] == 2000:
    assert model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
        'TimeAwareLSTM']
    batch_size_dict = {'AR': 84, 'Identity': 84, 'LSTM': 24, 'TimeLSTM': 15,
        'TimeAwareLSTM': 15}
    batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 1028:
    assert model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM']
    batch_size_dict = {'LSTM': 36, 'TimeLSTM': 18, 'TimeAwareLSTM': 30}
    batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 128:
    if hidden_dim_ls[1] == 64:
        assert model_selection in ['ConvLSTM', 'ConvTimeLSTM',
            'ConvTimeAwareLSTM']
        batch_size_dict = {'ConvLSTM': 320, 'ConvTimeLSTM': 216,
            'ConvTimeAwareLSTM': 320}
        batch_size = batch_size_dict[model_selection]
    else:
        assert model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM']
        batch_size_dict = {'LSTM': 360, 'TimeLSTM': 224, 'TimeAwareLSTM': 320}
        batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 64:
    assert model_selection in ['ConvLSTM', 'ConvTimeLSTM',
        'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet']
    batch_size_dict = {'ConvLSTM': 320, 'ConvTimeLSTM': 216,
        'ConvTimeAwareLSTM': 352, 'ConvTimeLSTMUnet': 80}
    batch_size = batch_size_dict[model_selection]
lag_dict = {'all': 6, 'Erebus': 3, 'ErtaAle': 9, 'Etna': 3, 'Kilauea': 10,
    'Masaya': 3, 'Nyamuragira': 3, 'Nyiragongo': 3, 'Pacaya': 3, 'Puuoo': 4}
num_input_scenes = lag_dict[training_data_set]
train_percent = 0.85
out_samp_perc = 0.15
try:
    volcanoes.remove('.ipynb_checkpoints')
except ValueError as e:
    do = 'nothing'
count = 0
vol_cutoff_indices = []
vol_cutoff_indices_test = []
vol_name_ls = []
for vol in os.listdir('data'):
    numpy_data_location = 'AST_08_data/' + vol + '/background_subtracted'
    table_data_location = 'data/' + vol + '/good_df.csv'
    volcano_scenes = np.load(numpy_data_location, allow_pickle=True)
    tabular_metadata = pd.read_csv(table_data_location)
    if vol == 'Erebus':
        tabular_metadata['acquisition_datetimes'] = [file.split('_')[2] for
            file in tabular_metadata['nighttime_volcano_files']]
        tabular_metadata = tabular_metadata[tabular_metadata.
            acquisition_datetimes != '00304102013143221']
        tabular_metadata = tabular_metadata[tabular_metadata.
            acquisition_datetimes != '00308162018142109']
        tabular_metadata = tabular_metadata.reset_index()
    train_n = int(np.floor((len(volcano_scenes) - num_input_scenes) *
        train_percent))
    out_n = int(np.floor((len(volcano_scenes) - num_input_scenes) *
        out_samp_perc))
    x_scenes_train = ma.zeros([train_n - 1, num_input_scenes,
        volcano_scenes.shape[1], volcano_scenes.shape[2]])
    x_scenes_test = ma.zeros([out_n, num_input_scenes, volcano_scenes.shape
        [1], volcano_scenes.shape[2]])
    time_differences_train = np.ones(x_scenes_train.shape)
    time_differences_test = np.ones(x_scenes_test.shape)
    y_scenes_train = ma.zeros([train_n - 1, 1, volcano_scenes.shape[1],
        volcano_scenes.shape[2]])
    y_scenes_test = ma.zeros([out_n, 1, volcano_scenes.shape[1],
        volcano_scenes.shape[2]])
    formatted_dates = [datetime.strptime(date, '%Y-%m-%d') for date in
        tabular_metadata['dates']]
    for i in range(num_input_scenes + 1, x_scenes_train.shape[0] +
        x_scenes_test.shape[0] + num_input_scenes + 1):
        if i < train_n + num_input_scenes:
            x_scenes_train[i - num_input_scenes - 1, :, :, :] = volcano_scenes[
                i - num_input_scenes:i, :, :]
            y_scenes_train[i - num_input_scenes - 1, 0, :, :] = volcano_scenes[
                i, :, :]
            if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
                dates_i_plus_1 = formatted_dates[i - num_input_scenes + 1:i + 1
                    ]
                dates_i = formatted_dates[i - num_input_scenes:i]
                for j in range(len(dates_i_plus_1)):
                    time_differences_train[i - num_input_scenes - 1, j] = (
                        dates_i_plus_1[j] - dates_i[j]).days
            else:
                dates_i = formatted_dates[i - num_input_scenes:i]
                dates_i_minus_1 = formatted_dates[i - num_input_scenes - 1:
                    i - 1]
                for j in range(len(dates_i)):
                    time_differences_train[i - num_input_scenes - 1, j] = (
                        dates_i[j] - dates_i_minus_1[j]).days
        else:
            x_scenes_test[i - train_n - num_input_scenes, :, :, :
                ] = volcano_scenes[i - num_input_scenes:i, :, :]
            y_scenes_test[i - train_n - num_input_scenes, 0, :, :
                ] = volcano_scenes[i, :, :]
            if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
                dates_i_plus_1 = formatted_dates[i - num_input_scenes + 1:i + 1
                    ]
                dates_i = formatted_dates[i - num_input_scenes:i]
                for j in range(len(dates_i_plus_1)):
                    time_differences_test[i - num_input_scenes - train_n - 1, j
                        ] = (dates_i_plus_1[j] - dates_i[j]).days
            else:
                dates_i = formatted_dates[i - num_input_scenes:i]
                dates_i_minus_1 = formatted_dates[i - num_input_scenes - 1:
                    i - 1]
                for j in range(len(dates_i)):
                    time_differences_test[i - num_input_scenes - train_n - 1, j
                        ] = (dates_i[j] - dates_i_minus_1[j]).days
    if count == 0:
        x_train = x_scenes_train
        t_train = time_differences_train
        y_train = y_scenes_train
        x_test = x_scenes_test
        t_test = time_differences_test
        y_test = y_scenes_test
    else:
        x_train = ma.append(x_train, x_scenes_train, axis=0)
        t_train = np.append(t_train, time_differences_train, axis=0)
        y_train = ma.append(y_train, y_scenes_train, axis=0)
        x_test = ma.append(x_test, x_scenes_test, axis=0)
        t_test = np.append(t_test, time_differences_test, axis=0)
        y_test = ma.append(y_test, y_scenes_test, axis=0)
    count += 1
    vol_cutoff_indices.append(y_train.shape[0])
    vol_cutoff_indices_test.append(y_test.shape[0])
    vol_name_ls.append(vol)
    print('\timported ' + str(x_scenes_train.shape[0]) +
        ' training scenes from ' + vol)
    print('\t\timported ' + str(x_scenes_test.shape[0]) +
        ' test scenes from ' + vol)
print('Processing data')
for i in range(len(x_train)):
    for j in range(x_train.shape[1]):
        ma = np.ma.masked_invalid(x_train[i, j, :, :])
        if not np.all(x_train[i, j, :, :].mask == False):
            if j == 0:
                if i == 0 or i in vol_cutoff_indices:
                    print('\t\tNo prior scene; nearest neighbor interpolation')
                    scene = x_train[i, j, :, :]
                    X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
                    xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
                    xybad = np.array((X[scene.mask], Y[scene.mask])).T
                    x_train[i, j, :, :][ma.mask == True] = scene[~scene.mask][
                        KDTree(xygood).query(xybad)[1]]
                else:
                    x_train[i, j, :, :][ma.mask == True] = x_train[i - 1, j,
                        :, :][ma.mask == True]
                    t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][
                        ma.mask == True] + t_train[i - 1, j, :, :][ma.mask ==
                        True]
            else:
                x_train[i, j, :, :][ma.mask == True] = x_train[i, j - 1, :, :][
                    ma.mask == True]
                t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][
                    ma.mask == True] + t_train[i, j - 1, :, :][ma.mask == True]
for i in range(len(x_test)):
    for j in range(x_test.shape[1]):
        ma = np.ma.masked_invalid(x_test[i, j, :, :])
        if not np.all(ma.mask == False):
            if j == 0:
                if i == 0 or i in vol_cutoff_indices_test:
                    scene = x_test[i, j, :, :]
                    X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
                    xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
                    xybad = np.array((X[scene.mask], Y[scene.mask])).T
                    x_test[i, j, :, :][ma.mask == True] = scene[~scene.mask][
                        KDTree(xygood).query(xybad)[1]]
                else:
                    x_test[i, j, :, :][ma.mask == True] = x_test[i - 1, j, :, :
                        ][ma.mask == True]
                    t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][
                        ma.mask == True] + t_test[i - 1, j, :, :][ma.mask ==
                        True]
            else:
                x_test[i, j, :, :][ma.mask == True] = x_test[i, j - 1, :, :][
                    ma.mask == True]
                t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma
                    .mask == True] + t_test[i, j - 1, :, :][ma.mask == True]
for i in range(len(y_train)):
    ma = np.ma.masked_invalid(y_train[i, :, :])
    if not np.all(ma.mask == False):
        if i == 0 or i in vol_cutoff_indices:
            y_train[i, :, :][ma.mask == True] = x_train[i, [-1], :, :][ma.
                mask == True]
        else:
            y_train[i, :, :][ma.mask == True] = y_train[i - 1, :, :, :][ma.
                mask == True]
for i in range(len(y_test)):
    ma = np.ma.masked_invalid(y_test[i, :, :, :])
    if not np.all(ma.mask == False):
        if i == 0 or i in vol_cutoff_indices_test:
            y_test[i, :, :, :][ma.mask == True] = x_test[i, [-1], :, :][ma.
                mask == True]
        else:
            y_test[i, :, :, :][ma.mask == True] = y_test[i - 1, :, :, :][ma
                .mask == True]
stored_parameters = np.zeros([2, 3])
x_train, stored_parameters = scale_and_remove_na(x_train, stored_parameters, 0)
x_test = (x_test - stored_parameters[0, 0]) / (stored_parameters[1, 0] -
    stored_parameters[0, 0])
t_train, stored_parameters = scale_and_remove_na(t_train, stored_parameters, 1)
t_test = (t_test - stored_parameters[0, 1]) / (stored_parameters[1, 1] -
    stored_parameters[0, 1])
y_train, stored_parameters = scale_and_remove_na(y_train, stored_parameters, 2)
y_test = (y_test - stored_parameters[0, 2]) / (stored_parameters[1, 2] -
    stored_parameters[0, 2])
np.save('outputs/transformation_parameters.npy', stored_parameters)
x_train = torch.from_numpy(x_train).type(torch.FloatTensor)
x_test = torch.from_numpy(x_test).type(torch.FloatTensor)
t_train = torch.from_numpy(t_train).type(torch.FloatTensor)
t_test = torch.from_numpy(t_test).type(torch.FloatTensor)
y_train = torch.from_numpy(y_train).type(torch.FloatTensor)
y_test = torch.from_numpy(y_test).type(torch.FloatTensor)
print('Setting up methods')
channels = 1
height = x_train.shape[2]
width = x_train.shape[3]
lstm_model = LSTM_Model(input_dim=channels, hidden_dim=hidden_dim_ls, GPU=
    True, input_size=(height, width), num_layers=n_layers)
total_params = sum(p.numel() for p in lstm_model.parameters() if p.
    requires_grad)
print('\tTotal number of model parameters:', total_params)
if model_selection != 'AR':
    lstm_model.cuda()
if training_data_set == 'all':
    training_set = efficient_Dataset(data_indices=range(y_train.shape[0]),
        x=x_train, t=t_train, y=y_train)
else:
    print('\tNote: imported all data, but only using training data for',
        training_data_set)
    vol_ID = vol_name_ls.index(training_data_set)
    index_max = vol_cutoff_indices[vol_ID]
    if vol_ID == 0:
        index_min = 0
    else:
        index_min = vol_cutoff_indices[vol_ID - 1]
    curr_data_indices = range(index_min, index_max)
    training_set = efficient_Dataset(data_indices=curr_data_indices, x=
        x_train, t=t_train, y=y_train)
test_set = efficient_Dataset(data_indices=range(y_test.shape[0]), x=x_test,
    t=t_test, y=y_test)
train_loader = torch.utils.data.DataLoader(dataset=training_set, batch_size
    =batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_set, batch_size=
    batch_size, shuffle=True)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if model_selection != 'AR':
    lstm_model = torch.nn.DataParallel(lstm_model)
loss = torch.nn.MSELoss()
assert model_selection in ['AR', 'LSTM', 'ConvTimeLSTMUnet']
if model_selection == 'AR':
    penalization = 1
elif model_selection == 'LSTM':
    penalization = 0.0001
elif model_selection == 'ConvTimeLSTMUnet':
    penalization = 0.1
optimizer = torch.optim.Adam(lstm_model.parameters(), weight_decay=penalization
    )
print('Beginning training for regularization strength:', penalization)
loss_list = []
epochs = 100
loop_begin_time = datetime.now()
for i in range(epochs):
    begin_time = datetime.now()
    for data in train_loader:
        batch_x, batch_t, batch_y = data
        batch_x.unsqueeze_(2)
        batch_t.unsqueeze_(2)
        batch_y.unsqueeze_(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_loss = loss(batch_y, batch_y_hat)
        loss_list.append(batch_loss.item())
        if model_selection != 'Identity':
            optimizer.zero_grad()
            batch_loss.backward()
            optimizer.step()
    end_time = datetime.now()
    time_diff = (end_time - loop_begin_time).total_seconds()
    if (i + 1) % 10 == 0:
        print('\tEpoch:', i, '\n\t\tMost recent batch loss:', batch_loss.
            item(), '\n\t\t' + str(time_diff) + ' seconds elapsed')
loop_end_time = datetime.now()
loop_time_diff = (loop_end_time - loop_begin_time).total_seconds()
print('\tTotal training-loop time:', loop_time_diff)
loss_array = np.asarray(loss_list)
np.save('outputs/loss_over_iterations' + str(penalization) + '.npy', loss_array
    )
del batch_x
del batch_t
del batch_y
del batch_y_hat
del batch_loss
torch.cuda.empty_cache()
print('Beginning final evaluation')
with torch.no_grad():
    count = 0
    for i in range(len(y_train)):
        batch_x = x_train[[i], :, :, :].unsqueeze(2)
        batch_t = t_train[[i], :, :, :].unsqueeze(2)
        batch_y = y_train[[i], :, :, :].unsqueeze(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = batch_y * (train_y_max - train_y_min) + train_y_min
        batch_y_hat = batch_y_hat * (train_y_max - train_y_min) + train_y_min
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim=0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim=0)
        count = count + 1
    train_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    train_set_loss = torch.sqrt(train_set_loss)
    train_set_loss = train_set_loss.item()
    print('\tTraining set loss:', train_set_loss)
    np.save('outputs/train_prediction' + str(penalization) + '.npy',
        cpu_y_hat_temps.numpy())
    np.save('outputs/train_truth' + str(penalization) + '.npy', cpu_y_temps
        .numpy())
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices[vol_ID - 1]
        index_max = vol_cutoff_indices[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTraining set loss for', vol, ':', vol_loss)
        vol_ID += 1
np.save('outputs/final_train_loss' + str(penalization) + '.npy', np.asarray
    (train_set_loss))
with torch.no_grad():
    count = 0
    for i in range(len(y_test)):
        batch_x = x_test[[i], :, :, :].unsqueeze(2)
        batch_t = t_test[[i], :, :, :].unsqueeze(2)
        batch_y = y_test[[i], :, :, :].unsqueeze(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = batch_y * (train_y_max - train_y_min) + train_y_min
        batch_y_hat = batch_y_hat * (train_y_max - train_y_min) + train_y_min
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim=0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim=0)
        count = count + 1
    test_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    test_set_loss = torch.sqrt(test_set_loss)
    test_set_loss = test_set_loss.item()
    print('\tTest set loss:', test_set_loss)
    np.save('outputs/test_prediction' + str(penalization) + '.npy',
        cpu_y_hat_temps.numpy())
    np.save('outputs/test_truth' + str(penalization) + '.npy', cpu_y_temps.
        numpy())
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices_test[vol_ID - 1]
        index_max = vol_cutoff_indices_test[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTest set loss for', vol, ':', vol_loss)
        vol_ID += 1
np.save('outputs/final_test_loss' + str(penalization) + '.npy', np.asarray(
    test_set_loss))
torch.save(lstm_model, 'outputs/model' + str(penalization) + '.pt')
print('\n')

import argparse
parser = argparse.ArgumentParser(description='level2 training')
parser.add_argument('--model_selection', type=str, help='model')
parser.add_argument('--training_data_set', type=str, help='training_data_set')
parser.add_argument('--n_layers', type=int, help='n_layers')
parser.add_argument('--hidden_dim_ls', nargs='+', type=int, help=
    'hidden_dim_ls')
args = parser.parse_args()
model_selection = args.model_selection
training_data_set = args.training_data_set
n_layers = args.n_layers
hidden_dim_ls = args.hidden_dim_ls
print('Model selected:', model_selection)
assert model_selection in ['Identity', 'AR', 'LSTM', 'TimeLSTM',
    'TimeAwareLSTM', 'ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM',
    'ConvTimeLSTMUnet']
print('Training data set:', training_data_set)
print('Hidden layers:', hidden_dim_ls)
import os
import numpy as np
import pandas as pd
from datetime import datetime
import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.utils import data
if model_selection == 'LSTM':
    from models.LSTM import StackedLSTM as LSTM_Model
elif model_selection == 'Identity':
    from models.Identity import Identity2 as LSTM_Model
elif model_selection == 'AR':
    from models.AR import AR as LSTM_Model
elif model_selection == 'TimeLSTM':
    from models.TimeLSTM import StackedTimeLSTM as LSTM_Model
elif model_selection == 'TimeAwareLSTM':
    from models.TimeAwareLSTM import StackedTimeAwareLSTM as LSTM_Model
elif model_selection == 'ConvLSTM':
    from models.ConvLSTM import ConvLSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTM':
    from models.ConvTimeLSTM2 import ConvTime_LSTM2 as LSTM_Model
elif model_selection == 'ConvTimeAwareLSTM':
    from models.ConvTimeAwareLSTM2 import ConvTimeAware_LSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTMUnet':
    from models.ConvTimeLSTMUnet import ConvTime_LSTM2_Unet as LSTM_Model
from helper_fns.processing import scale_and_remove_na
from helper_fns.efcnt_data_AST08 import efficient_Dataset
from optimization import ssim
import numpy.ma as ma
from scipy.spatial import KDTree
torch.manual_seed(0)
print('Importing and formatting data')
volcanoes = os.listdir('data')
assert training_data_set in volcanoes or training_data_set == 'all'
assert hidden_dim_ls[0] in [2000, 1028, 128, 64]
if hidden_dim_ls[0] == 2000:
    assert model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
        'TimeAwareLSTM']
    batch_size_dict = {'AR': 84, 'Identity': 84, 'LSTM': 24, 'TimeLSTM': 15,
        'TimeAwareLSTM': 15}
    batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 1028:
    assert model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM']
    batch_size_dict = {'LSTM': 36, 'TimeLSTM': 18, 'TimeAwareLSTM': 30}
    batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 128:
    if hidden_dim_ls[1] == 64:
        assert model_selection in ['ConvLSTM', 'ConvTimeLSTM',
            'ConvTimeAwareLSTM']
        batch_size_dict = {'ConvLSTM': 320, 'ConvTimeLSTM': 216,
            'ConvTimeAwareLSTM': 320}
        batch_size = batch_size_dict[model_selection]
    else:
        assert model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM']
        batch_size_dict = {'LSTM': 360, 'TimeLSTM': 224, 'TimeAwareLSTM': 320}
        batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 64:
    assert model_selection in ['ConvLSTM', 'ConvTimeLSTM',
        'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet']
    batch_size_dict = {'ConvLSTM': 320, 'ConvTimeLSTM': 216,
        'ConvTimeAwareLSTM': 352, 'ConvTimeLSTMUnet': 80}
    batch_size = batch_size_dict[model_selection]
lag_dict = {'all': 6, 'Erebus': 3, 'ErtaAle': 9, 'Etna': 3, 'Kilauea': 10,
    'Masaya': 3, 'Nyamuragira': 3, 'Nyiragongo': 3, 'Pacaya': 3, 'Puuoo': 4}
num_input_scenes = lag_dict[training_data_set]
train_percent = 0.85
out_samp_perc = 0.15
try:
    volcanoes.remove('.ipynb_checkpoints')
except ValueError as e:
    do = 'nothing'
count = 0
vol_cutoff_indices = []
vol_cutoff_indices_test = []
vol_name_ls = []
for vol in os.listdir('data'):
    numpy_data_location = 'AST_08_data/' + vol + '/background_subtracted'
    table_data_location = 'data/' + vol + '/good_df.csv'
    volcano_scenes = np.load(numpy_data_location, allow_pickle=True)
    tabular_metadata = pd.read_csv(table_data_location)
    if vol == 'Erebus':
        tabular_metadata['acquisition_datetimes'] = [file.split('_')[2] for
            file in tabular_metadata['nighttime_volcano_files']]
        tabular_metadata = tabular_metadata[tabular_metadata.
            acquisition_datetimes != '00304102013143221']
        tabular_metadata = tabular_metadata[tabular_metadata.
            acquisition_datetimes != '00308162018142109']
        tabular_metadata = tabular_metadata.reset_index()
    train_n = int(np.floor((len(volcano_scenes) - num_input_scenes) *
        train_percent))
    out_n = int(np.floor((len(volcano_scenes) - num_input_scenes) *
        out_samp_perc))
    x_scenes_train = ma.zeros([train_n - 1, num_input_scenes,
        volcano_scenes.shape[1], volcano_scenes.shape[2]])
    x_scenes_test = ma.zeros([out_n, num_input_scenes, volcano_scenes.shape
        [1], volcano_scenes.shape[2]])
    time_differences_train = np.ones(x_scenes_train.shape)
    time_differences_test = np.ones(x_scenes_test.shape)
    y_scenes_train = ma.zeros([train_n - 1, 1, volcano_scenes.shape[1],
        volcano_scenes.shape[2]])
    y_scenes_test = ma.zeros([out_n, 1, volcano_scenes.shape[1],
        volcano_scenes.shape[2]])
    formatted_dates = [datetime.strptime(date, '%Y-%m-%d') for date in
        tabular_metadata['dates']]
    for i in range(num_input_scenes + 1, x_scenes_train.shape[0] +
        x_scenes_test.shape[0] + num_input_scenes + 1):
        if i < train_n + num_input_scenes:
            x_scenes_train[i - num_input_scenes - 1, :, :, :] = volcano_scenes[
                i - num_input_scenes:i, :, :]
            y_scenes_train[i - num_input_scenes - 1, 0, :, :] = volcano_scenes[
                i, :, :]
            if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
                dates_i_plus_1 = formatted_dates[i - num_input_scenes + 1:i + 1
                    ]
                dates_i = formatted_dates[i - num_input_scenes:i]
                for j in range(len(dates_i_plus_1)):
                    time_differences_train[i - num_input_scenes - 1, j] = (
                        dates_i_plus_1[j] - dates_i[j]).days
            else:
                dates_i = formatted_dates[i - num_input_scenes:i]
                dates_i_minus_1 = formatted_dates[i - num_input_scenes - 1:
                    i - 1]
                for j in range(len(dates_i)):
                    time_differences_train[i - num_input_scenes - 1, j] = (
                        dates_i[j] - dates_i_minus_1[j]).days
        else:
            x_scenes_test[i - train_n - num_input_scenes, :, :, :
                ] = volcano_scenes[i - num_input_scenes:i, :, :]
            y_scenes_test[i - train_n - num_input_scenes, 0, :, :
                ] = volcano_scenes[i, :, :]
            if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
                dates_i_plus_1 = formatted_dates[i - num_input_scenes + 1:i + 1
                    ]
                dates_i = formatted_dates[i - num_input_scenes:i]
                for j in range(len(dates_i_plus_1)):
                    time_differences_test[i - num_input_scenes - train_n - 1, j
                        ] = (dates_i_plus_1[j] - dates_i[j]).days
            else:
                dates_i = formatted_dates[i - num_input_scenes:i]
                dates_i_minus_1 = formatted_dates[i - num_input_scenes - 1:
                    i - 1]
                for j in range(len(dates_i)):
                    time_differences_test[i - num_input_scenes - train_n - 1, j
                        ] = (dates_i[j] - dates_i_minus_1[j]).days
    if count == 0:
        x_train = x_scenes_train
        t_train = time_differences_train
        y_train = y_scenes_train
        x_test = x_scenes_test
        t_test = time_differences_test
        y_test = y_scenes_test
    else:
        x_train = ma.append(x_train, x_scenes_train, axis=0)
        t_train = np.append(t_train, time_differences_train, axis=0)
        y_train = ma.append(y_train, y_scenes_train, axis=0)
        x_test = ma.append(x_test, x_scenes_test, axis=0)
        t_test = np.append(t_test, time_differences_test, axis=0)
        y_test = ma.append(y_test, y_scenes_test, axis=0)
    count += 1
    vol_cutoff_indices.append(y_train.shape[0])
    vol_cutoff_indices_test.append(y_test.shape[0])
    vol_name_ls.append(vol)
    print('\timported ' + str(x_scenes_train.shape[0]) +
        ' training scenes from ' + vol)
    print('\t\timported ' + str(x_scenes_test.shape[0]) +
        ' test scenes from ' + vol)
print('Processing data')
for i in range(len(x_train)):
    for j in range(x_train.shape[1]):
        ma = np.ma.masked_invalid(x_train[i, j, :, :])
        if not np.all(x_train[i, j, :, :].mask == False):
            if j == 0:
                if i == 0 or i in vol_cutoff_indices:
                    print('\t\tNo prior scene; nearest neighbor interpolation')
                    scene = x_train[i, j, :, :]
                    X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
                    xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
                    xybad = np.array((X[scene.mask], Y[scene.mask])).T
                    x_train[i, j, :, :][ma.mask == True] = scene[~scene.mask][
                        KDTree(xygood).query(xybad)[1]]
                else:
                    x_train[i, j, :, :][ma.mask == True] = x_train[i - 1, j,
                        :, :][ma.mask == True]
                    t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][
                        ma.mask == True] + t_train[i - 1, j, :, :][ma.mask ==
                        True]
            else:
                x_train[i, j, :, :][ma.mask == True] = x_train[i, j - 1, :, :][
                    ma.mask == True]
                t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][
                    ma.mask == True] + t_train[i, j - 1, :, :][ma.mask == True]
for i in range(len(x_test)):
    for j in range(x_test.shape[1]):
        ma = np.ma.masked_invalid(x_test[i, j, :, :])
        if not np.all(ma.mask == False):
            if j == 0:
                if i == 0 or i in vol_cutoff_indices_test:
                    scene = x_test[i, j, :, :]
                    X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
                    xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
                    xybad = np.array((X[scene.mask], Y[scene.mask])).T
                    x_test[i, j, :, :][ma.mask == True] = scene[~scene.mask][
                        KDTree(xygood).query(xybad)[1]]
                else:
                    x_test[i, j, :, :][ma.mask == True] = x_test[i - 1, j, :, :
                        ][ma.mask == True]
                    t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][
                        ma.mask == True] + t_test[i - 1, j, :, :][ma.mask ==
                        True]
            else:
                x_test[i, j, :, :][ma.mask == True] = x_test[i, j - 1, :, :][
                    ma.mask == True]
                t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma
                    .mask == True] + t_test[i, j - 1, :, :][ma.mask == True]
for i in range(len(y_train)):
    ma = np.ma.masked_invalid(y_train[i, :, :])
    if not np.all(ma.mask == False):
        if i == 0 or i in vol_cutoff_indices:
            y_train[i, :, :][ma.mask == True] = x_train[i, [-1], :, :][ma.
                mask == True]
        else:
            y_train[i, :, :][ma.mask == True] = y_train[i - 1, :, :, :][ma.
                mask == True]
for i in range(len(y_test)):
    ma = np.ma.masked_invalid(y_test[i, :, :, :])
    if not np.all(ma.mask == False):
        if i == 0 or i in vol_cutoff_indices_test:
            y_test[i, :, :, :][ma.mask == True] = x_test[i, [-1], :, :][ma.
                mask == True]
        else:
            y_test[i, :, :, :][ma.mask == True] = y_test[i - 1, :, :, :][ma
                .mask == True]
stored_parameters = np.zeros([2, 3])
x_train, stored_parameters = scale_and_remove_na(x_train, stored_parameters, 0)
x_test = (x_test - stored_parameters[0, 0]) / (stored_parameters[1, 0] -
    stored_parameters[0, 0])
t_train, stored_parameters = scale_and_remove_na(t_train, stored_parameters, 1)
t_test = (t_test - stored_parameters[0, 1]) / (stored_parameters[1, 1] -
    stored_parameters[0, 1])
y_train, stored_parameters = scale_and_remove_na(y_train, stored_parameters, 2)
y_test = (y_test - stored_parameters[0, 2]) / (stored_parameters[1, 2] -
    stored_parameters[0, 2])
np.save('outputs/transformation_parameters.npy', stored_parameters)
x_train = torch.from_numpy(x_train).type(torch.FloatTensor)
x_test = torch.from_numpy(x_test).type(torch.FloatTensor)
t_train = torch.from_numpy(t_train).type(torch.FloatTensor)
t_test = torch.from_numpy(t_test).type(torch.FloatTensor)
y_train = torch.from_numpy(y_train).type(torch.FloatTensor)
y_test = torch.from_numpy(y_test).type(torch.FloatTensor)
print('Setting up methods')
channels = 1
height = x_train.shape[2]
width = x_train.shape[3]
lstm_model = LSTM_Model(input_dim=channels, hidden_dim=hidden_dim_ls, GPU=
    True, input_size=(height, width), num_layers=n_layers)
total_params = sum(p.numel() for p in lstm_model.parameters() if p.
    requires_grad)
print('\tTotal number of model parameters:', total_params)
if model_selection != 'AR':
    lstm_model.cuda()
if training_data_set == 'all':
    training_set = efficient_Dataset(data_indices=range(y_train.shape[0]),
        x=x_train, t=t_train, y=y_train)
else:
    print('\tNote: imported all data, but only using training data for',
        training_data_set)
    vol_ID = vol_name_ls.index(training_data_set)
    index_max = vol_cutoff_indices[vol_ID]
    if vol_ID == 0:
        index_min = 0
    else:
        index_min = vol_cutoff_indices[vol_ID - 1]
    curr_data_indices = range(index_min, index_max)
    training_set = efficient_Dataset(data_indices=curr_data_indices, x=
        x_train, t=t_train, y=y_train)
test_set = efficient_Dataset(data_indices=range(y_test.shape[0]), x=x_test,
    t=t_test, y=y_test)
train_loader = torch.utils.data.DataLoader(dataset=training_set, batch_size
    =batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_set, batch_size=
    batch_size, shuffle=True)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if model_selection != 'AR':
    lstm_model = torch.nn.DataParallel(lstm_model)
loss = torch.nn.MSELoss()
assert model_selection in ['AR', 'LSTM', 'ConvTimeLSTMUnet']
if model_selection == 'AR':
    penalization = 1
elif model_selection == 'LSTM':
    penalization = 0.0001
elif model_selection == 'ConvTimeLSTMUnet':
    penalization = 0.1
optimizer = torch.optim.Adam(lstm_model.parameters(), weight_decay=penalization
    )
print('Beginning training for regularization strength:', penalization)
loss_list = []
epochs = 100
loop_begin_time = datetime.now()
for i in range(epochs):
    begin_time = datetime.now()
    for data in train_loader:
        batch_x, batch_t, batch_y = data
        batch_x.unsqueeze_(2)
        batch_t.unsqueeze_(2)
        batch_y.unsqueeze_(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_loss = loss(batch_y, batch_y_hat)
        loss_list.append(batch_loss.item())
        if model_selection != 'Identity':
            optimizer.zero_grad()
            batch_loss.backward()
            optimizer.step()
    end_time = datetime.now()
    time_diff = (end_time - loop_begin_time).total_seconds()
    if (i + 1) % 10 == 0:
        print('\tEpoch:', i, '\n\t\tMost recent batch loss:', batch_loss.
            item(), '\n\t\t' + str(time_diff) + ' seconds elapsed')
loop_end_time = datetime.now()
loop_time_diff = (loop_end_time - loop_begin_time).total_seconds()
print('\tTotal training-loop time:', loop_time_diff)
loss_array = np.asarray(loss_list)
np.save('outputs/loss_over_iterations' + str(penalization) + '.npy', loss_array
    )
del batch_x
del batch_t
del batch_y
del batch_y_hat
del batch_loss
torch.cuda.empty_cache()
print('Beginning final evaluation')
with torch.no_grad():
    count = 0
    for i in range(len(y_train)):
        batch_x = x_train[[i], :, :, :].unsqueeze(2)
        batch_t = t_train[[i], :, :, :].unsqueeze(2)
        batch_y = y_train[[i], :, :, :].unsqueeze(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = batch_y * (train_y_max - train_y_min) + train_y_min
        batch_y_hat = batch_y_hat * (train_y_max - train_y_min) + train_y_min
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim=0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim=0)
        count = count + 1
    train_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    train_set_loss = torch.sqrt(train_set_loss)
    train_set_loss = train_set_loss.item()
    print('\tTraining set loss:', train_set_loss)
    np.save('outputs/train_prediction' + str(penalization) + '.npy',
        cpu_y_hat_temps.numpy())
    np.save('outputs/train_truth' + str(penalization) + '.npy', cpu_y_temps
        .numpy())
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices[vol_ID - 1]
        index_max = vol_cutoff_indices[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTraining set loss for', vol, ':', vol_loss)
        vol_ID += 1
np.save('outputs/final_train_loss' + str(penalization) + '.npy', np.asarray
    (train_set_loss))
with torch.no_grad():
    count = 0
    for i in range(len(y_test)):
        batch_x = x_test[[i], :, :, :].unsqueeze(2)
        batch_t = t_test[[i], :, :, :].unsqueeze(2)
        batch_y = y_test[[i], :, :, :].unsqueeze(2)
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0] * x_sh[1] * x_sh[2], x_sh[3],
                x_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0] * t_sh[1] * t_sh[2], t_sh[3],
                t_sh[4])
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM'
            ]:
            batch_t = batch_t[:, :, [0]]
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet'
            ]:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM',
            'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = batch_y * (train_y_max - train_y_min) + train_y_min
        batch_y_hat = batch_y_hat * (train_y_max - train_y_min) + train_y_min
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim=0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim=0)
        count = count + 1
    test_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    test_set_loss = torch.sqrt(test_set_loss)
    test_set_loss = test_set_loss.item()
    print('\tTest set loss:', test_set_loss)
    np.save('outputs/test_prediction' + str(penalization) + '.npy',
        cpu_y_hat_temps.numpy())
    np.save('outputs/test_truth' + str(penalization) + '.npy', cpu_y_temps.
        numpy())
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices_test[vol_ID - 1]
        index_max = vol_cutoff_indices_test[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTest set loss for', vol, ':', vol_loss)
        vol_ID += 1
np.save('outputs/final_test_loss' + str(penalization) + '.npy', np.asarray(
    test_set_loss))
torch.save(lstm_model, 'outputs/model' + str(penalization) + '.pt')
print('\n')

 



# High level summary:
# The goal of the script is to train one of many neural networks
# on previously process ASTER imagery.

# Input data:
# This input data is expected to be in the form of numpy arrays
# representing sequences of cropped thermal infrared images from
# 1999-2012. More specifically, this is anticipated to be temperature
# above the environmental background (kelvin/celsius) for volcanoes.
# This input data is gathered from the previously numbered scripts in
# the directory.

# Script arguments: (model selection and training data)
# Model selection:
# This training code has been designed with several different
# neural network architectures in mind. Namely,
#     1) "LSTM"
#     2 LSTM variants that handle irregular time intervals
#         * "Time-LSTM"
#         * "TimeAwareLSTM"
#     3) Convolutional LSTM (that allows images sequences)
#         * "ConvLSTM"
#     4) New/proposed models that allow image sequences that are
#        irregular through time
#         * Convolutional LSTM + Time-LSTM
#             - "ConvTimeLSTM"
#         * Convolutional LSTM + TimeAwareLSTM
#             - "ConvTimeAwareLSTM"
# The quoted keyword-names are the required input for this argument and will
# select which neural network is trained; this in turn affects preprocessing
# the data into different forms, the training batch size, etc...
#
# Training data:
# This training code has been designed with two sets of training data in mind.
# These two options are:
#     1) Single volcano; in this case, the argument should be the directory
#        name containing the volcano-specific data. All of these directories
#        are expected to be within the "data" directory which is located in
#        current working directory.
#     2) All volcanoes; in this case the argument should be "all".

# Hardware:
# This code was developed for Penn State's Advanced CyberInfrastructure's
# GPU instances which utilize dual NVIDIA Tesla K80 GPU cards resulting
# in 4 GPUs with 11 GB of RAM each
# Later adapted for more powerful hardware...
# 10 GPUs (Tesla V100s-PCI), 32 GB of RAM each


import argparse

parser = argparse.ArgumentParser(description='level2 training')
parser.add_argument('--model_selection', type=str, help='model')
parser.add_argument('--training_data_set', type=str, help='training_data_set')
parser.add_argument('--n_layers', type=int, help='n_layers')
parser.add_argument('--hidden_dim_ls', nargs='+', type=int, help='hidden_dim_ls')
args = parser.parse_args()

model_selection = args.model_selection
training_data_set = args.training_data_set
n_layers = args.n_layers
hidden_dim_ls = args.hidden_dim_ls


# Gathering function inputs
### model_selection = input("Which model do you select?")
print('Model selected:', model_selection)
assert(model_selection in ['Identity', 'AR', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM', 'ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet'])
### training_data_set = input("Which set of training data do you want to use?")
print('Training data set:', training_data_set)
### n_layers = int(input("Enter number of layers: ")) 
### hidden_dim_ls = []
# iterating till the range 
### for i in range(n_layers): 
###     layer_dim = int(input()) 
###     hidden_dim_ls.append(layer_dim)
print('Hidden layers:', hidden_dim_ls)


# Libraries and imports
import os
import numpy as np
import pandas as pd
from datetime import datetime
import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.utils import data
if model_selection == 'LSTM':
	from models.LSTM import StackedLSTM as LSTM_Model
elif model_selection == 'Identity':
	from models.Identity import Identity2 as LSTM_Model
elif model_selection == 'AR':
	from models.AR import AR as LSTM_Model
elif model_selection == 'TimeLSTM':
	from models.TimeLSTM import StackedTimeLSTM as LSTM_Model
elif model_selection == 'TimeAwareLSTM':
	from models.TimeAwareLSTM import StackedTimeAwareLSTM as LSTM_Model
elif model_selection == 'ConvLSTM':
	from models.ConvLSTM import ConvLSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTM':
	from models.ConvTimeLSTM2 import ConvTime_LSTM2 as LSTM_Model
elif model_selection == 'ConvTimeAwareLSTM':
	from models.ConvTimeAwareLSTM2 import ConvTimeAware_LSTM as LSTM_Model
elif model_selection == 'ConvTimeLSTMUnet':
	from models.ConvTimeLSTMUnet import ConvTime_LSTM2_Unet as LSTM_Model
from helper_fns.processing import scale_and_remove_na
from helper_fns.efcnt_data_AST08 import efficient_Dataset
from optimization import ssim
import numpy.ma as ma
from scipy.spatial import KDTree


# Setting seed for reproducibility
torch.manual_seed(0)


# Basic data import, step 0
print("Importing and formatting data")
volcanoes = os.listdir("data")
assert((training_data_set in volcanoes) or (training_data_set == "all"))


# Training parameters
# This needs to actually be variable, will do with later exploration
assert(hidden_dim_ls[0] in [2000, 1028, 128, 64])
if hidden_dim_ls[0] == 2000: # [2000, 2000, 1]
	assert(model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM'])
	# batch_size_dict = {'AR':84, 'Identity':84, 'LSTM':4, 'TimeLSTM':2, 'TimeAwareLSTM':2}
	batch_size_dict = {'AR':84, 'Identity':84, 'LSTM':24, 'TimeLSTM':15, 'TimeAwareLSTM':15}
	batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 1028: # [1028, 512, 1]
	assert(model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM'])
	# batch_size_dict = {'LSTM':12, 'TimeLSTM':6, 'TimeAwareLSTM':10}
	batch_size_dict = {'LSTM':36, 'TimeLSTM':18, 'TimeAwareLSTM':30}
	batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 128: 
	if hidden_dim_ls[1] == 64: # [128, 64, 64, 1] for ConvLSTMs
		assert(model_selection in ['ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM'])
		batch_size_dict = {'ConvLSTM':320, 'ConvTimeLSTM':216, 'ConvTimeAwareLSTM':320}
		batch_size = batch_size_dict[model_selection]
	else: # [128, 1] for LSTMs
		assert(model_selection in ['LSTM', 'TimeLSTM', 'TimeAwareLSTM'])
		# batch_size_dict = {'LSTM':128, 'TimeLSTM':80, 'TimeAwareLSTM':114}
		batch_size_dict = {'LSTM':360, 'TimeLSTM':224, 'TimeAwareLSTM':320}
		batch_size = batch_size_dict[model_selection]
if hidden_dim_ls[0] == 64: # [64, 64, 1]; [64, 128, 64, 1] for Unet
	assert(model_selection in ['ConvLSTM', 'ConvTimeLSTM', 'ConvTimeAwareLSTM', 'ConvTimeLSTMUnet'])
	# batch_size_dict = {'ConvLSTM':108, 'ConvTimeLSTM':76, 'ConvTimeAwareLSTM':124, 'ConvTimeLSTMUnet':36}
	batch_size_dict = {'ConvLSTM':320, 'ConvTimeLSTM':216, 'ConvTimeAwareLSTM':352, 'ConvTimeLSTMUnet':80}
	batch_size = batch_size_dict[model_selection]
lag_dict = {"all":6, "Erebus":3, "ErtaAle":9, "Etna":3, "Kilauea":10, "Masaya":3, "Nyamuragira":3, "Nyiragongo":3, "Pacaya":3, "Puuoo":4}

num_input_scenes = lag_dict[training_data_set]
train_percent = 0.85
out_samp_perc = 0.15


# Removing possible directory clutter
try:
	volcanoes.remove(".ipynb_checkpoints")
except ValueError as e:
	do = 'nothing'

count = 0
vol_cutoff_indices = []
vol_cutoff_indices_test = []
vol_name_ls = []
for vol in os.listdir('data'):
	### Basic data import ###
	numpy_data_location = "AST_08_data/" + vol + "/background_subtracted"
	table_data_location = "data/" + vol + "/good_df.csv"
	volcano_scenes = np.load(numpy_data_location, allow_pickle = True)
	tabular_metadata = pd.read_csv(table_data_location)
	### Taking care of the unretrievable scenes for Erebus ###
	### This should always be done when referencing processed data (AST_08_data) with metadata (data) ###
	if vol == 'Erebus':
		tabular_metadata['acquisition_datetimes'] = [file.split('_')[2] for file in tabular_metadata['nighttime_volcano_files']]
		tabular_metadata = tabular_metadata[tabular_metadata.acquisition_datetimes != '00304102013143221'] 
		tabular_metadata = tabular_metadata[tabular_metadata.acquisition_datetimes != '00308162018142109']
		tabular_metadata = tabular_metadata.reset_index()
	### Separate model inputs and outputs
	# Determine number in each partition
	train_n = int(np.floor((len(volcano_scenes) - num_input_scenes)*train_percent))
	out_n = int(np.floor((len(volcano_scenes) - num_input_scenes)*out_samp_perc))
	# For every data partition
	# Array for the prior scenes
	#   "train_n - 1" is to remove the first scene that wont have an associated TimeAwareLSTM time interval
	x_scenes_train = ma.zeros([train_n - 1, num_input_scenes, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	x_scenes_test = ma.zeros([out_n, num_input_scenes, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	# Array for the time differences between scenes
	time_differences_train = np.ones(x_scenes_train.shape)
	time_differences_test = np.ones(x_scenes_test.shape)
	# Array for the target scenes
	y_scenes_train = ma.zeros([train_n - 1, 1, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	y_scenes_test = ma.zeros([out_n, 1, volcano_scenes.shape[1], volcano_scenes.shape[2]])
	# Formatting the string dates as datetime objects
	formatted_dates = [datetime.strptime(date, '%Y-%m-%d') for date in tabular_metadata['dates']]
	# For all observations - acknowledging that the first (n-1) wont have n prior observations
	#     Also, the first data point wont have a TimeAwareLSTM time value, so it is omitted
	for i in range(num_input_scenes + 1, x_scenes_train.shape[0] + x_scenes_test.shape[0] + num_input_scenes+1):
		if i < (train_n + num_input_scenes):
			# Store the image data
			x_scenes_train[i - num_input_scenes - 1, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :]
			y_scenes_train[i - num_input_scenes - 1, 0, :, :] = volcano_scenes[i, :, :]
			# Compute the time differences and store
			# Time LSTM uses forward-time interval
			if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
				dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]
				dates_i = formatted_dates[(i - num_input_scenes):i]
				for j in range(len(dates_i_plus_1)):
					time_differences_train[i - num_input_scenes - 1, j] = (dates_i_plus_1[j] - dates_i[j]).days
			# While TimeAwareLSTM uses backwards-time interval
			else:
				dates_i = formatted_dates[(i - num_input_scenes):i]
				dates_i_minus_1 = formatted_dates[(i - num_input_scenes - 1):(i - 1)]
				for j in range(len(dates_i)):
					time_differences_train[i - num_input_scenes - 1, j] = (dates_i[j] - dates_i_minus_1[j]).days
		else:
			# Store the image data
			x_scenes_test[i - train_n - num_input_scenes, :, :, :] = volcano_scenes[(i - num_input_scenes):i, :, :]
			y_scenes_test[i - train_n - num_input_scenes, 0, :, :] = volcano_scenes[i, :, :]
			# Compute the time differences and store
			# Time LSTM uses forward-time interval
			if model_selection in ['TimeLSTM', 'ConvTimeLSTM']:
				dates_i_plus_1 = formatted_dates[(i - num_input_scenes + 1):(i + 1)]
				dates_i = formatted_dates[(i - num_input_scenes):i]
				for j in range(len(dates_i_plus_1)):
					time_differences_test[i - num_input_scenes - train_n - 1, j] = (dates_i_plus_1[j] - dates_i[j]).days
			# While TimeAwareLSTM uses backwards-time interval
			else:
				dates_i = formatted_dates[(i - num_input_scenes):i]
				dates_i_minus_1 = formatted_dates[(i - num_input_scenes - 1):(i - 1)]
				for j in range(len(dates_i)):
					time_differences_test[i - num_input_scenes - train_n - 1, j] = (dates_i[j] - dates_i_minus_1[j]).days
	if count == 0:
		x_train = x_scenes_train
		t_train = time_differences_train
		y_train = y_scenes_train
		x_test = x_scenes_test
		t_test = time_differences_test
		y_test = y_scenes_test
	else:
		x_train = ma.append(x_train, x_scenes_train, axis = 0)
		t_train = np.append(t_train, time_differences_train, axis = 0)
		y_train = ma.append(y_train, y_scenes_train, axis = 0)
		x_test = ma.append(x_test, x_scenes_test, axis = 0)
		t_test = np.append(t_test, time_differences_test, axis = 0)
		y_test = ma.append(y_test, y_scenes_test, axis = 0)
	count += 1
	vol_cutoff_indices.append(y_train.shape[0])
	vol_cutoff_indices_test.append(y_test.shape[0])
	vol_name_ls.append(vol)
	print('\timported ' + str(x_scenes_train.shape[0]) + ' training scenes from ' + vol)
	print('\t\timported ' + str(x_scenes_test.shape[0]) + ' test scenes from ' + vol)


print("Processing data")
# Gap fill missing values with previous ones or nearest neighbor
for i in range(len(x_train)):
	for j in range(x_train.shape[1]):
		# Identifying missing values
		ma = np.ma.masked_invalid(x_train[i, j, :, :])
		# If the mask found NA values
		if not np.all(x_train[i, j, :, :].mask == False):
			# Using previous value to fill
			if j == 0:
				# Unless there's no previous, then nearest neighbor interpolate
				if i == 0 or i in vol_cutoff_indices:
					print('\t\tNo prior scene; nearest neighbor interpolation')
					scene = x_train[i, j, :, :]
					# Explicitly retrieving good and bad locations from the mask
					# note that these differ from (x, y) which earlier defined bad pixel locations
					X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
					xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
					xybad = np.array((X[scene.mask], Y[scene.mask])).T
					# Performing the nearest neighbor gap-filling
					x_train[i, j, :, :][ma.mask == True] = scene[~scene.mask][KDTree(xygood).query(xybad)[1]]
				else:
					x_train[i, j, :, :][ma.mask == True] = x_train[i-1, j, :, :][ma.mask == True]
					t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][ma.mask == True] + t_train[i-1, j, :, :][ma.mask == True]
			else:
				x_train[i, j, :, :][ma.mask == True] = x_train[i, j-1, :, :][ma.mask == True]
				t_train[i, j, :, :][ma.mask == True] = t_train[i, j, :, :][ma.mask == True] + t_train[i, j-1, :, :][ma.mask == True]
for i in range(len(x_test)):
	for j in range(x_test.shape[1]):
		# Identifying missing values
		ma = np.ma.masked_invalid(x_test[i, j, :, :])
		# If the mask found NA values
		if not np.all(ma.mask == False):
			# Using previous value to fill
			if j == 0:
				# Unless there's no previous, then nearest neighbor interpolate
				if i == 0 or i in vol_cutoff_indices_test:
					scene = x_test[i, j, :, :]
					# Explicitly retrieving good and bad locations from the mask
					# note that these differ from (x, y) which earlier defined bad pixel locations
					X, Y = np.mgrid[0:scene.shape[0], 0:scene.shape[1]]
					xygood = np.array((X[~scene.mask], Y[~scene.mask])).T
					xybad = np.array((X[scene.mask], Y[scene.mask])).T
					# Performing the nearest neighbor gap-filling
					x_test[i, j, :, :][ma.mask == True] = scene[~scene.mask][KDTree(xygood).query(xybad)[1]]
				else:
					x_test[i, j, :, :][ma.mask == True] = x_test[i-1, j, :, :][ma.mask == True]
					t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma.mask == True] + t_test[i-1, j, :, :][ma.mask == True]
			else:
				x_test[i, j, :, :][ma.mask == True] = x_test[i, j-1, :, :][ma.mask == True]
				t_test[i, j, :, :][ma.mask == True] = t_test[i, j, :, :][ma.mask == True] + t_test[i, j-1, :, :][ma.mask == True]
for i in range(len(y_train)):
	ma = np.ma.masked_invalid(y_train[i, :, :])
	# If the mask found NA values
	if not np.all(ma.mask == False):
		if i == 0 or i in vol_cutoff_indices: 
			y_train[i, :, :][ma.mask == True] = x_train[i, [-1], :, :][ma.mask == True]
		else:
			y_train[i, :, :][ma.mask == True] = y_train[i-1, :, :, :][ma.mask == True]
for i in range(len(y_test)):
	ma = np.ma.masked_invalid(y_test[i, :, :, :])
	# If the mask found NA values
	if not np.all(ma.mask == False):
		if i == 0 or i in vol_cutoff_indices_test:
			y_test[i, :, :, :][ma.mask == True] = x_test[i, [-1], :, :][ma.mask == True]
		else:
			y_test[i, :, :, :][ma.mask == True] = y_test[i-1, :, :, :][ma.mask == True]


# Scale 0-1 using min and max from training set
# Also attempts to find NAs and replace with 0s, but those shouldnt exist anymore
stored_parameters = np.zeros([2, 3])
x_train, stored_parameters = scale_and_remove_na(x_train, stored_parameters, 0)
x_test = (x_test - stored_parameters[0, 0]) / (stored_parameters[1, 0] - stored_parameters[0, 0])
t_train, stored_parameters = scale_and_remove_na(t_train, stored_parameters, 1)
t_test = (t_test - stored_parameters[0, 1]) / (stored_parameters[1, 1] - stored_parameters[0, 1])
y_train, stored_parameters = scale_and_remove_na(y_train, stored_parameters, 2)
y_test = (y_test - stored_parameters[0, 2]) / (stored_parameters[1, 2] - stored_parameters[0, 2])
np.save("outputs/transformation_parameters.npy", stored_parameters)


# Convert to torch tensors
x_train = torch.from_numpy(x_train).type(torch.FloatTensor)
x_test = torch.from_numpy(x_test).type(torch.FloatTensor)
t_train = torch.from_numpy(t_train).type(torch.FloatTensor)
t_test = torch.from_numpy(t_test).type(torch.FloatTensor)
y_train = torch.from_numpy(y_train).type(torch.FloatTensor)
y_test = torch.from_numpy(y_test).type(torch.FloatTensor)


# Defining model parameters
# Picking one of the like-sequence tensors within the list to set parameters
print("Setting up methods")
channels = 1
height = x_train.shape[2]
width = x_train.shape[3]
lstm_model = LSTM_Model(input_dim=channels,hidden_dim=hidden_dim_ls,GPU=True,input_size=(height,width),num_layers=n_layers)
# Print number of model parameters
total_params = sum(p.numel() for p in lstm_model.parameters() if p.requires_grad)
print('\tTotal number of model parameters:', total_params)



# Passing to GPU
if model_selection != 'AR':
	lstm_model.cuda()


# Defining data sets and loaders for parallelization option
if training_data_set == 'all':
	training_set = efficient_Dataset(data_indices=range(y_train.shape[0]), x = x_train, t=t_train, y = y_train)
else:
	print('\tNote: imported all data, but only using training data for', training_data_set)
	vol_ID = vol_name_ls.index(training_data_set)
	index_max = vol_cutoff_indices[vol_ID]
	if vol_ID == 0:
		index_min = 0
	else:
		index_min = vol_cutoff_indices[vol_ID - 1]
	curr_data_indices = range(index_min, index_max)
	training_set = efficient_Dataset(data_indices=curr_data_indices, x = x_train, t=t_train, y = y_train)
test_set = efficient_Dataset(data_indices=range(y_test.shape[0]), x = x_test, t = t_test, y = y_test)
train_loader = torch.utils.data.DataLoader(dataset = training_set, batch_size = batch_size, shuffle = True)
test_loader = torch.utils.data.DataLoader(dataset = test_set, batch_size = batch_size, shuffle = True)


# Determining compute options (GPU? Parallel?)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if model_selection != 'AR':
	lstm_model = torch.nn.DataParallel(lstm_model)


# Setting optimization methods
loss = torch.nn.MSELoss()
# Defining many regularization strengths to iterate over to combat overfitting
assert(model_selection in ['AR', 'LSTM', "ConvTimeLSTMUnet"])
if model_selection == 'AR':
	penalization = 1
elif model_selection == 'LSTM':
	penalization = 0.0001
elif model_selection == 'ConvTimeLSTMUnet':
	penalization = 0.1
optimizer = torch.optim.Adam(lstm_model.parameters(), weight_decay = penalization)
# Training loop
print("Beginning training for regularization strength:", penalization)
loss_list = []
#epochs = int(np.ceil((7*10**5) / x_train.shape[0]))
epochs = 100
loop_begin_time = datetime.now()
for i in range(epochs):
	# Marking the beginning time of epoch
	begin_time = datetime.now()
	for data in train_loader:
		
		# data loader
		batch_x, batch_t, batch_y = data
		batch_x.unsqueeze_(2) # code was originally wrote for 5-band level-1 data
		batch_t.unsqueeze_(2) # this unsqueeze is just allowing the explicit acknowledgement
		batch_y.unsqueeze_(2) # that the level-2 data is 1-band deep
		
		# reshaping data if needed for non-spatial LSTMs
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_x = batch_x.permute(0, 3, 4, 1, 2)
			x_sh = batch_x.shape
			batch_x = batch_x.reshape(x_sh[0]*x_sh[1]*x_sh[2], x_sh[3], x_sh[4])
		# Only further processing time in a time-conscious, non-spatial LSTM
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
			batch_t = batch_t.permute(0, 3, 4, 1, 2)
			t_sh = batch_t.shape
			batch_t = batch_t.reshape(t_sh[0]*t_sh[1]*t_sh[2], t_sh[3], t_sh[4])
			# This next line is fragile to the assumption that
			# bands have the same sampling time difference
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_t = batch_t[:,:,[0]]
		
		# move to GPU
		if model_selection != 'AR':
			batch_x = batch_x.to(device)
			# Only move time tensors to GPU if time-conscious LSTM
			if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
				batch_t = batch_t.to(device)
			batch_y = batch_y.to(device)
		
		# Run the model, determining forward pass based on model selected
		if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
			batch_y_hat = lstm_model(batch_x)
		elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_y_hat = lstm_model(batch_x, batch_t)
		elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet']:
			batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
		
		# Extracting the target prediction based on model output
		if model_selection not in ['AR', 'Identity']:
			batch_y_hat = batch_y_hat[0][0]
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_y_hat = batch_y_hat.reshape(x_sh)
			batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
		batch_y_hat = batch_y_hat[:, [-1], :, :, :]
		
		# calculate and store the loss
		batch_loss = loss(batch_y, batch_y_hat)
		loss_list.append(batch_loss.item())
		
		# update parameters
		if model_selection != 'Identity':
			optimizer.zero_grad()
			batch_loss.backward()
			optimizer.step()
		
	# Marking the end time and computing time difference, also printing epoch information
	end_time = datetime.now()
	time_diff = (end_time - loop_begin_time).total_seconds()
	# With l2 regularization iterations, too many prints; only printing every 10 epochs
	if (i + 1) % 10 == 0:
		print('\tEpoch:', i, '\n\t\tMost recent batch loss:', batch_loss.item(), '\n\t\t' + str(time_diff) + ' seconds elapsed')
loop_end_time = datetime.now()
loop_time_diff = (loop_end_time - loop_begin_time).total_seconds()
print('\tTotal training-loop time:', loop_time_diff)


# Converting loss values into array and saving
loss_array = np.asarray(loss_list)
np.save("outputs/loss_over_iterations" + str(penalization) + ".npy", loss_array)


# Trying to free GPU memory
del batch_x
del batch_t
del batch_y
del batch_y_hat
del batch_loss
torch.cuda.empty_cache()


print("Beginning final evaluation")

	# Determine train set performance by volcano
with torch.no_grad():
    count = 0
    for i in range(len(y_train)):
        batch_x = x_train[[i], :, :, :].unsqueeze(2)
        batch_t = t_train[[i], :, :, :].unsqueeze(2)
        batch_y = y_train[[i], :, :, :].unsqueeze(2)
        
        # reshaping data if needed for non-spatial LSTMs
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
            batch_x = batch_x.permute(0, 3, 4, 1, 2)
            x_sh = batch_x.shape
            batch_x = batch_x.reshape(x_sh[0]*x_sh[1]*x_sh[2], x_sh[3], x_sh[4])
        # Only further processing time in a time-conscious, non-spatial LSTM
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
            batch_t = batch_t.permute(0, 3, 4, 1, 2)
            t_sh = batch_t.shape
            batch_t = batch_t.reshape(t_sh[0]*t_sh[1]*t_sh[2], t_sh[3], t_sh[4])
            # This next line is fragile to the assumption that
            # bands have the same sampling time difference
        if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_t = batch_t[:,:,[0]]
            
        # move to GPU
        if model_selection != 'AR':
            batch_x = batch_x.to(device)
            # Only move time tensors to GPU if time-conscious LSTM
            if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
                batch_t = batch_t.to(device)
            batch_y = batch_y.to(device)
            
        # Run the model, determining forward pass based on model selected
        if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
            batch_y_hat = lstm_model(batch_x)
        elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
            batch_y_hat = lstm_model(batch_x, batch_t)
        elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet']:
            batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
        
        # Extracting the target prediction based on model output
        if model_selection not in ['AR', 'Identity']:
            batch_y_hat = batch_y_hat[0][0]
        if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
            batch_y_hat = batch_y_hat.reshape(x_sh)
            batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
        batch_y_hat = batch_y_hat[:, [-1], :, :, :]
        
        # Moving data off GPU now that model has ran
        batch_y = batch_y.cpu()
        batch_y_hat = batch_y_hat.cpu()
        
        # Transformating the data to temperature values
        train_y_min = torch.tensor(stored_parameters[0, 0])
        train_y_max = torch.tensor(stored_parameters[1, 0])
        batch_y = (batch_y * (train_y_max - train_y_min)) + train_y_min
        batch_y_hat = (batch_y_hat * (train_y_max - train_y_min)) + train_y_min
        
        # Storing all temperature-valued truths and predictions for
        # one root mean squared error calculation to get error in
        # terms of temperature
        if count == 0:
            cpu_y_temps = batch_y
            cpu_y_hat_temps = batch_y_hat
        else:
            cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim = 0)
            cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim = 0)
        count = count + 1
    # calculate and store the loss
    train_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
    train_set_loss = torch.sqrt(train_set_loss)
    train_set_loss = train_set_loss.item()
    # print traination set loss
    print("\tTraining set loss:", train_set_loss)
    # Saving the predictions and corresponding truths
    np.save("outputs/train_prediction" + str(penalization) + ".npy", cpu_y_hat_temps.numpy())
    np.save("outputs/train_truth" + str(penalization) + ".npy", cpu_y_temps.numpy())
    
    vol_ID = 0
    for vol in vol_name_ls:
        if vol_ID == 0:
            index_min = 0
        else:
            index_min = vol_cutoff_indices[vol_ID - 1]
        index_max = vol_cutoff_indices[vol_ID]
        pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
        true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
        vol_loss = loss(pred_vol, true_vol)
        vol_loss = torch.sqrt(vol_loss)
        vol_loss = vol_loss.item()
        print('\t\tTraining set loss for', vol, ':', vol_loss)
        vol_ID += 1

# Saving the train set loss
np.save("outputs/final_train_loss" + str(penalization) + ".npy", np.asarray(train_set_loss))

# Determine test set performance by volcano
with torch.no_grad():
	count = 0
	for i in range(len(y_test)):
		batch_x = x_test[[i], :, :, :].unsqueeze(2)
		batch_t = t_test[[i], :, :, :].unsqueeze(2)
		batch_y = y_test[[i], :, :, :].unsqueeze(2)
		
		# reshaping data if needed for non-spatial LSTMs
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_x = batch_x.permute(0, 3, 4, 1, 2)
			x_sh = batch_x.shape
			batch_x = batch_x.reshape(x_sh[0]*x_sh[1]*x_sh[2], x_sh[3], x_sh[4])
		# Only further processing time in a time-conscious, non-spatial LSTM
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM']:
			batch_t = batch_t.permute(0, 3, 4, 1, 2)
			t_sh = batch_t.shape
			batch_t = batch_t.reshape(t_sh[0]*t_sh[1]*t_sh[2], t_sh[3], t_sh[4])
			# This next line is fragile to the assumption that
			# bands have the same sampling time difference
		if model_selection in ['TimeLSTM', 'TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_t = batch_t[:,:,[0]]
			
		# move to GPU
		if model_selection != 'AR':
			batch_x = batch_x.to(device)
			# Only move time tensors to GPU if time-conscious LSTM
			if model_selection not in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
				batch_t = batch_t.to(device)
			batch_y = batch_y.to(device)
			
		# Run the model, determining forward pass based on model selected
		if model_selection in ['AR', 'Identity', 'LSTM', 'ConvLSTM']:
			batch_y_hat = lstm_model(batch_x)
		elif model_selection in ['TimeAwareLSTM', 'ConvTimeAwareLSTM']:
			batch_y_hat = lstm_model(batch_x, batch_t)
		elif model_selection in ['TimeLSTM', 'ConvTimeLSTM', 'ConvTimeLSTMUnet']:
			batch_y_hat = lstm_model(batch_x, batch_x, batch_t)
		
		# Extracting the target prediction based on model output
		if model_selection not in ['AR', 'Identity']:
			batch_y_hat = batch_y_hat[0][0]
		if model_selection in ['AR', 'Identity', 'LSTM', 'TimeLSTM', 'TimeAwareLSTM']:
			batch_y_hat = batch_y_hat.reshape(x_sh)
			batch_y_hat = batch_y_hat.permute(0, 3, 4, 1, 2)
		batch_y_hat = batch_y_hat[:, [-1], :, :, :]
		
		# Moving data off GPU now that model has ran
		batch_y = batch_y.cpu()
		batch_y_hat = batch_y_hat.cpu()
		
		# Transformating the data to temperature values
		train_y_min = torch.tensor(stored_parameters[0, 0])
		train_y_max = torch.tensor(stored_parameters[1, 0])
		batch_y = (batch_y * (train_y_max - train_y_min)) + train_y_min
		batch_y_hat = (batch_y_hat * (train_y_max - train_y_min)) + train_y_min
		
		# Storing all temperature-valued truths and predictions for
		# one root mean squared error calculation to get error in
		# terms of temperature
		if count == 0:
			cpu_y_temps = batch_y
			cpu_y_hat_temps = batch_y_hat
		else:
			cpu_y_temps = torch.cat([cpu_y_temps, batch_y], dim = 0)
			cpu_y_hat_temps = torch.cat([cpu_y_hat_temps, batch_y_hat], dim = 0)
		count = count + 1
	# calculate and store the loss
	test_set_loss = loss(cpu_y_hat_temps, cpu_y_temps)
	test_set_loss = torch.sqrt(test_set_loss)
	test_set_loss = test_set_loss.item()
	# print test set loss
	print("\tTest set loss:", test_set_loss)
	# Saving the predictions and corresponding truths
	np.save("outputs/test_prediction" + str(penalization) + ".npy", cpu_y_hat_temps.numpy())
	np.save("outputs/test_truth" + str(penalization) + ".npy", cpu_y_temps.numpy())
	
	vol_ID = 0
	for vol in vol_name_ls:
		if vol_ID == 0:
			index_min = 0
		else:
			index_min = vol_cutoff_indices_test[vol_ID - 1]
		index_max = vol_cutoff_indices_test[vol_ID]
		pred_vol = cpu_y_hat_temps[index_min:index_max, :, :, :, :]
		true_vol = cpu_y_temps[index_min:index_max, :, :, :, :]
		vol_loss = loss(pred_vol, true_vol)
		vol_loss = torch.sqrt(vol_loss)
		vol_loss = vol_loss.item()
		print('\t\tTest set loss for', vol, ':', vol_loss)
		vol_ID += 1

# Saving the testation set loss
np.save("outputs/final_test_loss" + str(penalization) + ".npy", np.asarray(test_set_loss))

# Saving the model
torch.save(lstm_model, "outputs/model" + str(penalization) + ".pt")

# Print breakpoint for human reading between regularization strengths
print('\n')
