class TestSchnetPack:
pos1 = Posinp.from_file(os.path.join(pos_folder, "N2.xyz"))
model1 = os.path.join(model_folder, "ani1_N2_model")
model2 = os.path.join(model_folder, "wacsf_model")
model3 = os.path.join(model_folder, "H2O_model")
calc1 = SchnetPackCalculator(model_dir=model1)
calc2 = SchnetPackCalculator(model_dir=model2)
calc3 = SchnetPackCalculator(model_dir=model3)
job = Job(posinp=pos1, calculator=calc1)
jobwacsf = Job(posinp=pos1, calculator=calc2)
def test_only_energy(self):
self.job.run("energy")
assert np.array(-2979.6067) == self.job.results["energy"][0]
new_job = deepcopy(self.job)
new_job.calculator.units = {"energy": "hartree", "positions": "atomic"}
new_job.run("energy")
assert new_job.results["energy"][0] == np.array(-81079.228)
self.job.results["energy"] = None
def test_forces_from_deriv(self):
assert self.job.calculator.available_properties == ["energy"]
self.job.run("forces")
ref_forces = np.array([[-0.0, -0.0, -0.31532133],
[-0.0, -0.0, 0.31532133]])
assert np.isclose(self.job.results["forces"][0], ref_forces).all()
new_job = deepcopy(self.job)
new_job.calculator.units = {"energy": "hartree", "positions": "atomic"}
new_job.run("forces")
assert np.isclose(ref_forces * 51.42206334724,
new_job.results["forces"][0]).all()
def test_forces_from_finite_difference(self):
assert self.job.calculator.available_properties == ["energy"]
self.job.run("forces", finite_difference=True)
ref_forces = np.array([[-0.0, -0.0, -0.32416448],
[-0.0, -0.0, 0.32416448]])
assert np.isclose(self.job.results["forces"][0], ref_forces).all()
def test_wacsf(self):
self.jobwacsf.run("energy")
assert self.jobwacsf.results["energy"] is not None
def test_bad_property(self):
with pytest.raises(ValueError):
self.job.run("dipole")
def test_convert(self):
pos_angstrom = Posinp.from_file(
os.path.join(pos_folder, "H2O_unrelaxed.xyz"))
pos_atomic = Posinp.from_file(
os.path.join(pos_folder, "H2O_atomic.xyz"))
job1 = Job(posinp=pos_angstrom, calculator=self.calc3)
job2 = Job(posinp=pos_atomic, calculator=self.calc3)
job1.run("energy")
job2.run("energy")
assert job1.results["energy"] == job2.results["energy"]
def __init__(self,
model_dir,
available_properties=None,
device="cpu",
**kwargs):
Calculator.__init__(self, **kwargs)
self.schnetpackcalculator = SchnetPackCalculator(
model_dir=model_dir,
available_properties=available_properties,
device=device,
)
self.implemented_properties = self.schnetpackcalculator._get_available_properties(
)
if "energy" in self.implemented_properties and "forces" not in self.implemented_properties:
self.implemented_properties.append("forces")
class AseSpkCalculator(Calculator):
def __init__(self,
model_dir,
available_properties=None,
device="cpu",
**kwargs):
Calculator.__init__(self, **kwargs)
self.schnetpackcalculator = SchnetPackCalculator(
model_dir=model_dir,
available_properties=available_properties,
device=device,
)
self.implemented_properties = self.schnetpackcalculator._get_available_properties(
)
if "energy" in self.implemented_properties and "forces" not in self.implemented_properties:
self.implemented_properties.append("forces")
def calculate(self,
atoms=None,
properties=["energy"],
system_changes=all_changes):
if self.calculation_required(atoms, properties):
Calculator.calculate(self, atoms)
posinp = Posinp.from_ase(atoms)
job = Job(posinp=posinp, calculator=self.schnetpackcalculator)
for prop in properties:
job.run(prop)
results = {}
for prop, result in zip(job.results.keys(), job.results.values()):
results[prop] = np.squeeze(result)
self.results = results
class TestGeopt:
posN2 = Posinp.from_file(os.path.join(pos_folder, "N2_unrelaxed.xyz"))
posH2O = Posinp.from_file(os.path.join(pos_folder, "H2O_unrelaxed.xyz"))
calcN2 = SchnetPackCalculator(os.path.join(model_folder, "ani1_N2_model"))
calcH2O = SchnetPackCalculator(os.path.join(model_folder, "H2O_model"))
def test_N2(self, capsys):
init = Job(posinp=self.posN2, calculator=self.calcN2)
init.run("energy")
assert init.results["energy"][0] == -2978.2354
geo = Geopt(posinp=self.posN2, calculator=self.calcN2)
geo.run(recenter=True, verbose=2)
output = capsys.readouterr()
assert 1.101 < geo.final_posinp.distance(0, 1) < 1.103
assert np.isclose(
geo.final_posinp.positions[0, 2],
-geo.final_posinp.positions[1, 2],
atol=1e-6,
)
assert output.out is not None
final = Job(posinp=geo.final_posinp, calculator=self.calcN2)
final.run("energy")
assert final.results["energy"][0] == -2979.6070
def test_H2O(self):
init = Job(posinp=self.posH2O, calculator=self.calcH2O)
init.run("energy")
assert np.isclose(init.results["energy"][0], -2076.7576, atol=1e-4)
geo = Geopt(posinp=self.posH2O, calculator=self.calcH2O)
geo.run()
assert 0.964 < geo.final_posinp.distance(0, 1) < 0.965
assert (geo.final_posinp.distance(0, 1) -
geo.final_posinp.distance(0, 2)) < 0.0001
# TODO Check the angle
final = Job(posinp=geo.final_posinp, calculator=self.calcH2O)
final.run("energy")
assert np.isclose(final.results["energy"][0], -2078.6301, atol=1e-4)
def test_errors(self):
with pytest.raises(ValueError):
geo = Geopt(posinp=None, calculator=self.calcN2)
with pytest.raises(TypeError):
geo = Geopt(posinp=self.posN2, calculator=1)
class TestPhononAutoGrad:
posH2O = Posinp.from_file(os.path.join(pos_folder, "H2Orelaxed.xyz"))
calc_ener = SchnetPackCalculator(os.path.join(model_folder, "H2O_model"))
calc_for = SchnetPackCalculator(
os.path.join(model_folder, "H2O_forces_model"))
def test_ph_h2o_autograd_2nd_derivative(self):
ph1 = Phonon(posinp=self.posH2O, calculator=self.calc_ener)
ph1.run()
ph1.energies.sort()
assert np.allclose(ph1.energies[6:9], [1726, 3856, 3942], atol=1)
def test_ph_h2o_autograd_1st_derivative(self):
ph1 = Phonon(posinp=self.posH2O, calculator=self.calc_for)
ph1.run()
ph1.energies.sort()
assert np.allclose(ph1.energies[6:9], [1589, 3703, 3812], atol=1)
def __init__(self,
model_dir,
available_properties=None,
device="cpu",
md=False,
dropout=False,
**kwargs):
r"""
Parameters
----------
model_dir : str
Same as :class:`SchnetPackCalculator`.
available_properties : str or list of str
Same as :class:`SchnetPackCalculator`.
device : str
Same as :class:`SchnetPackCalculator`.
md : bool
Default is False. Should be set to True if the
calculator is used for molecular dynamics.
dropout : bool
Same as :class:`SchnetPackCalculator`.
units : dict
Same as :class:`SchnetPackCalculator`.
"""
Calculator.__init__(self, **kwargs)
self.schnetpackcalculator = SchnetPackCalculator(
model_dir=model_dir,
available_properties=available_properties,
device=device,
md=md,
dropout=dropout,
)
self.implemented_properties = (
self.schnetpackcalculator._get_available_properties())
if ("energy" in self.implemented_properties
and "forces" not in self.implemented_properties):
self.implemented_properties.append("forces")
def get_orthonormal_basis(modes=None, model=None, posinp=None, write=False):
if modes:
with h5py.File(modes, "r") as f:
normal_modes = f["modes"][:]
orthonormal_basis = gramschmidt(normal_modes)
else:
calculator = SchnetPackCalculator(model)
posinp = Posinp.from_file(posinp)
ph = Phonon(posinp=posinp, calculator=calculator)
ph.run()
normal_modes = ph.normal_modes.copy()
orthonormal_basis = gramschmidt(normal_modes)
if write:
fx = h5py.File("orthonormal_basis.h5", "w")
fx.create_dataset("basis", data=orthonormal_basis)
fx.close()
return orthonormal_basis
class TestGrapheneAngles:
graphene2_name = os.path.join(pos_folder, "gra2.xyz")
pos_2at = Posinp.from_file(graphene2_name)
graphene4_name = os.path.join(pos_folder, "gra4_red.xyz")
pos_4at_red = Posinp.from_file(graphene4_name)
calc = SchnetPackCalculator(os.path.join(model_folder, "H2O_model"))
def test_same_result(self):
j1 = Job(posinp=self.pos_2at, calculator=self.calc)
j1.run("energy")
e1 = j1.results["energy"]
j2 = Job(posinp=self.pos_4at_red, calculator=self.calc)
j2.run("energy")
e2 = j2.results["energy"]
pos_4at = deepcopy(self.pos_4at_red)
pos_4at.convert_units("angstroem")
j3 = Job(posinp=pos_4at, calculator=self.calc)
j3.run("energy")
e3 = j3.results["energy"]
assert e2 == e3
assert np.isclose(e2, 2 * e1, atol=0.0002)
def test_cell_with_angles(self):
assert np.isclose(
self.pos_2at.cell.array,
np.array([[2.46, 0.0, 0.0], [0.0, 0.0, 0.0],
[1.23, 0.0, 2.13042249]]),
).all()
def test_verify_values(self):
assert np.isclose(self.pos_2at.angles, np.array([90.0, 60.0,
90.0])).all()
assert self.pos_4at_red.orthorhombic
assert not self.pos_2at.orthorhombic
assert np.isclose(self.pos_2at.cell.lengths(),
np.array([2.46, 0.0, 2.46])).all()
class TestPhononFinite:
posN2 = Posinp.from_file(os.path.join(pos_folder, "N2_unrelaxed.xyz"))
calcN2 = SchnetPackCalculator(os.path.join(model_folder, "myN2_model"))
def test_ph_N2(self):
ph1 = Phonon(posinp=self.posN2,
calculator=self.calcN2,
finite_difference=True)
ph1.run(batch_size=1)
assert np.isclose(ph1.energies.max(), 2339.53, atol=0.01)
assert all(
np.abs(np.delete(ph1.energies, np.argmax(ph1.energies))) < 30)
ph2 = Phonon(
posinp=ph1._ground_state,
calculator=ph1.calculator,
relax=False,
finite_difference=True,
translation_amplitudes=0.03,
)
ph2.run()
assert np.allclose(ph1.energies, ph2.energies)
ph3 = Phonon(
posinp=self.posN2,
calculator=self.calcN2,
finite_difference=True,
relax=False,
)
ph3.run()
assert not np.allclose(ph3.energies, ph1.energies, atol=100)
def test_phonon_posinp_error(self):
with pytest.raises(TypeError):
ph = Phonon(posinp=None, calculator=self.calcN2)
def test_phonon_calc_error(self):
with pytest.raises(TypeError):
ph = Phonon(posinp=self.posN2, calculator=None)
def get_normal_modes(posinp, model, device="cpu", rotate=False):
calculator = SchnetPackCalculator(model, device=device)
posinp = Posinp.from_file(posinp)
ph = Phonon(posinp=posinp, calculator=calculator)
ph.run()
if rotate:
ndim = 3 * len(posinp)
final_modes = np.zeros((ndim, ndim))
idx = np.arange(ndim)
for i in range(ndim):
init = ph.normal_modes[:, i]
final = np.empty_like(init)
final[np.where(idx % 3 == 0)[0]] = init[np.where(idx % 3 == 0)[0]]
final[np.where(idx % 3 == 1)[0]] = init[np.where(idx % 3 == 2)[0]]
final[np.where(idx %
3 == 2)[0]] = -1.0 * init[np.where(idx % 3 == 1)[0]]
final_modes[:, i] = final
return ph.energies, final_modes
else:
return ph.energies, ph.normal_modes
class AseSpkCalculator(Calculator):
r"""
Wrapper :class:`Calculator` class around the :class:`SchnetPackCalculator`
class to use directly inside ASE funtions.
"""
def __init__(self,
model_dir,
available_properties=None,
device="cpu",
md=False,
dropout=False,
**kwargs):
r"""
Parameters
----------
model_dir : str
Same as :class:`SchnetPackCalculator`.
available_properties : str or list of str
Same as :class:`SchnetPackCalculator`.
device : str
Same as :class:`SchnetPackCalculator`.
md : bool
Default is False. Should be set to True if the
calculator is used for molecular dynamics.
dropout : bool
Same as :class:`SchnetPackCalculator`.
units : dict
Same as :class:`SchnetPackCalculator`.
"""
Calculator.__init__(self, **kwargs)
self.schnetpackcalculator = SchnetPackCalculator(
model_dir=model_dir,
available_properties=available_properties,
device=device,
md=md,
dropout=dropout,
)
self.implemented_properties = (
self.schnetpackcalculator._get_available_properties())
if ("energy" in self.implemented_properties
and "forces" not in self.implemented_properties):
self.implemented_properties.append("forces")
def calculate(self,
atoms=None,
properties=["energy"],
system_changes=all_changes):
r"""
This method will be called by ASE functions.
"""
if self.calculation_required(atoms, properties):
from mlcalcdriver.base.posinp import Posinp
Calculator.calculate(self, atoms)
posinp = Posinp.from_ase(atoms)
from mlcalcdriver.base.job import Job
job = Job(posinp=posinp, calculator=self.schnetpackcalculator)
for prop in properties:
job.run(prop)
results = {}
for prop, result in zip(job.results.keys(), job.results.values()):
results[prop] = np.squeeze(result)
self.results = results
def get_normal_modes(posinp, model, device="cpu"):
calculator = SchnetPackCalculator(model, device=device)
posinp = Posinp.from_file(posinp)
ph = Phonon(posinp=posinp, calculator=calculator)
ph.run()
return ph.energies, ph.normal_modes