#!/usr/bin/env python3

import tensorflow as tf

import numpy as np

import os

import sys

import argparse

import logging

import string

import random

from shutil import copyfile

from datetime import datetime

from neural_network.NeuralNetwork import *

from neural_network.activation_fn import *

from training.Trainer import *

from training.DataContainer import *

from training.DataProvider import *

from training.DataQueue import *

#used for creating a "unique" id for a run (almost impossible to generate the same twice)

def id_generator(size=8, chars=string.ascii_uppercase + string.ascii_lowercase + string.digits):

return ''.join(random.SystemRandom().choice(chars) for _ in range(size))

logging.basicConfig(filename='train.log',level=logging.DEBUG)

#define command line arguments

parser = argparse.ArgumentParser(fromfile_prefix_chars='@')

parser.add_argument("--restart", type=str, default=None, help="restart training from a specific folder")

parser.add_argument("--num_features", type=int, help="dimensionality of feature vectors")

parser.add_argument("--num_basis", type=int, help="number of radial basis functions")

parser.add_argument("--num_blocks", type=int, help="number of interaction blocks")

parser.add_argument("--num_residual_atomic", type=int, help="number of residual layers for atomic refinements")

parser.add_argument("--num_residual_interaction", type=int, help="number of residual layers for the message phase")

parser.add_argument("--num_residual_output", type=int, help="number of residual layers for the output blocks")

parser.add_argument("--cutoff", default=10.0, type=float, help="cutoff distance for short range interactions")

parser.add_argument("--use_electrostatic", default=1, type=int, help="use electrostatics in energy prediction (0/1)")

parser.add_argument("--use_dispersion", default=1, type=int, help="use dispersion in energy prediction (0/1)")

parser.add_argument("--grimme_s6", default=None, type=float, help="grimme s6 dispersion coefficient")

parser.add_argument("--grimme_s8", default=None, type=float, help="grimme s8 dispersion coefficient")

parser.add_argument("--grimme_a1", default=None, type=float, help="grimme a1 dispersion coefficient")

parser.add_argument("--grimme_a2", default=None, type=float, help="grimme a2 dispersion coefficient")

parser.add_argument("--dataset", type=str, help="file path to dataset")

parser.add_argument("--num_train", type=int, help="number of training samples")

parser.add_argument("--num_valid", type=int, help="number of validation samples")

parser.add_argument("--seed", default=42, type=int, help="seed for splitting dataset into training/validation/test")

parser.add_argument("--max_steps", type=int, help="maximum number of training steps")

parser.add_argument("--learning_rate", default=0.001, type=float, help="learning rate used by the optimizer")

parser.add_argument("--max_norm", default=1000.0, type=float, help="max norm for gradient clipping")

parser.add_argument("--ema_decay", default=0.999, type=float, help="exponential moving average decay used by the trainer")

parser.add_argument("--keep_prob", default=1.0, type=float, help="keep probability for dropout regularization of rbf layer")

parser.add_argument("--l2lambda", type=float, help="lambda multiplier for l2 loss (regularization)")

parser.add_argument("--nhlambda", type=float, help="lambda multiplier for non-hierarchicality loss (regularization)")

parser.add_argument("--decay_steps", type=int, help="decay the learning rate every N steps by decay_rate")

parser.add_argument("--decay_rate", type=float, help="factor with which the learning rate gets multiplied by every decay_steps steps")

parser.add_argument("--batch_size", type=int, help="batch size used per training step")

parser.add_argument("--valid_batch_size", type=int, help="batch size used for going through validation_set")

parser.add_argument('--force_weight', default=52.91772105638412, type=float, help="this defines the force contribution to the loss function relative to the energy contribution (to take into account the different numerical range)")

parser.add_argument('--charge_weight', default=14.399645351950548, type=float, help="this defines the charge contribution to the loss function relative to the energy contribution (to take into account the different numerical range)")

parser.add_argument('--dipole_weight', default=27.211386024367243, type=float, help="this defines the dipole contribution to the loss function relative to the energy contribution (to take into account the different numerical range)")

parser.add_argument('--summary_interval', type=int, help="write a summary every N steps")

parser.add_argument('--validation_interval', type=int, help="check performance on validation set every N steps")

parser.add_argument('--save_interval', type=int, help="save progress every N steps")

parser.add_argument('--record_run_metadata', type=int, help="records metadata like memory consumption etc.")

#if no command line arguments are present, config file is parsed

config_file='config.txt'

if len(sys.argv) == 1:

if os.path.isfile(config_file):

args = parser.parse_args(["@"+config_file])

else:

args = parser.parse_args(["--help"])

else:

args = parser.parse_args()

#create directories

#a unique directory name is created for this run based on the input

if args.restart is None:

directory=datetime.utcnow().strftime("%Y%m%d%H%M%S") + "_" + id_generator() +"_F"+str(args.num_features)+"K"+str(args.num_basis)+"b"+str(args.num_blocks)+"a"+str(args.num_residual_atomic)+"i"+str(args.num_residual_interaction)+"o"+str(args.num_residual_output)+"cut"+str(args.cutoff)+"e"+str(args.use_electrostatic)+"d"+str(args.use_dispersion)+"l2"+str(args.l2lambda)+"nh"+str(args.nhlambda)+"keep"+str(args.keep_prob)

else:

directory=args.restart

logging.info("creating directories...")

if not os.path.exists(directory):

os.makedirs(directory)

best_dir = os.path.join(directory, 'best')

if not os.path.exists(best_dir):

os.makedirs(best_dir)

log_dir = os.path.join(directory, 'logs')

if not os.path.exists(log_dir):

os.makedirs(log_dir)

best_loss_file = os.path.join(best_dir, 'best_loss.npz')

best_checkpoint = os.path.join(best_dir, 'best_model.ckpt')

step_checkpoint = os.path.join(log_dir, 'model.ckpt')

#write config file (to restore command line arguments)

logging.info("writing args to file...")

with open(os.path.join(directory, config_file), 'w') as f:

for arg in vars(args):

f.write('--'+ arg + '='+ str(getattr(args, arg)) + "\n")

#load dataset

logging.info("loading dataset...")

data = DataContainer(args.dataset)

#generate DataProvider (splits dataset into training, validation and test set based on seed)

data_provider = DataProvider(data, args.num_train, args.num_valid, args.batch_size, args.valid_batch_size, seed=args.seed)

#create neural network

logging.info("creating neural network...")

nn = NeuralNetwork(F=args.num_features,

scope="neural_network")

logging.info("prepare training...")

#generate data queues for efficient training

train_queue = DataQueue(data_provider.next_batch, capacity=1000, scope="train_data_queue")

valid_queue = DataQueue(data_provider.next_valid_batch, capacity=args.num_valid//args.valid_batch_size, scope="valid_data_queue")

#get values for training and validation set

Eref_t, Earef_t, Fref_t, Z_t, Dref_t, Qref_t, Qaref_t, R_t, idx_i_t, idx_j_t, batch_seg_t = train_queue.dequeue_op

Eref_v, Earef_v, Fref_v, Z_v, Dref_v, Qref_v, Qaref_v, R_v, idx_i_v, idx_j_v, batch_seg_v = valid_queue.dequeue_op

#calculate all necessary quantities (unscaled partial charges, energies, forces)

Ea_t, Qa_t, Dij_t, nhloss_t = nn.atomic_properties(Z_t, R_t, idx_i_t, idx_j_t)

Ea_v, Qa_v, Dij_v, nhloss_v = nn.atomic_properties(Z_v, R_v, idx_i_v, idx_j_v)

energy_t, forces_t = nn.energy_and_forces_from_atomic_properties(Ea_t, Qa_t, Dij_t, Z_t, R_t, idx_i_t, idx_j_t, Qref_t, batch_seg_t)

energy_v, forces_v = nn.energy_and_forces_from_atomic_properties(Ea_v, Qa_v, Dij_v, Z_v, R_v, idx_i_v, idx_j_v, Qref_v, batch_seg_v)

#total charge

Qtot_t = tf.segment_sum(Qa_t, batch_seg_t)

Qtot_v = tf.segment_sum(Qa_v, batch_seg_v)

#dipole moment vector

QR_t = tf.stack([Qa_t*R_t[:,0], Qa_t*R_t[:,1], Qa_t*R_t[:,2]],1)

QR_v = tf.stack([Qa_v*R_v[:,0], Qa_v*R_v[:,1], Qa_v*R_v[:,2]],1)

D_t = tf.segment_sum(QR_t, batch_seg_t)

D_v = tf.segment_sum(QR_v, batch_seg_v)

#function to calculate loss, mean squared error, mean absolute error between two values

def calculate_errors(val1, val2, weights=1):

return loss, mse, mae

with tf.name_scope("loss"):

#calculate energy, force, charge and dipole errors/loss

#energy

if data.E is not None:

eloss_t, emse_t, emae_t = calculate_errors(Eref_t, energy_t)

eloss_v, emse_v, emae_v = calculate_errors(Eref_v, energy_v)

else:

eloss_t, emse_t, emae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

eloss_v, emse_v, emae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

#atomic energies

if data.Ea is not None:

ealoss_t, eamse_t, eamae_t = calculate_errors(Earef_t, Ea_t)

ealoss_v, eamse_v, eamae_v = calculate_errors(Earef_v, Ea_v)

else:

ealoss_t, eamse_t, eamae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

ealoss_v, eamse_v, eamae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

#forces

if data.F is not None:

floss_t, fmse_t, fmae_t = calculate_errors(Fref_t, forces_t)

floss_v, fmse_v, fmae_v = calculate_errors(Fref_v, forces_v)

else:

floss_t, fmse_t, fmae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

floss_v, fmse_v, fmae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

#charge

if data.Q is not None:

qloss_t, qmse_t, qmae_t = calculate_errors(Qref_t, Qtot_t)

qloss_v, qmse_v, qmae_v = calculate_errors(Qref_v, Qtot_v)

else:

qloss_t, qmse_t, qmae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

qloss_v, qmse_v, qmae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

#atomic charges

if data.Qa is not None:

qaloss_t, qamse_t, qamae_t = calculate_errors(Qaref_t, Qa_t)

qaloss_v, qamse_v, qamae_v = calculate_errors(Qaref_v, Qa_v)

else:

qaloss_t, qamse_t, qamae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

qaloss_v, qamse_v, qamae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

#dipole

if data.D is not None:

dloss_t, dmse_t, dmae_t = calculate_errors(Dref_t, D_t)

dloss_v, dmse_v, dmae_v = calculate_errors(Dref_v, D_v)

else:

dloss_t, dmse_t, dmae_t = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

dloss_v, dmse_v, dmae_v = tf.constant(0.0), tf.constant(0.0), tf.constant(0.0)

#define additional variables (such that certain losses can be overwritten)

eloss_train = eloss_t

floss_train = floss_t

qloss_train = qloss_t

dloss_train = dloss_t

eloss_valid = eloss_v

floss_valid = floss_v

qloss_valid = qloss_v

dloss_valid = dloss_v

#atomic energies are present, so they replace the normal energy loss

if data.Ea is not None:

eloss_train = ealoss_t

eloss_valid = ealoss_v

#atomic charges are present, so they replace the normal charge loss and nullify dipole loss

if data.Qa is not None:

qloss_train = qaloss_t

qloss_valid = qaloss_v

dloss_train = tf.constant(0.0)

dloss_valid = tf.constant(0.0)

#define loss function (used to train the model)

l2loss = tf.reduce_mean(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))

loss_t = eloss_train + args.force_weight*floss_train + args.charge_weight*qloss_train + args.dipole_weight*dloss_train + args.nhlambda*nhloss_t + args.l2lambda*l2loss

loss_v = eloss_valid + args.force_weight*floss_valid + args.charge_weight*qloss_valid + args.dipole_weight*dloss_valid + args.nhlambda*nhloss_v + args.l2lambda*l2loss

#create trainer

trainer = Trainer(args.learning_rate, args.decay_steps, args.decay_rate, scope="trainer")

with tf.name_scope("trainer_ops"):

train_op = trainer.build_train_op(loss_t, args.ema_decay, args.max_norm)

save_variable_backups_op = trainer.save_variable_backups()

load_averaged_variables_op = trainer.load_averaged_variables()

restore_variable_backups_op = trainer.restore_variable_backups()

#creates a summary from key-value pairs given a dictionary

def create_summary(dictionary):

return summary

#create summary writer

nn_summary_op = tf.summary.merge_all()

summary_writer = tf.summary.FileWriter(logdir=log_dir, graph=tf.get_default_graph())

#create saver

with tf.name_scope("saver"):

saver = tf.train.Saver(max_to_keep=50)

#save/load best recorded loss (only the best model is saved)

if os.path.isfile(best_loss_file):

loss_file = np.load(best_loss_file)

best_loss = loss_file["loss"].item()

best_emae = loss_file["emae"].item()

best_ermse = loss_file["ermse"].item()

best_fmae = loss_file["fmae"].item()

best_frmse = loss_file["frmse"].item()

best_qmae = loss_file["qmae"].item()

best_qrmse = loss_file["qrmse"].item()

best_dmae = loss_file["dmae"].item()

best_drmse = loss_file["drmse"].item()

best_step = loss_file["step"].item()

else:

best_loss = np.Inf #initialize best loss to infinity

best_emae = np.Inf

best_ermse = np.Inf

best_fmae = np.Inf

best_frmse = np.Inf

best_qmae = np.Inf

best_qrmse = np.Inf

best_dmae = np.Inf

best_drmse = np.Inf

best_step = 0.

np.savez(best_loss_file, loss=best_loss, emae=best_emae, ermse=best_ermse,

fmae=best_fmae, frmse=best_frmse,

qmae=best_qmae, qrmse=best_qrmse,

dmae=best_dmae, drmse=best_drmse,

step=best_step)

#for calculating average performance on the training set

def reset_averages():

return 0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

def update_averages(num, tmploss_avg, tmploss, emse_avg, emse, emae_avg, emae, fmse_avg, fmse, fmae_avg, fmae,

return num, tmploss_avg, emse_avg, emae_avg, fmse_avg, fmae_avg, qmse_avg, qmae_avg, dmse_avg, dmae_avg

#initialize training set error averages

num_t, tmploss_avg_t, emse_avg_t, emae_avg_t, fmse_avg_t, fmae_avg_t, qmse_avg_t, qmae_avg_t, dmse_avg_t, dmae_avg_t = reset_averages()

#create tensorflow session

with tf.Session() as sess:

if (args.record_run_metadata > 0):

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)

run_metadata = tf.RunMetadata()

else:

run_options = None

run_metadata = None

#start data queues

coord = tf.train.Coordinator()

train_queue.create_thread(sess, coord)

valid_queue.create_thread(sess, coord)

#initialize variables

tf.global_variables_initializer().run()

#restore latest checkpoint

checkpoint = tf.train.latest_checkpoint(log_dir)

if checkpoint is not None:

step = int(checkpoint.split('-')[-1]) #reads step from checkpoint filename

saver.restore(sess, checkpoint)

sess.run(trainer.global_step.assign(step))

else:

step = 0

#training loop

logging.info("starting training...")

while not coord.should_stop():

#finish training when maximum number of iterations is reached

if step > args.max_steps:

coord.request_stop()

break

#perform training step

step += 1

_, tmploss, emse, emae, fmse, fmae, qmse, qmae, dmse, dmae = sess.run([train_op, loss_t, emse_t, emae_t, fmse_t, fmae_t, qmse_t, qmae_t, dmse_t, dmae_t], options=run_options, feed_dict={nn.keep_prob: args.keep_prob}, run_metadata=run_metadata)

#update averages

num_t, tmploss_avg_t, emse_avg_t, emae_avg_t, fmse_avg_t, fmae_avg_t, qmse_avg_t, qmae_avg_t, dmse_avg_t, dmae_avg_t = update_averages(num_t, tmploss_avg_t, tmploss, emse_avg_t, emse, emae_avg_t, emae, fmse_avg_t, fmse, fmae_avg_t, fmae, qmse_avg_t, qmse, qmae_avg_t, qmae, dmse_avg_t, dmse, dmae_avg_t, dmae)

#save progress

if (step % args.save_interval == 0):

saver.save(sess, step_checkpoint, global_step=step)

#check performance on the validation set

if (step % args.validation_interval == 0):

#save backup variables and load averaged variables

sess.run(save_variable_backups_op)

sess.run(load_averaged_variables_op)

#initialize averages to 0

num_v, tmploss_avg_v, emse_avg_v, emae_avg_v, fmse_avg_v, fmae_avg_v, qmse_avg_v, qmae_avg_v, dmse_avg_v, dmae_avg_v = reset_averages()

#compute averages

for i in range(args.num_valid//args.valid_batch_size):

tmploss, emse, emae, fmse, fmae, qmse, qmae, dmse, dmae = sess.run([loss_v, emse_v, emae_v, fmse_v, fmae_v, qmse_v, qmae_v, dmse_v, dmae_v])

num_v, tmploss_avg_v, emse_avg_v, emae_avg_v, fmse_avg_v, fmae_avg_v, qmse_avg_v, qmae_avg_v, dmse_avg_v, dmae_avg_v = update_averages(num_v, tmploss_avg_v, tmploss, emse_avg_v, emse, emae_avg_v, emae, fmse_avg_v, fmse, fmae_avg_v, fmae, qmse_avg_v, qmse, qmae_avg_v, qmae, dmse_avg_v, dmse, dmae_avg_v, dmae)

#store results in dictionary

results = {}

results["loss_valid"] = tmploss_avg_v

if data.E is not None:

results["energy_mae_valid"] = emae_avg_v

results["energy_rmse_valid"] = np.sqrt(emse_avg_v)

if data.F is not None:

results["forces_mae_valid"] = fmae_avg_v

results["forces_rmse_valid"] = np.sqrt(fmse_avg_v)

if data.Q is not None:

results["charge_mae_valid"] = qmae_avg_v

results["charge_rmse_valid"] = np.sqrt(qmse_avg_v)

if data.D is not None:

results["dipole_mae_valid"] = dmae_avg_v

results["dipole_rmse_valid"] = np.sqrt(dmse_avg_v)

if results["loss_valid"] < best_loss:

best_loss = results["loss_valid"]

if data.E is not None:

best_emae = results["energy_mae_valid"]

best_ermse = results["energy_rmse_valid"]

else:

best_emae = np.Inf

best_ermse = np.Inf

if data.F is not None:

best_fmae = results["forces_mae_valid"]

best_frmse = results["forces_rmse_valid"]

else:

best_fmae = np.Inf

best_frmse = np.Inf

if data.Q is not None:

best_qmae = results["charge_mae_valid"]

best_qrmse = results["charge_rmse_valid"]

else:

best_qmae = np.Inf

best_qrmse = np.Inf

if data.D is not None:

best_dmae = results["dipole_mae_valid"]

best_drmse = results["dipole_rmse_valid"]

else:

best_dmae = np.Inf

best_drmse = np.Inf

best_step = step

np.savez(best_loss_file, loss=best_loss, emae=best_emae, ermse=best_ermse,

fmae=best_fmae, frmse=best_frmse,

qmae=best_qmae, qrmse=best_qrmse,

dmae=best_dmae, drmse=best_drmse,

step=best_step)

nn.save(sess, best_checkpoint, global_step=step)

results["loss_best"] = best_loss

if data.E is not None:

results["energy_mae_best"] = best_emae

results["energy_rmse_best"] = best_ermse

if data.F is not None:

results["forces_mae_best"] = best_fmae

results["forces_rmse_best"] = best_frmse

if data.Q is not None:

results["charge_mae_best"] = best_qmae

results["charge_rmse_best"] = best_qrmse

if data.D is not None:

results["dipole_mae_best"] = best_dmae

results["dipole_rmse_best"] = best_drmse

summary = create_summary(results)

summary_writer.add_summary(summary, global_step=step)

#restore backup variables

sess.run(restore_variable_backups_op)

#generate summaries

if (step % args.summary_interval == 0) and (step > 0):

results = {}

results["loss_train"] = tmploss_avg_t

if data.E is not None:

results["energy_mae_train"] = emae_avg_t

results["energy_rmse_train"] = np.sqrt(emse_avg_t)

if data.F is not None:

results["forces_mae_train"] = fmae_avg_t

results["forces_rmse_train"] = np.sqrt(fmse_avg_t)

if data.Q is not None:

results["charge_mae_train"] = qmae_avg_t

results["charge_rmse_train"] = np.sqrt(qmse_avg_t)

if data.D is not None:

results["dipole_mae_train"] = dmae_avg_t

results["dipole_rmse_train"] = np.sqrt(dmse_avg_t)

num_t, tmploss_avg_t, emse_avg_t, emae_avg_t, fmse_avg_t, fmae_avg_t, qmse_avg_t, qmae_avg_t, dmse_avg_t, dmae_avg_t = reset_averages()

summary = create_summary(results)

summary_writer.add_summary(summary, global_step=step)

nn_summary = nn_summary_op.eval()

summary_writer.add_summary(nn_summary, global_step=step)

if (args.record_run_metadata > 0):

summary_writer.add_run_metadata(run_metadata, 'step %d' % step, global_step=step)

if data.E is not None:

print(str(step)+'/'+str(args.max_steps), "loss:", results["loss_train"], "best:", best_loss, "emae:", results["energy_mae_train"], "best:", best_emae)