def test_predict_on_structure_en(n_cpus, write_to_structure,
selective_atoms):
old_structure = deepcopy(_fake_structure)
old_structure.forces = np.random.uniform(-1, 1, (3, 3))
used_structure = deepcopy(old_structure)
forces, stds, energies = predict_on_structure_par_en(
structure=used_structure,
gp=_fake_gp,
n_cpus=n_cpus,
write_to_structure=write_to_structure,
selective_atoms=selective_atoms,
skipped_atom_value=0)
assert np.array_equal(forces.shape, old_structure.positions.shape)
assert np.array_equal(stds.shape, old_structure.positions.shape)
if write_to_structure:
if selective_atoms == [1]:
assert np.array_equal(old_structure.forces[0],
used_structure.forces[0])
assert np.array_equal(old_structure.forces[2],
used_structure.forces[2])
assert np.array_equal(used_structure.forces[1], forces[1])
else:
assert not np.array_equal(old_structure.forces[0],
used_structure.forces[0])
assert not np.array_equal(old_structure.forces[2],
used_structure.forces[2])
assert np.array_equal(forces, used_structure.forces)
# These will be unequal no matter what
assert not np.array_equal(old_structure.forces[1],
used_structure.forces[1])
else:
assert np.array_equal(old_structure.forces, used_structure.forces)
assert np.array_equal(forces.shape, (3, 3))
assert np.array_equal(stds.shape, (3, 3))
assert len(energies) == len(old_structure)
def calculate_gp(self, atoms):
nat = len(atoms)
struc_curr = Structure(np.array(atoms.cell),
atoms.get_atomic_numbers(), atoms.positions)
if self.par:
forces, stds, local_energies = \
predict_on_structure_par_en(struc_curr, self.gp_model)
else:
forces, stds, local_energies = \
predict_on_structure_en(struc_curr, self.gp_model)
self.results['forces'] = forces
self.results['stds'] = stds
self.results['local_energies'] = local_energies
self.results['energy'] = np.sum(local_energies)
atoms.get_uncertainties = self.get_uncertainties
return forces