import numpy as np
try:
from xtb.interface import Environment, Param, Calculator as XTBCalculator
from xtb.libxtb import VERBOSITY_MINIMAL, VERBOSITY_FULL, VERBOSITY_MUTED
except ModuleNotFoundError:
can_pyxtb = False
from pysisyphus.calculators.Calculator import Calculator
from pysisyphus.elem_data import ATOMIC_NUMBERS
[docs]
class PyXTB(Calculator):
def __init__(self, *args, gfn=2, acc=None, verbosity=0, keep_calculator=False, **kwargs):
super().__init__(*args, **kwargs)
self.env = Environment()
self.gfn = str(gfn)
self.acc = acc
avail_verbosities = {
0: VERBOSITY_MUTED,
1: VERBOSITY_MINIMAL,
2: VERBOSITY_FULL,
}
self.verbosity = avail_verbosities[verbosity]
avail_params = {
"0": Param.GFN0xTB,
"1": Param.GFN1xTB,
"2": Param.GFN2xTB,
"ff": Param.GFNFF,
}
self.keep_calculator = keep_calculator
self._calculator = None
self.param = avail_params[self.gfn]
self.uhf = self.mult - 1
[docs]
def get_calculator(self, atoms, coords):
# Reuse old calculator and updatec coordinates
if self._calculator:
self._calculator.update(coords.reshape(-1, 3))
return self._calculator
# Create new calculator
numbers = np.array([ATOMIC_NUMBERS[atom.lower()] for atom in atoms])
calc = XTBCalculator(
self.param, numbers, coords.reshape(-1, 3), charge=self.charge, uhf=self.uhf
)
calc.set_verbosity(self.verbosity)
if self.acc:
calc.set_accuracy(self.acc)
# Keep calculator, if requested
if self.keep_calculator and (self._calculator is None):
self._calculator = calc
return calc
[docs]
def get_energy(self, atoms, coords, **prepare_kwargs):
calc = self.get_calculator(atoms, coords)
xtb_result = calc.singlepoint(copy=True)
results = {
"energy": xtb_result.get_energy(),
}
return results
[docs]
def get_forces(self, atoms, coords, **prepare_kwargs):
calc = self.get_calculator(atoms, coords)
xtb_result = calc.singlepoint(copy=True)
results = {
"energy": xtb_result.get_energy(),
"forces": -xtb_result.get_gradient().flatten(),
}
return results