"""Support for using QCEngine as an MDI engine.
For details regarding MDI, see https://molssi.github.io/MDI_Library/html/index.html.
"""
from typing import Any, Dict, List, Optional
import numpy as np
import qcelemental as qcel
from qcelemental.util import which_import
from .compute import compute
try:
from mdi import (
MDI_CHAR,
MDI_COMMAND_LENGTH,
MDI_DOUBLE,
MDI_INT,
MDI_MAJOR_VERSION,
MDI_Accept_Communicator,
MDI_Init,
MDI_MPI_get_world_comm,
MDI_Recv,
MDI_Recv_Command,
MDI_Register_Command,
MDI_Register_Node,
MDI_Send,
)
use_mdi = True
except ImportError:
use_mdi = False
[docs]class MDIServer:
def __init__(
self,
mdi_options: str,
program: str,
molecule,
model,
keywords,
raise_error: bool = False,
local_options: Optional[Dict[str, Any]] = None,
):
"""Initialize an MDIServer object for communication with MDI
Parameters
----------
mdi_options: str
Options used during MDI initialization.
program : str
The program to execute the input with.
molecule
The initial state of the molecule.
model
The simulation model to use.
keywords
Program-specific keywords.
raise_error : bool, optional
Determines if compute should raise an error or not.
local_options : Optional[Dict[str, Any]], optional
A dictionary of local configuration options
"""
if not use_mdi:
raise Exception("Trying to run as an MDI engine, but the MDI Library was not found")
if MDI_MAJOR_VERSION < 1:
raise Exception("QCEngine requires version 1.0.0 or higher of the MDI Library")
# Confirm that the MDI library has been located
which_import("mdi", raise_error=True, raise_msg="Please install via 'conda install pymdi -c conda-forge'")
# Initialize MDI
MDI_Init(mdi_options)
# Input variables
self.molecule = molecule
self.model = model
self.keywords = keywords
self.program = program
self.raise_error = raise_error
self.local_options = local_options
# The MDI interface does not currently support multiple fragments
if len(self.molecule.fragments) != 1:
raise Exception("The MDI interface does not support multiple fragments")
# Molecule charge and multiplicity
self.total_charge = self.molecule.molecular_charge
self.multiplicity = self.molecule.molecular_multiplicity
# Flag to track whether the latest molecule specification has been validated
self.molecule_validated = True
# Output of most recent compute call
self.compute_return = None
# Set whether the current energy is valid
self.energy_is_current = False
# MPI variables
self.mpi_world = None
self.world_rank = 0
# Flag to stop listening for MDI commands
self.stop_listening = False
# Dictionary of all supported MDI commands
self.commands = {
"<@": self.send_node,
"<NATOMS": self.send_natoms,
">NATOMS": self.recv_natoms,
"<COORDS": self.send_coords,
"<ENERGY": self.send_energy,
"<FORCES": self.send_forces,
">COORDS": self.recv_coords,
"SCF": self.run_energy,
"<ELEMENTS": self.send_elements,
">ELEMENTS": self.recv_elements,
"<MASSES": self.send_masses,
">MASSES": self.recv_masses,
"<TOTCHARGE": self.send_total_charge,
">TOTCHARGE": self.recv_total_charge,
"<ELEC_MULT": self.send_multiplicity,
">ELEC_MULT": self.recv_multiplicity,
"EXIT": self.stop,
}
# Register the @DEFAULT node
MDI_Register_Node("@DEFAULT")
# Register all supported commands
for c in self.commands.keys():
MDI_Register_Command("@DEFAULT", c)
# Set the current node
self.current_node = "@DEFAULT"
# Accept a communicator to the driver code
self.comm = MDI_Accept_Communicator()
[docs] def update_molecule(self, key: str, value):
"""Update the molecule
Parameters
----------
key : str
Key of the molecular element to update
value
Update value
"""
if key == "molecular_charge" or key == "molecular_multiplicity":
# In order to validate correctly, the charges and multiplicities must be set simultaneously
try:
self.molecule = qcel.models.Molecule(
**{
**self.molecule.dict(),
**{
"molecular_charge": self.total_charge,
"fragment_charges": [self.total_charge],
"molecular_multiplicity": self.multiplicity,
"fragment_multiplicities": [self.multiplicity],
},
}
)
self.molecule_validated = True
except qcel.exceptions.ValidationError:
# The molecule didn't validate correctly, but a future >TOTCHARGE or >ELEC_MULT command might fix it
self.molecule_validated = False
else:
try:
self.molecule = qcel.models.Molecule(**{**self.molecule.dict(), **{key: value}})
self.molecule_validated = True
except qcel.exceptions.ValidationError:
if self.molecule_validated:
# This update caused the validation error
raise Exception("MDI command caused a validation error")
# Respond to the <@ command
[docs] def send_node(self) -> str:
"""Send the name of the current node through MDI
Returns
-------
node : str
Name of the current node
"""
node = self.current_node
MDI_Send(node, MDI_COMMAND_LENGTH, MDI_CHAR, self.comm)
return node
# Respond to the <NATOMS command
[docs] def send_natoms(self) -> int:
"""Send the number of atoms through MDI
Returns
-------
natom : int
Number of atoms
"""
natom = len(self.molecule.geometry)
MDI_Send(natom, 1, MDI_INT, self.comm)
return natom
# Respond to the >NATOMS command
[docs] def recv_natoms(self, natoms: Optional[int] = None) -> None:
"""Receive the number of atoms in the system through MDI and create a new molecule with them
Parameters
----------
natoms : int, optional
New number of atoms. If None, receive through MDI.
"""
natom = len(self.molecule.geometry)
if natoms is None:
natoms = MDI_Recv(1, MDI_INT, self.comm)
mol_string = ""
for iatom in range(natoms):
mol_string += "He " + str(1.0 * iatom) + " 0.0 0.0\n"
self.molecule = qcel.models.Molecule.from_data(mol_string)
self.energy_is_current = False
# Respond to the <COORDS command
[docs] def send_coords(self) -> np.ndarray:
"""Send the nuclear coordinates through MDI
Returns
-------
coords : np.ndarray
Nuclear coordinates
"""
natom = len(self.molecule.geometry)
coords = np.reshape(self.molecule.geometry, (3 * natom))
MDI_Send(coords, 3 * natom, MDI_DOUBLE, self.comm)
return coords
# Respond to the >COORDS command
[docs] def recv_coords(self, coords: Optional[np.ndarray] = None) -> None:
"""Receive a set of nuclear coordinates through MDI and assign them to the atoms in the current molecule
Parameters
----------
coords : np.ndarray, optional
New nuclear coordinates. If None, receive through MDI.
"""
natom = len(self.molecule.geometry)
if coords is None:
coords = np.zeros(3 * natom)
MDI_Recv(3 * natom, MDI_DOUBLE, self.comm, buf=coords)
new_geometry = np.reshape(coords, (-1, 3))
self.molecule = qcel.models.Molecule(**{**self.molecule.dict(), **{"geometry": new_geometry}})
self.energy_is_current = False
# Respond to the <ENERGY command
[docs] def send_energy(self) -> float:
"""Send the total energy through MDI
Returns
-------
energy : float
Energy of the system
"""
# Ensure that the molecule currently passes validation
if not self.molecule_validated:
raise Exception("MDI attempting to compute energy on an unvalidated molecule")
self.run_energy()
# Confirm that the calculation completed successfully
if not hasattr(self.compute_return, "properties"):
raise Exception("MDI Calculation failed: \n\n" + str(self.compute_return.error.error_message))
properties = self.compute_return.properties.dict()
energy = properties["return_energy"]
MDI_Send(energy, 1, MDI_DOUBLE, self.comm)
return energy
# Respond to the <FORCES command
[docs] def send_forces(self) -> np.ndarray:
"""Send the nuclear forces through MDI
Returns
-------
forces : np.ndarray
Forces on the nuclei
"""
# Ensure that the molecule currently passes validation
if not self.molecule_validated:
raise Exception("MDI attempting to compute gradients on an unvalidated molecule")
self.run_energy()
properties = self.compute_return.properties.dict()
forces = np.reshape(-1.0 * properties["return_gradient"], (-1,))
if len(forces) != 3 * len(self.molecule.geometry):
raise Exception(
"MDI: The length of the forces is not what was expected. Expected: "
+ str(3 * len(self.molecule.geometry))
+ " Actual: "
+ str(len(forces))
)
MDI_Send(forces, len(forces), MDI_DOUBLE, self.comm)
return forces
# Respond to the SCF command
[docs] def run_energy(self) -> None:
if not self.energy_is_current:
"""Ensure that the orientation of the molecule remains fixed"""
self.update_molecule("fix_com", True)
self.update_molecule("fix_orientation", True)
"""Run an energy calculation"""
input = qcel.models.AtomicInput(
molecule=self.molecule, driver="gradient", model=self.model, keywords=self.keywords
)
self.compute_return = compute(
input_data=input, program=self.program, raise_error=self.raise_error, local_options=self.local_options
)
# If there is an error message, print it out
if hasattr(self.compute_return, "error"):
if self.compute_return.error is not None:
print("---------------- QCEngine Compute Error ----------------\n\n")
print(str(self.compute_return.error.error_message))
print("\n\n-------------- End QCEngine Compute Error --------------", flush=True)
self.energy_is_current = True
# Respond to the <ELEMENTS command
[docs] def send_elements(self):
"""Send the atomic number of each nucleus through MDI
Returns
-------
elements : :obj:`list` of :obj:`int`
Element of each atom
"""
natom = len(self.molecule.geometry)
elements = [qcel.periodictable.to_atomic_number(self.molecule.symbols[iatom]) for iatom in range(natom)]
MDI_Send(elements, natom, MDI_INT, self.comm)
return elements
# Respond to the >ELEMENTS command
[docs] def recv_elements(self, elements: Optional[List[int]] = None):
"""Receive a set of atomic numbers through MDI and assign them to the atoms in the current molecule
Parameters
----------
elements : :obj:`list` of :obj:`int`, optional
New element numbers. If None, receive through MDI.
"""
natom = len(self.molecule.geometry)
if elements is None:
elements = MDI_Recv(natom, MDI_INT, self.comm)
for iatom in range(natom):
self.molecule.symbols[iatom] = qcel.periodictable.to_symbol(elements[iatom])
return elements
# Respond to the <MASSES command
[docs] def send_masses(self) -> np.ndarray:
"""Send the nuclear masses through MDI
Returns
-------
masses : np.ndarray
Atomic masses
"""
natom = len(self.molecule.geometry)
masses = self.molecule.masses
MDI_Send(masses, natom, MDI_DOUBLE, self.comm)
return masses
# Respond to the >MASSES command
[docs] def recv_masses(self, masses: Optional[List[float]] = None) -> None:
"""Receive a set of nuclear masses through MDI and assign them to the atoms in the current molecule
Parameters
----------
masses : :obj:`list` of :obj:`float`, optional
New nuclear masses. If None, receive through MDI.
"""
natom = len(self.molecule.geometry)
if masses is None:
masses = MDI_Recv(natom, MDI_DOUBLE, self.comm)
self.update_molecule("masses", masses)
self.energy_is_current = False
# Respond to the <TOTCHARGE command
[docs] def send_total_charge(self) -> float:
"""Send the total system charge through MDI
Returns
-------
charge : float
Total charge of the system
"""
charge = self.molecule.molecular_charge
MDI_Send(charge, 1, MDI_DOUBLE, self.comm)
return charge
# Respond to the >TOTCHARGE command
[docs] def recv_total_charge(self, charge: Optional[float] = None) -> None:
"""Receive the total system charge through MDI
Parameters
----------
charge : float, optional
New charge of the system. If None, receive through MDI.
"""
if charge is None:
charge = MDI_Recv(1, MDI_DOUBLE, self.comm)
self.total_charge = charge
self.energy_is_current = False
# Allow a validation error here, because a future >ELEC_MULT command might resolve it
try:
self.update_molecule("molecular_charge", self.total_charge)
except qcel.exceptions.ValidationError:
pass
# Respond to the <ELEC_MULT command
[docs] def send_multiplicity(self) -> int:
"""Send the electronic multiplicity through MDI
Returns
-------
multiplicity : int
Multiplicity of the system
"""
multiplicity = self.molecule.molecular_multiplicity
MDI_Send(multiplicity, 1, MDI_INT, self.comm)
return multiplicity
# Respond to the >ELEC_MULT command
[docs] def recv_multiplicity(self, multiplicity: Optional[int] = None) -> None:
"""Receive the electronic multiplicity through MDI
Parameters
----------
multiplicity : int, optional
New multiplicity of the system. If None, receive through MDI.
"""
if multiplicity is None:
multiplicity = MDI_Recv(1, MDI_INT, self.comm)
self.multiplicity = multiplicity
self.energy_is_current = False
# Allow a validation error here, because a future >TOTCHARGE command might resolve it
try:
self.update_molecule("molecular_multiplicity", self.multiplicity)
except qcel.exceptions.ValidationError:
pass
# Respond to the EXIT command
[docs] def stop(self) -> None:
"""Stop listening for MDI commands"""
self.stop_listening = True
# Enter server mode, listening for commands from the driver
[docs] def start(self) -> None:
"""Receive commands through MDI and respond to them as defined by the MDI Standard"""
while not self.stop_listening:
if self.world_rank == 0:
command = MDI_Recv_Command(self.comm)
else:
command = None
if self.world_rank == 0:
print("MDI command received: " + str(command))
# Search for this command in self.commands
found_command = False
for supported_command in self.commands:
if not found_command and command == supported_command:
# Run the function corresponding to this command
self.commands[supported_command]()
found_command = True
if not found_command:
raise Exception("Unrecognized command: " + str(command))