xnn.common.deploy.ase_calculator.XNNCalculator#

class xnn.common.deploy.ase_calculator.XNNCalculator(model, cutoff, device='cpu', **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.

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.

Raises:

ImportError – If ASE is not installed.

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

Properties calculator can handle (energy, forces, …)

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].

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