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)[source]#

Build an engine from a trainer checkpoint (best.pt).

The checkpoint is the dictionary written by save(): {"model": state_dict, "cfg": Config}. The model is rebuilt with build_model() from the stored config, wrapped in ForceStressOutput (with the stress head enabled, matching the state-dict layout the trainer saves), 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 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()). Refused when the checkpoint already carries dispersion, which would count it twice. Defaults to None (serve the checkpoint as is).

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

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() (tensors are created in the model’s dtype so float32 and float64 models both work; 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