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)
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_force_xyz(chemenv)
std = np.sqrt(np.abs(var))
forces[n] = force
stds[n] = std
if write_to_structure and structure.forces is not None:
structure.forces[n] = force
structure.stds[n] = std
local_energies[n] = gp.predict_local_energy(chemenv)
return forces, stds, local_energies
def predict_on_structure(
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"):
"""
Return the forces/std. dev. uncertainty 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 write_to_structure: Write results to structure's forces,
std attributes
:param selective_atoms: Only predict on these atoms; e.g. [0,1,2] will
only predict and return for those atoms
:param skipped_atom_value: What value to use for atoms that are skipped.
Defaults to 0 but other options could be e.g. NaN. Will NOT
write this to the structure if write_to_structure is True.
:return: N x 3 numpy array of foces, Nx3 numpy array of uncertainties
:rtype: (np.ndarray, np.ndarray)
"""
forces = np.zeros((structure.nat, 3))
stds = np.zeros((structure.nat, 3))
if selective_atoms:
forces.fill(skipped_atom_value)
stds.fill(skipped_atom_value)
else:
selective_atoms = []
for n in range(structure.nat):
# Skip the atoms which we aren't predicting on if
# selective atoms is on.
if n not in selective_atoms and selective_atoms:
continue
chemenv = AtomicEnvironment(structure, n, gp.cutoffs, cutoffs_mask=gp.hyps_mask)
force, var = gp.predict_force_xyz(chemenv)
std = np.sqrt(np.abs(var))
forces[n] = force
stds[n] = std
if write_to_structure:
structure.forces[n] = force
structure.stds[n] = std
return forces, stds
