# This file is part of tblite.
# SPDX-Identifier: LGPL-3.0-or-later
#
# tblite is free software: you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# tblite is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with tblite. If not, see <https://www.gnu.org/licenses/>.
"""
The Python API of *tblite* natively support integration with the atomic
simulation environment (`ASE`_).
By constructing a calculator most functionality of ASE is readily available.
For details on building the Python API checkout the
:ref:`installation guide <python-build>`.
.. _ase: https://wiki.fysik.dtu.dk/ase/
Example
-------
The calculator can be used like other ASE calculators by attaching it to an atoms object.
>>> from ase.build import molecule
>>> from tblite.ase import TBLite
>>> atoms = molecule("H2O")
>>> atoms.calc = TBLite(method="GFN2-xTB")
>>> atoms.get_potential_energy() # in eV
-137.96777594361677
The calculator will perform single point calculations as necessary, if properties for
the same atoms are requested the properties will be returned without additional single point
calculations.
>>> atoms.get_dipole_moment() # in Debye
array([-7.83154860e-17, 1.10157136e-16, -4.76020974e-01])
>>> atoms.get_forces() # in eV/A
array([[-3.31859412e-15, 7.83871951e-15, -7.49527168e-01],
[ 3.08351626e-15, -1.53506496e-01, 3.74763584e-01],
[ 2.35077857e-16, 1.53506496e-01, 3.74763584e-01]])
Changed properties of the atoms object will lead to an reevalution of the single point,
for example updating the initial charges of the atoms to set a total charge.
>>> atoms.set_initial_charges([0, 0, -1])
>>> atoms.get_potential_energy() # in eV
-131.50015843891816
Some properties, like the total charge, can also be set via the parameters of the
calculator.
>>> atoms = molecule("H2O")
>>> atoms.calc = TBLite(method="GFN2-xTB", charge=-1)
>>> atoms.get_potential_energy() # in eV
-131.50015843891822
Finally, the parameters of the calculator can be updated at any time using the *set*
method, which resets the calculation as necessary.
>>> atoms.calc.set(method="GFN1-xTB", charge=0)
>>> atoms.get_potential_energy() # in eV
-156.96750590276173
>>> atoms.get_dipole_moment() # in Debye
array([ 3.71500789e-16, 3.08439980e-16, -5.97048616e-01])
>>> atoms.get_forces() # in eV/A
array([[ 2.88997963e-15, 5.70620869e-15, -7.84252969e-01],
[-4.90812472e-15, -2.28835628e-01, 3.92126485e-01],
[ 2.01814509e-15, 2.28835628e-01, 3.92126485e-01]])
The calculator also support atom partitioned properties, including energies and charges.
>>> atoms.get_charges() # in e
array([-0.66558478, 0.33279239, 0.33279239])
>>> atoms.get_potential_energies() # in eV
array([-132.37935247, -12.29407672, -12.29407672])
"""
try:
import ase
import ase.calculators.calculator
from ase.atoms import Atoms
from ase.units import Bohr, Hartree, kB
except ModuleNotFoundError as e:
raise ModuleNotFoundError("This submodule requires ASE installed") from e
from typing import Any, Dict, List, Optional
from .interface import Calculator
[docs]
class TBLite(ase.calculators.calculator.Calculator):
r"""
ASE calculator for using xTB Hamiltonians from the tblite library.
Supported properties by this calculator are:
- *energy* (*free_energy*): The total energy of the system (in eV)
- *energies*: The atom-partitioned energy of the system (in eV)
- *forces*: The derivative of the energy with respect to the atomic positions (in eV/Ang)
- *stress*: The derivative of the energy with respect to cell deformations (PBC only, in eV/Ang³)
- *dipole*: The dipole moment of the system (in Debye)
- *charges*: Mulliken charges of the system (in e)
Supported keywords are
======================== ================= ====================================================
Keyword Default Description
======================== ================= ====================================================
method "GFN2-xTB" Underlying method for energy and forces
charge None Total charge of the system
multiplicity None Total multiplicity of the system
accuracy 1.0 Numerical accuracy of the calculation
electronic_temperature 300.0 Electronic temperatur in Kelvin
max_iterations 250 Iterations for self-consistent evaluation
initial_guess "sad" Initial guess for wavefunction (sad, eeq, or eeqbc)
mixer_damping 0.4 Damping parameter for self-consistent mixer
electric_field None Uniform electric field vector (in V/A)
spin_polarization None Spin polarization (scaling factor)
solvation None Solvation model to use (see below for details)
cache_api True Reuse generate API objects (recommended)
verbosity 1 Set verbosity of printout
xtb_config None Additional configuration for the API calculator (dict)
======================== ================= ====================================================
Solvation models are supported by passing a tuple to the *solvation* keyword.
The first element of the tuple is the name of the solvation model, the following
arguments are passed to initialize the solvation model.
================== ====================================================
Solvation model Arguments
================== ====================================================
"alpb" Solvent name (str), solution state (optional)
"gbsa" Solvent name (str), solution state (optional)
"cpcm" Epsilon (float)
"gbe" Epsilon (float), Born kernel
"gb" Epsilon (float), Born kernel
================== ====================================================
Example
-------
An ASE calculator can be constructed by using the *TBLite* class provided
by the *tblite.ase* module.
For example to perform a single point calculation for a CO\ :sub:`2`
crystal use
>>> from tblite.ase import TBLite
>>> from ase.atoms import Atoms
>>> import numpy as np
>>> atoms = Atoms(
... symbols="C4O8",
... positions=np.array(
... [
... [0.9441259872, 0.9437851680, 0.9543505632],
... [3.7179966528, 0.9556570368, 3.7316862240],
... [3.7159517376, 3.7149292800, 0.9692330016],
... [0.9529872864, 3.7220864832, 3.7296981120],
... [1.6213905408, 1.6190616096, 1.6313879040],
... [0.2656685664, 0.2694175776, 0.2776540416],
... [4.3914553920, 1.6346256864, 3.0545920000],
... [3.0440834880, 0.2764611744, 4.4080419264],
... [4.3910577696, 3.0416409504, 0.2881058304],
... [3.0399936576, 4.3879335936, 1.6497353376],
... [0.2741322432, 4.4003734944, 3.0573754368],
... [1.6312174944, 3.0434586528, 4.4023048032],
... ]
... ),
... cell=np.array([5.68032, 5.68032, 5.68032]),
... pbc=np.array([True, True, True]),
... )
>>> atoms.calc = TBLite(method="GFN1-xTB", verbosity=0)
>>> atoms.get_potential_energy() # result in eV
-1257.0943962462964
The resulting calculator can be used like most ASE calculator, *e.g.* for optimizing geometries.
"""
implemented_properties = [
"energy",
"energies",
"forces",
"charges",
"dipole",
"stress",
]
default_parameters = {
"method": "GFN2-xTB",
"charge": None,
"multiplicity": None,
"accuracy": 1.0,
"guess": "sad",
"max_iterations": 250,
"mixer_damping": 0.4,
"electric_field": None,
"spin_polarization": None,
"solvation": None,
"electronic_temperature": 300.0,
"cache_api": True,
"verbosity": 1,
"xtb_config": None,
}
_res = None
_xtb = None
def __init__(
self,
atoms: Optional[Atoms] = None,
**kwargs,
):
"""
Construct the TBLite base calculator object.
"""
ase.calculators.calculator.Calculator.__init__(self, atoms=atoms, **kwargs)
[docs]
def set(self, **kwargs) -> dict:
"""
Set new parameters to TBLite. Will automatically reconstruct the underlying
model in case critical parameters change.
Example
-------
>>> from ase.build import molecule
>>> from tblite.ase import TBLite
>>> atoms = molecule("H2O")
>>> atoms.calc = TBLite(method="GFN2-xTB")
>>> atoms.get_potential_energy()
-137.96777625229421
>>> atoms.calc.set(method="GFN1-xTB")
{'method': 'GFN1-xTB'}
>>> atoms.get_potential_energy()
-156.9675057724589
"""
_update_parameters(kwargs)
changed_parameters = ase.calculators.calculator.Calculator.set(self, **kwargs)
# Always reset the calculation if parameters change
if changed_parameters:
self.reset()
# If the method is changed, invalidate the cached calculator as well
if (
"method" in changed_parameters
or "electric_field" in changed_parameters
or "spin_polarization" in changed_parameters
or "solvation" in changed_parameters
or "xtb_config" in changed_parameters
):
self._xtb = None
self._res = None
# Minor changes can be updated in the API calculator directly
if self._xtb is not None:
if "accuracy" in changed_parameters:
self._xtb.set("accuracy", self.parameters.accuracy)
if "electronic_temperature" in changed_parameters:
self._xtb.set(
"temperature",
self.parameters.electronic_temperature * kB / Hartree,
)
if "max_iterations" in changed_parameters:
self._xtb.set("max-iter", self.parameters.max_iterations)
if "initial_guess" in changed_parameters:
self._xtb.set(
"guess", {"sad": 0, "eeq": 1, "eeqbc": 2}[self.parameters.guess]
)
if "mixer_damping" in changed_parameters:
self._xtb.set("mixer-damping", self.parameters.mixer_damping)
if (
"charge" in changed_parameters or "multiplicity" in changed_parameters
) and self.atoms is not None:
self._xtb.update(
charge=_get_charge(self.atoms, self.parameters),
uhf=_get_uhf(self.atoms, self.parameters),
)
return changed_parameters
[docs]
def reset(self) -> None:
"""
Clear all information from old calculation. This will only remove the cached
API objects in case the `cache_api` is set to False.
"""
ase.calculators.calculator.Calculator.reset(self)
if not self.parameters.cache_api:
self._xtb = None
self._res = None
def _check_api_calculator(self, system_changes: List[str]) -> None:
"""Check state of API calculator and reset if necessary"""
# Changes in positions, cell, charges and magnetic moments can use a normal update
_reset = system_changes.copy()
for _change in system_changes:
if _change in ("positions", "cell", "initial_charges", "initial_magmoms"):
_reset.remove(_change)
# Invalidate cached calculator and results object
if _reset:
self._xtb = None
self._res = None
else:
if system_changes and self._xtb is not None:
try:
_cell = self.atoms.cell
self._xtb.update(
positions=self.atoms.positions / Bohr,
lattice=_cell / Bohr,
charge=_get_charge(self.atoms, self.parameters),
uhf=_get_uhf(self.atoms, self.parameters),
)
# An exception in this part means the geometry is bad,
# still we will give a complete reset a try as well
except RuntimeError:
self._xtb = None
self._res = None
[docs]
def calculate(
self,
atoms: Optional[Atoms] = None,
properties: Optional[List[str]] = None,
system_changes: List[str] = ase.calculators.calculator.all_changes,
) -> None:
"""
Perform actual calculation by calling the TBLite API
Example
-------
>>> from ase.build import molecule
>>> from tblite.ase import TBLite
>>> calc = TBLite(method="GFN2-xTB")
>>> calc.calculate(molecule("H2O"))
>>> calc.get_potential_energy()
-137.96777625229421
>>> calc.calculate(molecule("CH4"))
>>> calc.get_potential_energy()
-113.60956621093894
Raises
------
ase.calculators.calculator.InputError
on invalid input passed to the interface module
ase.calculators.calculator.CalculationFailed
in case of an `RuntimeError` in the library
"""
if not properties:
properties = ["energy"]
ase.calculators.calculator.Calculator.calculate(
self, atoms, properties, system_changes
)
self._check_api_calculator(system_changes)
if self._xtb is None:
self._xtb = _create_api_calculator(self.atoms, self.parameters)
try:
self._res = self._xtb.singlepoint(self._res)
except RuntimeError as e:
raise ase.calculators.calculator.CalculationFailed(str(e)) from e
# These properties are garanteed to exist for all implemented calculators
self.results["energy"] = self._res["energy"] * Hartree
self.results["energies"] = self._res["energies"] * Hartree
self.results["free_energy"] = self.results["energy"]
self.results["forces"] = -self._res["gradient"] * Hartree / Bohr
self.results["charges"] = self._res["charges"]
self.results["dipole"] = self._res["dipole"] * Bohr
self.results["bond-orders"] = self._res["bond-orders"]
# stress tensor is only returned for periodic systems
if self.atoms.pbc.any():
_stress = self._res["virial"] * Hartree / self.atoms.get_volume()
self.results["stress"] = _stress.flat[[0, 4, 8, 5, 2, 1]]
def _create_api_calculator(
atoms: ase.Atoms,
parameters: ase.calculators.calculator.Parameters,
) -> Calculator:
"""Create a new API calculator object"""
try:
_cell = atoms.cell
_periodic = atoms.pbc
_charge = _get_charge(atoms, parameters)
_uhf = _get_uhf(atoms, parameters)
calc = Calculator(
parameters.method,
atoms.numbers,
atoms.positions / Bohr,
_charge,
_uhf,
_cell / Bohr,
_periodic,
xtb_config=parameters.xtb_config,
)
calc.set("accuracy", parameters.accuracy)
calc.set(
"temperature",
parameters.electronic_temperature * kB / Hartree,
)
calc.set("max-iter", parameters.max_iterations)
calc.set("guess", {"sad": 0, "eeq": 1, "eeqbc": 2}[parameters.guess])
calc.set("mixer-damping", parameters.mixer_damping)
calc.set("verbosity", parameters.verbosity)
if parameters.electric_field is not None:
calc.add(
"electric-field",
parameters.electric_field * Bohr / Hartree,
)
if parameters.spin_polarization is not None:
calc.add("spin-polarization", parameters.spin_polarization)
if parameters.solvation is not None:
solvation_model, *solvation_args = parameters.solvation
if isinstance(solvation_args, (tuple, list)):
calc.add(f"{solvation_model}-solvation", *solvation_args)
else:
calc.add(f"{solvation_model}-solvation", solvation_args)
except RuntimeError as e:
raise ase.calculators.calculator.InputError(str(e)) from e
return calc
def _get_charge(
atoms: ase.Atoms, parameters: ase.calculators.calculator.Parameters
) -> int:
"""
Get the total charge of the system.
If no charge is provided, the total charge of the system is calculated
by summing the initial charges of all atoms.
"""
return (
atoms.get_initial_charges().sum()
if parameters.charge is None
else parameters.charge
)
def _get_uhf(
atoms: ase.Atoms, parameters: ase.calculators.calculator.Parameters
) -> int:
"""
Get the number of unpaired electrons.
If no multiplicity is provided, the number of unpaired electrons
is calculated by summing the initial magnetic moments of all atoms.
"""
return (
int(atoms.get_initial_magnetic_moments().sum().round())
if parameters.multiplicity is None
else parameters.multiplicity - 1
)
def _update_parameters(kwargs: Dict[str, Any]) -> None:
"""
Update the parameters of the calculator with the provided keyword arguments.
This function is used to update the parameters of the calculator
when new values are provided.
"""
for key in (
"alpb",
"gbsa",
"cpcm",
"gbe",
"gb",
):
if f"{key}_solvation" in kwargs:
value = kwargs.pop(f"{key}_solvation")
if isinstance(value, (tuple, list)):
value = (key, *value)
else:
value = (key, value)
if "solvation" in kwargs:
raise ase.calculators.calculator.InputError(
"Multiple solvation models provided, can only use one"
)
kwargs["solvation"] = value
if "tblite" not in ase.calculators.calculator.external_calculators:
ase.calculators.calculator.register_calculator_class("tblite", TBLite)