xnn.common.deploy.mdi_engine.MDIEngine#

class xnn.common.deploy.mdi_engine.MDIEngine(model, cutoff, device='cpu', total_charge=0.0)[source]#

Bases: object

MDI engine exposing a trained xnn model to an external driver.

Holds the driver-supplied system state (atom count, elements, cell, coordinates) and lazily re-evaluates the model whenever a result is requested after the geometry changed. Each evaluation converts the MDI atomic-unit inputs to angstrom, builds an AtomicGraph via structure_to_graph() (which handles molecular and periodic systems alike), runs the model, and converts the eV / angstrom outputs back to Hartree / Bohr.

Parameters:

  • model (torch.nn.Module) – A trained model. It is moved to device, put in eval mode and its parameter gradients are disabled. Its forward must accept an AtomicGraph and return a dict with an "energy" key and, for the <FORCES / <STRESS commands, "forces" / "stress" keys; wrap a bare model in ForceStressOutput to provide them (as from_checkpoint() does).

  • cutoff (float) – Neighbor-list cutoff radius in angstrom used when building the graph. Use the model’s own cutoff (a dispersion wrapper widens it beyond the core model’s radius), as from_checkpoint() does.

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

  • total_charge (float, optional) – Net charge of the system in units of e, passed to the model as total_charge (D4’s EEQ charges and charge-aware models such as PhysNet use it). Defaults to 0; a driver can change it at run time with >TOTCHARGE.

Variables:

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

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

  • device (torch.device) – The torch device.

  • total_charge (float) – The current net charge (e).

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

  • stress (energy, forces,) – Results of the latest evaluation, in MDI atomic units (Hartree, Hartree/Bohr, Hartree/Bohr^3). None before the first evaluation (stress also for non-periodic systems).

classmethod from_checkpoint(path, device='cpu', dtype=None, dispersion=None, total_charge=0.0, eeq_reuse=False, cache_dir=None, use_fast='auto', head=None)[source]#

Build an engine from a checkpoint or a pre-trained model.

path is anything from_pretrained() loads: a trainer checkpoint (best.pt, the {"model": state_dict, "cfg": Config} dictionary written by save()), a portable model directory, a registered model name such as mace-off23-small, a URL or a Zenodo DOI. The model is rebuilt with build_model() from the stored config, wrapped in ForceStressOutput with the stress head enabled, and the weights are loaded.

The model is built in float64, so the constant tables of the physics terms (D3 / D4 reference data, LES kernels) hold their exact values, and cast once afterwards, to dtype or to the checkpoint’s own floating-point dtype. Building in float32 and upcasting would keep float32-rounded tables, which costs about 5e-8 hartree in a D4 energy even when serving in float64.

Parameters:

  • path (str) – Path to the checkpoint file.

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

  • dtype (torch.dtype, optional) – Floating-point dtype to serve in (torch.float64 for NVE energy conservation with a float32-trained model, torch.float32 for speed). Defaults to the dtype of the checkpoint’s weights.

  • dispersion (dict, str, bool or None, optional) – Add a D3 / D4 dispersion correction to a checkpoint that was trained without one: "d4" / "d3" for the defaults or a mapping as in the config’s extra["dispersion"] (see add_dispersion()). A checkpoint whose config records subtracted_dispersion (what the labels had removed) gets that term added by default, None; an explicit mapping then overrides the recorded keys it names, and False serves the checkpoint as is. Refused when the checkpoint already carries dispersion, which would count it twice.

  • total_charge (float, optional) – Net charge of the system in units of e, by default 0. A driver can change it at run time with >TOTCHARGE.

  • eeq_reuse (bool, optional) – Carry the large-regime D4 EEQ solve over from one step to the next (enable_eeq_reuse()), for molecular dynamics and optimizations where each structure follows the previous one closely. Results agree with the fresh solve to the solver tolerance (1e-9 relative residual in float64). By default False.

  • cache_dir (str, optional) – Model hub cache for a registered name, URL or DOI (see default_model_cache_dir()).

  • use_fast (bool or str, optional) – "auto" (default), True or False: the fused GPU kernels or the reference implementation (see xnn.common.models.fast).

  • head (str, optional) – The head of a multi-head (replay fine-tuned) checkpoint to serve.

Returns:

An engine wrapping the restored model, with the neighbor-list cutoff taken from the built model (a dispersion wrapper widens it beyond the config’s core-model radius).

Return type:

MDIEngine

Raises:

ValueError – If dispersion is given for a checkpoint whose model already includes a dispersion correction.

calculate()[source]#

Evaluate the model on the current system state.

Converts positions and cell from Bohr to angstrom, builds the graph with structure_to_graph() (positions and cell in float64, the model computing in its own dtype; the current total_charge rides along), runs the model and stores energy (Hartree), forces (Hartree/Bohr) and, for periodic systems, stress (Hartree/Bohr^3).

Raises:

RuntimeError – If no coordinates or elements have been received yet.

Return type:

None

run(mdi_options, mpi_comm=None)[source]#

Run the MDI engine loop until the driver sends EXIT.

Initializes the MDI library, registers the supported commands on the @DEFAULT node, accepts the driver connection and then services commands: system updates (>NATOMS, >ELEMENTS, >CELL, >COORDS, >TOTCHARGE) mark the results stale, result requests (<ENERGY, <FORCES, <STRESS) trigger a model evaluation when needed and send the values in MDI atomic units. <STRESS sends zeros for non-periodic systems. SCF forces an immediate evaluation.

Parameters:

  • mdi_options (str) – The MDI option string, e.g. "-role ENGINE -name xnn -method TCP -port 8021 -hostname localhost".

  • mpi_comm (mpi4py.MPI.Comm, optional) – MPI communicator to hand to MDI_Init (required for -method MPI). Defaults to None (TCP method).

Raises:

  • ImportError – If the pymdi package is not installed.

  • RuntimeError – If the driver sends a command this engine does not support, or requests forces/stress from a model that does not produce them.

Return type:

None