def prepare_model_megnet(individuals, epochs, outfile, excl=[]):
# prepares model file
# prepares Megnet model based on list of individuals
# uses total energy per atom
# excl - excluding particular stoichiometry - important for network learning
structures = []
energies = []
adapt = AseAtomsAdaptor()
empty = 0
if not excl:
empty = 1
i = 0
for ind in individuals:
struct_ase = ind.get_init_structure()
chem_sym = struct_ase.get_chemical_symbols()
e_tot = ind.e_tot
struct_pymatgen = adapt.get_structure(struct_ase)
flag = 1
if empty == 0 and chem_sym == excl:
flag = 0
if flag == 1:
structures.append(struct_pymatgen)
energies.append(e_tot)
i = i + 1
print("read data of " + str(i) + " structures total")
# standard vales as taken from Megnet manual
nfeat_bond = 100
nfeat_global = 2
r_cutoff = 5
gaussian_centers = np.linspace(0, r_cutoff + 1, nfeat_bond)
gaussian_width = 0.5
distance_converter = GaussianDistance(gaussian_centers, gaussian_width)
graph_converter = CrystalGraph(bond_converter=distance_converter, cutoff=r_cutoff)
model = MEGNetModel(nfeat_bond, nfeat_global, graph_converter=graph_converter)
# model training
model.train(structures, energies, epochs=epochs)
model.save_model(outfile)
train_targets,
val_graphs,
val_targets,
epochs=EPOCHS,
verbose=2,
initial_epoch=0,
callbacks=callbacks)
# 6. Model testing
## load the best model with lowest validation error
files = glob("./callback/*.hdf5")
best_model = sorted(files, key=os.path.getctime)[-1]
model.load_weights(best_model)
model.save_model("best_model.hdf5")
def evaluate(test_graphs, test_targets):
"""
Evaluate the test errors using test_graphs and test_targets
Args:
test_graphs (list): list of graphs
test_targets (list): list of target properties
Returns:
mean absolute errors
"""
test_data = model.graph_converter.get_flat_data(test_graphs, test_targets)
gen = GraphBatchDistanceConvert(
targets[it][i] = prdc
# targets[i] = (model.predict_structure(structures[i]).ravel() + targets[i])/2
logging.info('Data count: {dc}, std orig dft value: {std_orig}, std of model output: {std_model}'.format(
dc=len(targets_lst), std_orig=np.std(targets_lst), std_model=np.std(prediction_lst)))
logging.info('Data count: {dc}, Mean orig: {mean_orig}, Mean_model: {mean_model}'.format(
dc=len(targets_lst), mean_orig=np.mean(targets_lst), mean_model=np.mean(prediction_lst)))
f = open(dump_model_name + '_'+ it + '.txt', 'wb') # to store and analyze the error
pickle.dump(error_lst, f)
f.close()
# model = MEGNetModel(10, 2, nblocks=3, lr=1e-3,
# n1=4, n2=4, n3=4, npass=1, ntarget=1,
# graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5)))
model = MEGNetModel(nfeat_edge=10, nfeat_global=2, graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5)))
model.save_model(dump_model_name+'_1by1_init_randomly' + '.hdf5')
init_model_tag = 'EGPHS'
ep = 5000
callback = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=10, restore_best_weights=True)
for s in test_structures:
test_input.append(model.graph_converter.graph_to_input(model.graph_converter.convert(s)))
db_short_full_dict = {'G': 'gllb-sc', 'H': 'hse', 'S': 'scan', 'P': 'pbe', 'E': 'E1'}
def construct_dataset_from_str(db_short_str):
s = []
t = []
for i in range(len(db_short_str)):
s.extend(structures[db_short_full_dict[db_short_str[i]]])
Xtrain = inputs.iloc[0:boundary]['structure']
ytrain = inputs.iloc[0:boundary]['band_gap']
Xtest = inputs.iloc[boundary:]['structure']
ytest = inputs.iloc[boundary:]['band_gap']
nfeat_bond = 10
nfeat_global = 2
r_cutoff = 5
gaussian_centers = np.linspace(0, 5, 10)
gaussian_width = 0.5
distance_convertor = GaussianDistance(gaussian_centers, gaussian_width)
bond_convertor = CrystalGraph(bond_convertor=distance_convertor,
cutoff=r_cutoff)
graph_convertor = CrystalGraph(
bond_convertor=GaussianDistance(np.linspace(0, 5, 10), 0.5))
model = MEGNetModel(nfeat_bond, nfeat_global, graph_convertor=graph_convertor)
model.from_file('fitted_gap_model.hdf5')
model.train(Xtrain,
ytrain,
epochs=epochs,
batch_size=batch_size,
validation_structures=Xtest,
validation_targets=ytest,
scrub_failed_structures=True)
model.save_model('fitted_gap_model.hdf5')
epochs=ep,
save_checkpoint=False,
automatic_correction=False)
idx += data_size[i]
prediction(model)
elif training_mode == 6: # PBE -> HSE ... -> part EXP, one by one, with 20% validation
idx = 0
for i in range(len(data_size)):
model.train(structures[idx:idx+int(0.8*data_size[i])], targets[idx:idx+int(0.8*data_size[i])],
validation_structures=structures[idx+int(0.8*data_size[i]):(idx+data_size[i])],
validation_targets=targets[idx+int(0.8*data_size[i]):(idx+data_size[i])],
callbacks=[callback, XiaotongCB((test_input, test_targets), commit_id)],
epochs=ep,
save_checkpoint=False,
automatic_correction=False)
model.save_model(commit_id+'_'+str(training_mode)+'_'+str(i)+'.hdf5')
idx += data_size[i]
prediction(model)
elif training_mode == 7: # all training set together with 20% validation
l = len(structures)
c = list(zip(structures, targets))
random.shuffle(c)
structures, targets = zip(*c)
model.train(structures[:int(0.8 * l)], targets[:int(0.8 * l)],
validation_structures=structures[int(0.8 * l):],
validation_targets=targets[int(0.8 * l):],
callbacks=[callback, XiaotongCB((test_input, test_targets), commit_id)],
epochs=ep*len(data_size),
save_checkpoint=False,
automatic_correction=False)
prediction(model)
def main() -> None:
"""Execute main script."""
parser = ArgumentParser()
parser.add_argument(
"--train",
action="store_true",
help="Whether to train the model.",
dest="do_train",
)
parser.add_argument(
"--eval",
action="store_true",
help="Whether to evaluate the model.",
dest="do_eval",
)
parser.add_argument(
"--which",
choices=["MEGNet", "VGP", "ProbNN"],
required=("--train" in sys.argv),
help=(
"Which components to train: "
"MEGNet -- Just the MEGNetModel; "
"VGP -- Just the VGP part of the ProbNN; "
"ProbNN -- The whole ProbNN."
),
dest="which",
)
parser.add_argument(
"--epochs",
"-n",
type=int,
required=("--train" in sys.argv),
help="Number of training epochs.",
dest="epochs",
)
parser.add_argument(
"--inducing",
"-i",
type=int,
help="Number of inducing index points.",
default=500,
dest="num_inducing",
)
args = parser.parse_args()
do_train: bool = args.do_train
do_eval: bool = args.do_eval
which_model: str = args.which
epochs: int = args.epochs
num_inducing: int = args.num_inducing
# Load the MEGNetModel into memory
try:
meg_model: MEGNetModel = MEGNetModel.from_file(str(MEGNET_MODEL_DIR))
except FileNotFoundError:
meg_model = MEGNetModel(**default_megnet_config())
# Load the data into memory
df = download_data(PHONONS_URL, PHONONS_SAVE_DIR)
structures = df["structure"]
targets = df["last phdos peak"]
num_data = len(structures)
print(f"{num_data} datapoints loaded.")
num_training = floor(num_data * TRAINING_RATIO)
print(f"{num_training} training data, {num_data-num_training} test data.")
train_structs = structures[:num_training]
train_targets = targets[:num_training]
test_structs = structures[num_training:]
test_targets = targets[num_training:]
if which_model == "MEGNet":
if do_train:
tf_callback = TensorBoard(MEGNET_LOGS / NOW, write_graph=False)
meg_model.train(
train_structs,
train_targets,
test_structs,
test_targets,
automatic_correction=False,
dirname="meg_checkpoints",
epochs=epochs,
callbacks=[tf_callback],
verbose=VERBOSITY,
)
meg_model.save_model(str(MEGNET_MODEL_DIR))
if do_eval:
train_predicted = meg_model.predict_structures(train_structs).flatten()
train_mae = MAE(train_predicted, None, train_targets)
metric_logger.info("MEGNet train MAE = %f", train_mae)
test_predicted = meg_model.predict_structures(test_structs).flatten()
test_mae = MAE(test_predicted, None, test_targets)
metric_logger.info("MEGNet test MAE = %f", test_mae)
else:
# Load the ProbNN into memory
try:
prob_model: MEGNetProbModel = MEGNetProbModel.load(PROB_MODEL_DIR)
except FileNotFoundError:
prob_model = MEGNetProbModel(meg_model, num_inducing, metrics=["MAE"])
if do_train:
if which_model == "VGP":
prob_model.set_frozen("NN", recompile=False)
prob_model.set_frozen(["VGP", "Norm"], freeze=False)
tf_callback = TensorBoard(VGP_LOGS / NOW, write_graph=False)
else:
prob_model.set_frozen(["VGP", "NN", "Norm"], freeze=False)
tf_callback = TensorBoard(FULL_MODEL_LOGS / NOW, write_graph=False)
prob_model.train(
train_structs,
train_targets,
epochs,
test_structs,
test_targets,
callbacks=[tf_callback],
verbose=VERBOSITY,
)
prob_model.save(PROB_MODEL_DIR)
if do_eval:
train_metrics = evaluate_uq_metrics(
prob_model, train_structs, train_targets
)
log_metrics(train_metrics, "training")
test_metrics = evaluate_uq_metrics(prob_model, test_structs, test_targets)
log_metrics(test_metrics, "test")
train_targets,
val_graphs,
val_targets,
epochs=EPOCHS,
verbose=2,
initial_epoch=0,
callbacks=callbacks)
# 6. Model testing
## load the best model with lowest validation error
files = glob('./callback/*.hdf5')
best_model = sorted(files, key=os.path.getctime)[-1]
model.load_weights(best_model)
model.save_model('best_model.hdf5')
def evaluate(test_graphs, test_targets):
"""
Evaluate the test errors using test_graphs and test_targets
Args:
test_graphs (list): list of graphs
test_targets (list): list of target properties
Returns:
mean absolute errors
"""
test_data = model.graph_converter.get_flat_data(test_graphs, test_targets)
gen = GraphBatchDistanceConvert(
for i in range(len(data_size)):
model.train(
structures[idx:idx + int(0.8 * data_size[i])],
targets[idx:idx + int(0.8 * data_size[i])],
validation_structures=structures[idx +
int(0.8 * data_size[i]):(
idx + data_size[i])],
validation_targets=targets[idx + int(0.8 * data_size[i]):(
idx + data_size[i])],
# callbacks=[callback, XiaotongCB((test_input, test_targets), commit_id)],
callbacks=[callback],
epochs=ep,
save_checkpoint=False,
batch_size=512,
automatic_correction=False)
model.save_model(dump_model_name + '_' + str(i) + '.hdf5')
idx += data_size[i]
prediction(model)
elif training_mode == 11: # PBE -> HSE ... -> part EXP, one by one, with 20% of last dataset as validation
idx = 0
for i in range(len(data_size)):
model.train(
structures[idx:idx + int(0.8 * data_size[i])],
targets[idx:idx + int(0.8 * data_size[i])],
validation_structures=structures[sum(data_size[:-1]) +
int(0.8 * data_size[-1]):],
validation_targets=targets[sum(data_size[:-1]) +
int(0.8 * data_size[-1]):],
# validation_structures=structures[idx+int(0.8*data_size[i]):(idx+data_size[i])],
# validation_targets=targets[idx+int(0.8*data_size[i]):(idx+data_size[i])],
# callbacks=[callback, XiaotongCB((test_input, test_targets), commit_id)],
