def ml_lj(
IMAGES,
filename,
count,
temp,
GSF,
dir="MD_results/",
const_t=False,
lj=False,
fine_tune=None,
loss_fn="l2amp",
save_logs=True,
):
if not os.path.exists(dir):
os.mkdir(dir)
# lj optimization
lj_data = None
cutoff = GSF["cutoff"]
eV_kcalmol = 0.043372093
if lj:
p0 = [
1.33905162,
0.12290683,
6.41914719,
0.64021468,
0.08010004,
8.26082762,
2.29284676,
0.29639983,
0.08071821,
]
params_dict = {"C": [], "O": [], "Cu": []}
lj_model = lj_optim(filename, IMAGES, p0, params_dict, cutoff)
fitted_params = lj_model.fit()
# fitted_params = p0
lj_energies, lj_forces, num_atoms = lj_model.lj_pred(
IMAGES, fitted_params, params_dict)
lj_data = [
lj_energies,
lj_forces,
num_atoms,
fitted_params,
params_dict,
lj_model,
]
# define the number of threads to parallelize training across
torch.set_num_threads(1)
calc = AMP(model=AMPTorch(
IMAGES,
descriptor=Gaussian,
Gs=GSF,
force_coefficient=0.3,
lj_data=lj_data,
label=filename,
save_logs=save_logs,
))
calc.model.lr = 1e-2
if loss_fn == "l2amp":
calc.model.criterion = CustomLoss
elif loss_fn == "tanh":
calc.model.criterion = TanhLoss
calc.model.convergence = {
"energy": 0.02,
"force": 0.02,
"epochs": 500,
"early_stop": False,
}
calc.model.loader_params = {
"batch_size": None,
"shuffle": False,
"num_workers": 0
}
calc.model.val_frac = 0.1
calc.model.structure = [2, 2]
# calc.model.optimizer = optim.Adam
# train the model
calc.train(overwrite=True)
parity_plot(calc, IMAGES, filename, data="forces")
parity_plot(calc, IMAGES, filename, data="energy")
md_run(IMAGES, count, calc, filename, dir, temp, const_t)
cores=1,
force_coefficient=0.3,
label=label,
save_logs=True,
lj_data=lj_data,
))
# define model settings
calc.model.device = "cpu"
calc.model.structure = [2, 2]
calc.model.val_frac = 0
calc.model.convergence = {
"energy": 0.02,
"force": 0.02,
"epochs": 1e10,
"early_stop": False,
}
calc.model.loader_params = {
"batch_size": None,
"shuffle": False,
"num_workers": 0
}
calc.model.criterion = CustomLoss
calc.model.optimizer = optim.LBFGS
calc.model.lr = 1e-2
calc.model.fine_tune = None
# train the model
calc.train(overwrite=True)
parity_plot(calc, images, data="energy", label=label)
parity_plot(calc, images, data="forces", label=label)
def test_training():
distances = np.linspace(2, 5, 100)
label = "test_training"
images = []
energies = []
forces = []
for l in distances:
image = Atoms(
"CuCO",
[
(-l * np.sin(0.65), l * np.cos(0.65), 0),
(0, 0, 0),
(l * np.sin(0.65), l * np.cos(0.65), 0),
],
)
image.set_cell([10, 10, 10])
image.wrap(pbc=True)
image.set_calculator(EMT())
images.append(image)
energies.append(image.get_potential_energy())
forces.append(image.get_forces())
energies = np.array(energies)
forces = np.concatenate(np.array(forces))
Gs = {}
Gs["G2_etas"] = np.logspace(np.log10(0.05), np.log10(5.0), num=2)
Gs["G2_rs_s"] = [0] * 2
Gs["G4_etas"] = [0.005]
Gs["G4_zetas"] = [1.0]
Gs["G4_gammas"] = [+1.0, -1]
Gs["cutoff"] = 6.5
torch.set_num_threads(1)
calc = AMP(model=AMPTorch(
images,
descriptor=SNN_Gaussian,
Gs=Gs,
force_coefficient=0.3,
label=label,
save_logs=True,
))
calc.model.device = "cpu"
calc.model.structure = [2, 2]
calc.model.val_frac = 0
calc.model.convergence = {
"energy": 0.005,
"force": 0.005,
"early_stop": False,
"epochs": 1e10,
}
calc.model.loader_params = {
"batch_size": None,
"shuffle": False,
"num_workers": 0
}
calc.model.criterion = CustomLoss
calc.model.optimizer = optim.LBFGS
calc.model.lr = 1e-2
calc.model.fine_tune = None
calc.train(overwrite=True)
train_hashes = list(calc.model.training_data.hashed_images.keys())
train_fp_hashes = {}
for hash in train_hashes:
with open("amp-data-fingerprints.ampdb/loose/" + hash, "rb") as f:
fp = load(f)
train_fp_hashes[hash] = fp
os.system("rm amp-data-fingerprints.ampdb/loose/" + hash)
train_prime_hashes = {}
for hash in train_hashes:
with open("amp-data-fingerprint-primes.ampdb/loose/" + hash,
"rb") as f:
prime = load(f)
train_prime_hashes[hash] = prime
os.system("rm amp-data-fingerprint-primes.ampdb/loose/" + hash)
dataset = TestDataset(
images,
descriptor=calc.model.descriptor,
unique_atoms=calc.model.training_data.elements,
Gs=calc.model.training_data.Gs,
fprange=calc.model.training_data.fprange,
)
test_hashes = list(dataset.hashed_images.keys())
test_fp_hashes = {}
for hash in test_hashes:
with open("amp-data-fingerprints.ampdb/loose/" + hash, "rb") as f:
fp = load(f)
test_fp_hashes[hash] = fp
os.system("rm amp-data-fingerprints.ampdb/loose/" + hash)
test_prime_hashes = {}
for hash in test_hashes:
with open("amp-data-fingerprint-primes.ampdb/loose/" + hash,
"rb") as f:
prime = load(f)
test_prime_hashes[hash] = prime
os.system("rm amp-data-fingerprint-primes.ampdb/loose/" + hash)
# test fingerprints are identical
for train_hash, test_hash in zip(train_hashes, test_hashes):
for train_fp, test_fp in zip(train_fp_hashes[train_hash],
test_fp_hashes[test_hash]):
for i, j in zip(train_fp[1], test_fp[1]):
assert abs(i - j) <= 1e-5, "Fingerprints do not match!"
# test fingerprint primes are identical
for train_hash, test_hash in zip(train_hashes, test_hashes):
for train_prime, test_prime in zip(
list(train_prime_hashes[train_hash].values()),
list(test_prime_hashes[test_hash].values()),
):
for i, j in zip(train_prime, test_prime):
assert abs(i - j) <= 1e-5, "Fingerprint primes do not match!"
num_of_atoms = 3
calculated_energies = np.array(
[calc.get_potential_energy(image) for image in images])
energy_rmse = np.sqrt(
(((calculated_energies - energies) / num_of_atoms)**2).sum() /
len(images))
assert (energy_rmse <= calc.model.convergence["energy"]
), "Energy training convergence not met!"
calculated_forces = np.concatenate(
np.array([calc.get_forces(image) for image in images]))
force_rmse = np.sqrt((((calculated_forces - forces))**2).sum() /
(3 * num_of_atoms * len(images)))
assert (force_rmse <= calc.model.convergence["force"]
), "Force training convergence not met!"
# test plot creation
parity_plot(calc, images, data="energy", label=label)
parity_plot(calc, images, data="forces", label=label)