def train_calc(inputs):
images, filename, file_dir, Gs, lj, forcesonly, scaling = inputs
class train_end_load_best_valid_loss(skorch.callbacks.base.Callback):
def on_train_end(self, net, X, y):
net.load_params(
"./results/checkpoints/{}_params.pt".format(filename))
cp = Checkpoint(
monitor="forces_score_best",
fn_prefix="./results/checkpoints/{}_".format(filename),
)
if not os.path.exists(file_dir):
os.makedirs(file_dir, exist_ok=True)
forcetraining = True
training_data = AtomsDataset(
images,
SNN_Gaussian,
Gs,
forcetraining=forcetraining,
label=filename,
cores=1,
lj_data=None,
scaling=scaling,
)
unique_atoms = training_data.elements
fp_length = training_data.fp_length
device = "cpu"
torch.set_num_threads(1)
net = NeuralNetRegressor(
module=FullNN(unique_atoms, [fp_length, 3, 20],
device,
forcetraining=forcetraining),
criterion=CustomMSELoss,
criterion__force_coefficient=0.04,
optimizer=torch.optim.LBFGS,
lr=1e-1,
batch_size=len(training_data),
max_epochs=200,
iterator_train__collate_fn=collate_amp,
iterator_train__shuffle=False,
iterator_valid__collate_fn=collate_amp,
iterator_valid__shuffle=False,
device=device,
train_split=CVSplit(cv=5, random_state=1),
callbacks=[
EpochScoring(
forces_score,
on_train=False,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=False,
use_caching=True,
target_extractor=target_extractor,
),
],
)
calc = AMP(training_data, net, label=filename)
calc.train()
return [training_data, net, filename]
iterator_valid__collate_fn=collate_amp,
device=device,
# train_split=0,
verbose=1,
callbacks=[
EpochScoring(
forces_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
],
)
calc_amptorch = AMP(training_data, net, label=label)
calc_amptorch.train(overwrite=True)
# for evaluation in cpu
# net_load = copy.copy(net)
# calc_cpu = AMP(training_data, net_load, label)
# calc_cpu.load(filename='./results/trained_models/{}.pt'.format(label))
parity_plot(calc_amptorch, images, data="energy", label=label)
parity_plot(calc_amptorch, images, data="forces", label=label)
criterion__force_coefficient=0.3,
optimizer=torch.optim.LBFGS,
optimizer__line_search_fn="strong_wolfe",
lr=1e-1,
batch_size=len(images),
max_epochs=50,
iterator_train__collate_fn=collate_amp,
iterator_train__shuffle=False,
iterator_valid__collate_fn=collate_amp,
device=device,
train_split=0,
callbacks=[
EpochScoring(
forces_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
],
)
calc = AMP(training_data, net, label)
calc.train(overwrite=True)
parity_plot(calc, images, data="energy", label=label)
parity_plot(calc, images, data="forces", label=label)
def test_skorch():
distances = np.linspace(2, 5, 100)
label = "skorch_example"
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
forcetraining = True
training_data = AtomsDataset(
images,
SNN_Gaussian,
Gs,
forcetraining=forcetraining,
label=label,
cores=1,
delta_data=None,
)
unique_atoms = training_data.elements
fp_length = training_data.fp_length
device = "cpu"
net = NeuralNetRegressor(
module=FullNN(unique_atoms, [fp_length, 2, 2],
device,
forcetraining=forcetraining),
criterion=CustomMSELoss,
criterion__force_coefficient=0.3,
optimizer=torch.optim.LBFGS,
optimizer__line_search_fn="strong_wolfe",
lr=1,
batch_size=100,
max_epochs=150,
iterator_train__collate_fn=collate_amp,
iterator_train__shuffle=True,
iterator_valid__collate_fn=collate_amp,
device=device,
train_split=0,
verbose=0,
callbacks=[
EpochScoring(
forces_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
],
)
calc = AMP(training_data, net, "test")
calc.train(overwrite=True)
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))
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)))
l1_force = np.sum(np.abs(calculated_forces - forces) / num_of_atoms, 1)
idx = 0
force_loss_image = np.zeros((len(calculated_energies), 1))
for i in range(len(calculated_energies)):
force_loss_image[i] = np.sum(l1_force[idx:idx + 3])
idx += 3
force_loss_image /= 3
reported_energy_score = net.history[-1]["energy_score"]
reported_forces_score = net.history[-1]["forces_score"]
assert force_rmse <= 0.005, "Force training convergence not met!"
assert energy_rmse <= 0.005, "Energy training convergence not met!"
assert round(reported_energy_score,
4) == round(energy_rmse,
4), "Shuffled reported energy scores incorrect!"
assert round(reported_forces_score,
4) == round(force_rmse,
4), "Shuffled reported forces score incorrect!"
def test_skorch_delta():
from amptorch.skorch_model import AMP
cp = Checkpoint(monitor="valid_loss_best", fn_prefix="valid_best_")
distances = np.linspace(2, 5, 10)
label = "skorch_example"
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())
image = Atoms("CuC", [(-1, 1, 0), (1, 1, 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
params = {
"C": {"re": 0.972, "D": 6.379, "sig": 0.477},
"O": {"re": 1.09, "D": 8.575, "sig": 0.603},
"Cu": {"re": 2.168, "D": 3.8386, "sig": 1.696},
}
morse_model = morse_potential(images, params, Gs["cutoff"], label)
morse_energies, morse_forces, num_atoms = morse_model.morse_pred(
images, params)
morse_data = [morse_energies, morse_forces, num_atoms, params, morse_model]
forcetraining = True
training_data = AtomsDataset(
images,
SNN_Gaussian,
Gs,
forcetraining=forcetraining,
label=label,
cores=1,
delta_data=morse_data,
)
batch_size = len(training_data)
unique_atoms = training_data.elements
fp_length = training_data.fp_length
device = "cpu"
net = NeuralNetRegressor(
module=FullNN(
unique_atoms, [fp_length, 2, 2], device, forcetraining=forcetraining
),
criterion=CustomMSELoss,
criterion__force_coefficient=0.3,
optimizer=torch.optim.LBFGS,
optimizer__line_search_fn="strong_wolfe",
lr=1e-2,
batch_size=batch_size,
max_epochs=100,
iterator_train__collate_fn=collate_amp,
iterator_train__shuffle=False,
iterator_valid__collate_fn=collate_amp,
device=device,
train_split=0,
verbose=0,
callbacks=[
EpochScoring(
forces_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
],
)
calc = AMP(training_data, net, "test")
calc.train(overwrite=True)
num_of_atoms = 3
calculated_energies = np.array(
[calc.get_potential_energy(image) for idx, image in enumerate(images)]
)
energy_rmse = np.sqrt(
(((calculated_energies - energies) / num_of_atoms) ** 2).sum() / len(images)
)
last_energy_score = net.history[-1]["energy_score"]
assert round(energy_rmse, 4) == round(
last_energy_score, 4
), "Energy errors incorrect!"
last_forces_score = net.history[-1]["forces_score"]
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 round(force_rmse, 4) == round(
last_forces_score, 4
), "Force errors incorrect!"
def test_skorch_val():
distances = np.linspace(2, 5, 100)
label = "example"
images = []
energies = []
forces = []
for l in distances:
image = Atoms(
"CuCCu",
[
(-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)
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
forcetraining = True
training_data = AtomsDataset(
images,
SNN_Gaussian,
Gs,
forcetraining=forcetraining,
label=label,
cores=1,
delta_data=None,
)
batch_size = len(training_data)
unique_atoms = training_data.elements
fp_length = training_data.fp_length
device = "cpu"
net = NeuralNetRegressor(
module=FullNN(unique_atoms, [fp_length, 2, 2],
device,
forcetraining=forcetraining),
criterion=CustomMSELoss,
criterion__force_coefficient=0.3,
optimizer=torch.optim.LBFGS,
optimizer__line_search_fn="strong_wolfe",
lr=1e-2,
batch_size=10,
max_epochs=20,
iterator_train__collate_fn=collate_amp,
iterator_train__shuffle=True,
iterator_valid__collate_fn=collate_amp,
device=device,
train_split=CVSplit(0.1, random_state=1),
verbose=0,
callbacks=[
EpochScoring(
forces_score,
on_train=False,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=False,
use_caching=True,
target_extractor=target_extractor,
),
],
)
val_indices = [80, 84, 33, 81, 93, 17, 36, 82, 69, 65]
val_images = [images[idx] for idx in val_indices]
val_energies = energies[val_indices]
val_forces = np.concatenate(np.array([forces[idx] for idx in val_indices]))
calc = AMP_skorch(training_data, net, "test")
calc.train(overwrite=True)
num_of_atoms = 3
last_energy_score = net.history[-1]["energy_score"]
last_forces_score = net.history[-1]["forces_score"]
calculated_energies = np.array(
[calc.get_potential_energy(image) for image in val_images])
energy_rmse = np.sqrt(
(((calculated_energies - val_energies) / num_of_atoms)**2).sum() /
len(val_images))
assert round(energy_rmse,
5) == round(last_energy_score,
5), "Validation energy errors incorrect!"
calculated_forces = np.concatenate(
np.array([calc.get_forces(image) for image in val_images]))
force_rmse = np.sqrt((((calculated_forces - val_forces))**2).sum() /
(3 * num_of_atoms * len(val_images)))
assert round(force_rmse,
5) == round(last_forces_score,
5), "Validation force errors incorrect!"
def test_load():
distances = np.linspace(2, 5, 100)
label = "example"
images = []
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)
# define symmetry functions to be used
Gs = {}
Gs["G2_etas"] = np.logspace(np.log10(0.05), np.log10(5.0), num=4)
Gs["G2_rs_s"] = [0] * 4
Gs["G4_etas"] = [0.005]
Gs["G4_zetas"] = [1.0]
Gs["G4_gammas"] = [+1.0, -1]
Gs["cutoff"] = 6.5
forcetraining = True
training_data = AtomsDataset(images, SNN_Gaussian, Gs, forcetraining=forcetraining,
label=label, cores=1, delta_data=None)
unique_atoms = training_data.elements
fp_length = training_data.fp_length
device = "cpu"
net = NeuralNetRegressor(
module=FullNN(unique_atoms, [fp_length, 3, 10], device, forcetraining=forcetraining),
criterion=CustomMSELoss,
criterion__force_coefficient=0.3,
optimizer=torch.optim.LBFGS,
optimizer__line_search_fn="strong_wolfe",
lr=1e-1,
batch_size=len(images),
max_epochs=5,
iterator_train__collate_fn=collate_amp,
iterator_train__shuffle=False,
iterator_valid__collate_fn=collate_amp,
device=device,
train_split=0,
verbose=0,
callbacks=[
EpochScoring(
forces_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
EpochScoring(
energy_score,
on_train=True,
use_caching=True,
target_extractor=target_extractor,
),
],
)
net_2 = copy.copy(net)
calc_1 = AMP(training_data, net, 'test')
calc_1.train(overwrite=True)
energy_1 = calc_1.get_potential_energy(images[0])
forces_1 = calc_1.get_forces(images[0])
calc_2 = AMP(training_data, net_2, 'test')
calc_2.load(filename='./results/trained_models/test.pt')
energy_2 = calc_2.get_potential_energy(images[0])
forces_2 = calc_2.get_forces(images[0])
assert energy_1 == energy_2, "Energies do not match!"
assert (forces_1 == forces_2).all(), "Forces do not match!"