import numpy as np
import pprint
from pyNNsMD.src.device import set_gpu
# No GPU for prediciton or the main class
set_gpu([-1])
from pyNNsMD.NNsMD import NeuralNetEnsemble
from pyNNsMD.hypers.hyper_mlp_eg import DEFAULT_HYPER_PARAM_ENERGY_GRADS as hyper
pprint.pprint(hyper)
# list of angles
anglist = [[1, 0, 2], [1, 0, 4], [2, 0, 4], [0, 1, 3], [0, 1, 8], [3, 1, 8],
[0, 4, 5], [0, 4, 6], [0, 4, 7], [6, 4, 7], [5, 4, 7], [5, 4, 6],
[9, 8, 10], [1, 8, 10], [9, 8, 11], [1, 8, 9], [1, 8, 11],
[10, 8, 11]]
dihedlist = [[5, 1, 2, 9], [3, 1, 2, 4]]
# Load data
atoms = [["C", "C", "H", "H", "C", "F", "F", "F", "C", "F", "H", "H"]] * 2701
geos = np.load("butene/butene_x.npy")
energy = np.load("butene/butene_energy.npy")
grads = np.load("butene/butene_force.npy")
nac = np.load("butene/butene_nac.npy")
print(geos.shape, energy.shape, grads.shape, nac.shape)
hyper["model"]["config"].update({
"atoms": 12,
"states": 2,
parser.add_argument("-m",
"--mode",
default="training",
required=True,
help="Which mode to use train or retrain")
args = vars(parser.parse_args())
file_std_out = open(os.path.join(args['filepath'], "fitlog.txt"), 'w')
sys.stderr = file_std_out
sys.stdout = file_std_out
print("Input argpars:", args)
from pyNNsMD.src.device import set_gpu
set_gpu([int(args['gpus'])])
print("Logic Devices:", tf.config.experimental.list_logical_devices('GPU'))
import pyNNsMD.utils.callbacks
import pyNNsMD.utils.activ
from pyNNsMD.models.schnet_e import SchnetEnergy
from pyNNsMD.utils.data import load_json_file, read_xyz_file, save_json_file
from pyNNsMD.scaler.energy import EnergyStandardScaler
from pyNNsMD.utils.loss import ScaledMeanAbsoluteError, get_lr_metric, r2_metric
from pyNNsMD.plots.loss import plot_loss_curves, plot_learning_curve
from pyNNsMD.plots.pred import plot_scatter_prediction
from kgcnn.utils.adj import define_adjacency_from_distance, coordinates_to_distancematrix
from kgcnn.utils.data import ragged_tensor_from_nested_numpy
from kgcnn.mol.methods import global_proton_dict