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,
