xnn.common.deploy.ase_calculator.XNNCalculator#

class xnn.common.deploy.ase_calculator.XNNCalculator(model, cutoff, device='cpu', eeq_reuse=False, **kwargs)[source]#

Bases: Calculator

ASE Calculator wrapping a trained xnn model.

This is an ase.calculators.calculator.Calculator subclass that lets a trained model drive standard ASE workflows (energy, force and stress evaluation, dynamics, optimization, etc.). It converts an Atoms object into the model’s graph representation, runs the model and stores the results in the ASE-expected units and layout.

Requires ase to be installed (optional extra: pip install "xnn[ase]"); constructing an instance raises ImportError if ASE is not available.

Parameters:

  • model (torch.nn.Module) – A trained model. It is moved to device and put in eval mode. Its forward is expected to accept an AtomicGraph and return a dict with an "energy" key and optionally "forces" and "stress".

  • cutoff (float) – Neighbor-list cutoff radius (in ASE length units) used when building the graph from an Atoms object.

  • device (str, optional) – Torch device the model runs on. Defaults to "cpu".

  • **kwargs – Additional keyword arguments forwarded to the base Calculator.

  • eeq_reuse (bool)

Variables:

  • implemented_properties (list of str) – Properties this calculator can produce: "energy", "forces" and "stress".

  • model (torch.nn.Module) – The wrapped model (on device, in eval mode).

  • cutoff (float) – The neighbor-list cutoff radius.

  • device (str) – The torch device string.

  • dtype (torch.dtype) – Floating-point dtype of the model; graph tensors are built in it.

  • eeq_reuse (bool, optional) – Carry the large-regime EEQ solve of a D4 term from one step to the next (enable_eeq_reuse()), for molecular dynamics and geometry optimization; results agree with the fresh solve to the solver tolerance. Default False.

Raises:

ImportError – If ASE is not installed.

implemented_properties: list[str] = ['energy', 'forces', 'stress', 'charges', 'dipole']#

Properties calculator can handle (energy, forces, …)

classmethod from_pretrained(source, device='cpu', eeq_reuse=False, calculator_kwargs=None, **kwargs)[source]#

Build a calculator for a pre-trained model.

Loads source with load_pretrained() (a registered name such as "mace-off23-small", a local checkpoint or model directory, a URL or a Zenodo DOI) with the stress head on, and takes the neighbor-list cutoff from the loaded model:

atoms.calc = XNNCalculator.from_pretrained("mace-mp-0-medium")
Parameters:

  • source (str or pathlib.Path) – The model to load.

  • device (str, optional) – Torch device. Defaults to "cpu".

  • eeq_reuse (bool, optional) – As for the constructor.

  • calculator_kwargs (dict, optional) – Keyword arguments for the base ASE Calculator.

  • **kwargs – Options of load_pretrained() (cache_dir, head, filename, dtype, dispersion, local_files_only, use_fast, …).

Returns:

The calculator.

Return type:

XNNCalculator

check_state(atoms, tol=1e-15)[source]#

ASE’s change detection, plus the net charge and spin multiplicity.

ASE compares positions, numbers, cell and periodicity only, so the same geometry with another atoms.info["charge"] would otherwise be served the cached result.

calculate(atoms=None, properties=('energy',), system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])[source]#

Compute requested properties for an Atoms object.

Builds an AtomicGraph from the atoms’ positions, atomic numbers, cell and periodic-boundary flags (the cell is passed only when any PBC direction is active), runs the model, and populates self.results. The total energy is the sum of the model’s per-atom (node) energies. Forces, when produced, are stored as an (n_atoms, 3) NumPy array. Stress, when produced and the system is periodic, is converted from the model’s 3 x 3 tensor to ASE’s Voigt 6-vector ordering [xx, yy, zz, yz, xz, xy]. A model that predicts partial charges (PhysNet, BAMBOO, AIMNet2) also fills "charges" and "dipole" (in the model’s units); so does an LES model from its les_charges and, for molecules, les_dipole (e and e Angstrom). The net charge and spin multiplicity of the structure are read from atoms.info (keys "charge" / "total_charge" and "spin_multiplicity" / "multiplicity"), as atoms_to_structure() does.

Parameters:

  • atoms (ase.Atoms or None, optional) – The atomic structure to evaluate. Passed through to the base Calculator.

  • properties (sequence of str, optional) – Names of the properties to compute. Defaults to ("energy",).

  • system_changes (list of str, optional) – Which aspects of the system have changed since the last call. Defaults to ASE’s all_changes.

Returns:

Results are written into self.results.

Return type:

None