def predict_on_structure_en( structure: Structure, gp: GaussianProcess, n_cpus: int = None) -> ('np.ndarray', 'np.ndarray', 'np.ndarray'): """ Return the forces/std. dev. uncertainty / local energy associated with each individual atom in a structure. Forces are stored directly to the structure and are also returned. :param structure: FLARE structure to obtain forces for, with N atoms :param gp: Gaussian Process model :param n_cpus: Dummy parameter passed as an argument to allow for flexibility when the callable may or may not be parallelized :return: N x 3 array of forces, N x 3 array of uncertainties, N-length array of energies :rtype: (np.ndarray, np.ndarray, np.ndarray) """ # Set up local energy array local_energies = np.array([0 for _ in range(structure.nat)]) # Loop through atoms in structure and predict forces, uncertainties, # and energies for n in range(structure.nat): chemenv = AtomicEnvironment(structure, n, gp.cutoffs) for i in range(3): force, var = gp.predict(chemenv, i + 1) structure.forces[n][i] = float(force) structure.stds[n][i] = np.sqrt(np.abs(var)) local_energies[n] = gp.predict_local_energy(chemenv) forces = np.array(structure.forces) stds = np.array(structure.stds) return forces, stds, local_energies
def predict_on_structure_en( structure: Structure, gp: GaussianProcess, n_cpus: int = None, write_to_structure: bool = True, selective_atoms: List[int] = None, skipped_atom_value=0) -> ('np.ndarray', 'np.ndarray', 'np.ndarray'): """ Return the forces/std. dev. uncertainty / local energy associated with each individual atom in a structure. Forces are stored directly to the structure and are also returned. :param structure: FLARE structure to obtain forces for, with N atoms :param gp: Gaussian Process model :param n_cpus: Dummy parameter passed as an argument to allow for flexibility when the callable may or may not be parallelized :return: N x 3 array of forces, N x 3 array of uncertainties, N-length array of energies :rtype: (np.ndarray, np.ndarray, np.ndarray) """ # Set up local energy array forces = np.zeros((structure.nat, 3)) stds = np.zeros((structure.nat, 3)) local_energies = np.zeros(structure.nat) forces = np.zeros(shape=(structure.nat, 3)) stds = np.zeros(shape=(structure.nat, 3)) if selective_atoms: forces.fill(skipped_atom_value) stds.fill(skipped_atom_value) local_energies.fill(skipped_atom_value) else: selective_atoms = [] # Loop through atoms in structure and predict forces, uncertainties, # and energies for n in range(structure.nat): if selective_atoms and n not in selective_atoms: continue chemenv = AtomicEnvironment(structure, n, gp.cutoffs, cutoffs_mask=gp.hyps_mask) for i in range(3): force, var = gp.predict(chemenv, i + 1) forces[n][i] = float(force) stds[n][i] = np.sqrt(np.abs(var)) if write_to_structure and structure.forces is not None: structure.forces[n][i] = float(force) structure.stds[n][i] = np.sqrt(np.abs(var)) local_energies[n] = gp.predict_local_energy(chemenv) return forces, stds, local_energies
def predict_on_atom_en(structure: Structure, atom: int, gp: GaussianProcess): chemenv = AtomicEnvironment(structure, atom, gp.cutoffs) comps = [] stds = [] # predict force components and standard deviations for i in range(3): force, var = gp.predict(chemenv, i + 1) comps.append(float(force)) stds.append(np.sqrt(np.abs(var))) # predict local energy local_energy = gp.predict_local_energy(chemenv) return comps, stds, local_energy
def predict_on_structure_en(structure: Structure, gp: GaussianProcess): local_energies = [0 for _ in range(structure.nat)] for n in range(structure.nat): chemenv = AtomicEnvironment(structure, n, gp.cutoffs) for i in range(3): force, var = gp.predict(chemenv, i + 1) structure.forces[n][i] = float(force) structure.stds[n][i] = np.sqrt(np.abs(var)) local_energies[n] = gp.predict_local_energy(chemenv) forces = np.array(structure.forces) stds = np.array(structure.stds) return forces, stds, local_energies
test_structure, forces = \ get_random_structure(np.eye(3), [1, 2], 3) energy = 3.14 gp_model.update_db(test_structure, forces, energy=energy) yield gp_model del gp_model _fake_gp = GaussianProcess(kernel_name='2+3', cutoffs=[5., 5.], hyps=[1., 1., 1., 1., 1.]) _fake_structure = Structure(cell=np.eye(3), species=[1, 1, 1], positions=np.random.uniform(0, 1, size=(3, 3))) _fake_gp.predict = fake_predict _fake_gp.predict_local_energy = fake_predict_local_energy assert isinstance(_fake_gp.predict(1, 1), tuple) assert isinstance(_fake_gp.predict_local_energy(1), float) @pytest.mark.parametrize('n_cpu', [None, 1, 2]) def test_predict_on_structure_par(n_cpu): # Predict only on the first atom, and make rest NAN selective_atoms = [0] skipped_atom_value = np.nan forces, stds = predict_on_structure_par( _fake_structure, _fake_gp, n_cpus=n_cpu, write_to_structure=False, selective_atoms=selective_atoms,