ASE Calculator Tutorial#

This tutorial demonstrates how to use a TACE model as a calculator within ASE (Atomic Simulation Environment).

ASE Calculator documentation: ASE Calculator

For detailed usage and scripts (e.g., opt, and other scripts), see xvzemin/tace

from ase.io import read
from tace.interface.ase import TACEAseCalc, add_dispersion

device = 'cuda'            # Compute device, e.g., 'cpu' or 'cuda'
dtype = 'float32'          # model dtype 'float32' or 'float64'
MODEL_PATH = '.pt'         # Path to the model checkpoint, file ends with .pt, .pth or .ckpt
level = 0  # first fidelity
atoms = read('*.xyz', 0)   #  Any ase readable files

dispersion = False

calc = TACEAseCalc(
    MODEL_PATH,
    device=device,
    dtype=dtype,
    level = level,
)
if dispersion: # pip install torch-dftd
    calc = add_dispersion(
        base_calc=calc,
        damping= "bj",  # choices: ["zero", "bj", "zerom", "bjm"]
        dispersion_xc="pbe",
        dispersion_cutoff= 40.0 * units.Bohr,
    )
atoms.calc = calc
class tace.interface.ase.calculator.TACEAseCalc(model: str, *, dtype: str | None = None, device: str | None = None, level: int | None = None, spin_on: bool | None = None, target_property: list[str] | None = None, neighborlist_backend: str = 'matscipy', **kwargs)[source]#

Bases: Calculator

Initialize a TACEAseCalc. We support the most fundamental potential energy surface property and multi-fidelity, multi-head, etc. For some advanced features, you need to store the attributes that need to be embedded in atoms.info, atoms.arrays or add a funciton by yourself. If you only need to predict, you can directly use the tace-eval command. It will output the predicted files, and if you add the –test option, it will also output the errors.

Parameters:

  • model_path (str) – Path to the trained model, file ends with pt, .pth or .ckpt.

  • device (str | torch.device, optional) – The device to run computations on, e.g., cpu or cuda. If None, the device is automatically inferred.

  • dtype (str, optional, default=None) – Model dtype for computations, e.g., float32 or float64.

  • level (int) – Specify which fidelity level to use.

  • spin_on (bool) – If your model uses spin_on uie embedding, you can control whether your calculation enables spin polarization.

  • target_property (list(str)) – Extra caculate hessians, atomic_virials, Conservative polarizability, etc, If you want to use this parameter, you must provide all the required physical quantities.

  • neighborlist_backend (str) – Support backend in one of [ase, matscipy, vesin], recommend matscipy

  • **kwargs – Additional keyword arguments passed to the ASE Calculator base class.

On this page