Exemplo n.º 1
0
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!"
Exemplo n.º 2
0
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_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_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!"