def __init__(self) -> None:
# Due to some mysterious bug in Sparrow, calculations get slower and slower over time.
# Therefore, we generate a new Sparrow object every time.
self.calculator = Sparrow('PM6')
self.settings = {
'molecular_charge': 0,
'max_scf_iterations': 128,
'unrestricted_calculation': 1,
}
self.atom_energies: Dict[str, float] = {}
def test_calculator(self):
calculator = Sparrow('PM6')
calculator.set_elements(list(self.atoms.symbols))
calculator.set_positions(self.atoms.positions)
calculator.set_settings({
'molecular_charge': 0,
'spin_multiplicity': 1
})
gradients = calculator.calculate_gradients()
energy = calculator.calculate_energy()
self.assertAlmostEqual(energy, -0.9379853016)
self.assertEqual(gradients.shape, (2, 3))
def calculate(self, atoms: Atoms, new_atom: Atom) -> Tuple[float, dict]:
start = time.time()
self.calculator = Sparrow('PM6')
all_atoms = atoms.copy()
all_atoms.append(new_atom)
e_tot = self._calculate_energy(all_atoms)
e_parts = self._calculate_energy(
atoms) + self._calculate_atomic_energy(new_atom)
delta_e = e_tot - e_parts
elapsed = time.time() - start
reward = -1 * delta_e
info = {
'elapsed_time': elapsed,
}
return reward, info
def test_energy_gradients(self):
calculator = Sparrow('PM6')
atoms = ase.io.read(filename=os.path.join(self.RESOURCES, 'h2o.xyz'),
format='xyz',
index=0)
calculator.set_positions(atoms.positions)
calculator.set_elements(list(atoms.symbols))
calculator.set_settings({
'molecular_charge': 0,
'spin_multiplicity': 1
})
energy = calculator.calculate_energy()
gradients = calculator.calculate_gradients()
energy_file = os.path.join(self.RESOURCES, 'energy.dat')
expected_energy = float(np.genfromtxt(energy_file))
self.assertAlmostEqual(energy, expected_energy)
gradients_file = os.path.join(self.RESOURCES, 'gradients.dat')
expected_gradients = np.genfromtxt(gradients_file)
self.assertTrue(np.allclose(gradients, expected_gradients))
class InteractionReward(MolecularReward):
def __init__(self) -> None:
# Due to some mysterious bug in Sparrow, calculations get slower and slower over time.
# Therefore, we generate a new Sparrow object every time.
self.calculator = Sparrow('PM6')
self.settings = {
'molecular_charge': 0,
'max_scf_iterations': 128,
'unrestricted_calculation': 1,
}
self.atom_energies: Dict[str, float] = {}
def calculate(self, atoms: Atoms, new_atom: Atom) -> Tuple[float, dict]:
start = time.time()
self.calculator = Sparrow('PM6')
all_atoms = atoms.copy()
all_atoms.append(new_atom)
e_tot = self._calculate_energy(all_atoms)
e_parts = self._calculate_energy(
atoms) + self._calculate_atomic_energy(new_atom)
delta_e = e_tot - e_parts
elapsed = time.time() - start
reward = -1 * delta_e
info = {
'elapsed_time': elapsed,
}
return reward, info
def _calculate_atomic_energy(self, atom: Atom) -> float:
if atom.symbol not in self.atom_energies:
atoms = Atoms()
atoms.append(atom)
self.atom_energies[atom.symbol] = self._calculate_energy(atoms)
return self.atom_energies[atom.symbol]
def _calculate_energy(self, atoms: Atoms) -> float:
if len(atoms) == 0:
return 0.0
self.calculator.set_elements(list(atoms.symbols))
self.calculator.set_positions(atoms.positions)
self.settings[
'spin_multiplicity'] = self.get_minimum_spin_multiplicity(atoms)
self.calculator.set_settings(self.settings)
return self.calculator.calculate_energy()
def test_minimize_fail(self):
calculator = Sparrow('PM6')
calculator.set_elements(list(self.atoms.symbols))
calculator.set_positions(self.atoms.positions)
calculator.set_settings({
'molecular_charge': self.charge,
'spin_multiplicity': self.spin_multiplicity
})
opt_atoms, success = minimize(
calculator=calculator,
atoms=self.atoms,
charge=self.charge,
spin_multiplicity=self.spin_multiplicity,
max_iter=1,
)
self.assertFalse(success)
def test_minimize_fixed(self):
calculator = Sparrow('PM6')
calculator.set_elements(list(self.atoms.symbols))
calculator.set_positions(self.atoms.positions)
calculator.set_settings({
'molecular_charge': self.charge,
'spin_multiplicity': self.spin_multiplicity
})
fixed_index = 2
opt_atoms, success = minimize(
calculator=calculator,
atoms=self.atoms,
charge=self.charge,
spin_multiplicity=self.spin_multiplicity,
fixed_indices=[fixed_index],
)
self.assertTrue(
np.all((self.atoms.positions -
opt_atoms.positions)[fixed_index] < 1E-6))
def calculate(self, atoms: Atoms, new_atom: Atom) -> Tuple[float, dict]:
start_time = time.time()
self.calculator = Sparrow('PM6')
all_atoms = atoms.copy()
all_atoms.append(new_atom)
e_tot = self._calculate_energy(all_atoms)
e_parts = self._calculate_energy(
atoms) + self._calculate_atomic_energy(new_atom)
delta_e = e_tot - e_parts
distance = np.linalg.norm(new_atom.position)
reward = -1 * (delta_e + self.distance_penalty * distance)
info = {
'elapsed_time': time.time() - start_time,
}
return reward, info
def test_minimize(self):
calculator = Sparrow('PM6')
calculator.set_elements(list(self.atoms.symbols))
calculator.set_positions(self.atoms.positions)
calculator.set_settings({
'molecular_charge': self.charge,
'spin_multiplicity': self.spin_multiplicity
})
energy1 = calculator.calculate_energy()
gradients1 = calculator.calculate_gradients()
opt_atoms, success = minimize(calculator=calculator,
atoms=self.atoms,
charge=self.charge,
spin_multiplicity=self.spin_multiplicity)
calculator.set_positions(opt_atoms.positions)
energy2 = calculator.calculate_energy()
gradients2 = calculator.calculate_gradients()
self.assertTrue(energy1 > energy2)
self.assertTrue(
np.sum(np.square(gradients1)) > np.sum(np.square(gradients2)))
self.assertTrue(np.all(gradients2 < 1E-3))
def test_atomic_energies(self):
calculator = Sparrow('PM6')
calculator.set_positions([(0, 0, 0)])
calculator.set_elements(['H'])
calculator.set_settings({
'molecular_charge': 0,
'spin_multiplicity': 2
})
self.assertAlmostEqual(calculator.calculate_energy(), -0.4133180865)
calculator.set_elements(['C'])
calculator.set_settings({
'molecular_charge': 0,
'spin_multiplicity': 1
})
self.assertAlmostEqual(calculator.calculate_energy(), -4.162353543)
calculator.set_elements(['O'])
calculator.set_settings({
'molecular_charge': 0,
'spin_multiplicity': 1
})
self.assertAlmostEqual(calculator.calculate_energy(), -10.37062419)