"""Compute GDMA and multipole data using Psi4."""
import os
import subprocess
from typing import TYPE_CHECKING
import jinja2
import numpy as np
from openff.recharge.esp.exceptions import Psi4Error
from openff.units import Quantity, unit
from openff.units.elements import SYMBOLS
from openff.utilities import get_data_file_path, temporary_cd
from pympfit.gdma import GDMAGenerator, GDMASettings
if TYPE_CHECKING:
from openff.toolkit import Molecule
[docs]
class Psi4GDMAGenerator(GDMAGenerator):
"""Compute the multipole moments of a molecule using Psi4."""
@classmethod
def _generate_input(
cls,
molecule: "Molecule",
conformer: Quantity,
settings: GDMASettings,
minimize: bool,
compute_mp: bool,
memory: Quantity = 500 * unit.mebibytes,
) -> str:
"""Generate the input files for Psi4.
Parameters
----------
molecule
The molecule to generate the GDMA for.
conformer
The conformer of the molecule to generate the GDMA for.
settings
The settings to use when generating the GDMA.
minimize
Whether to energy minimize the conformer prior to computing the GDMA using
the same level of theory that the GDMA will be computed at.
compute_esp
Whether to compute the multipoles.
compute_field
Whether to compute the field at each grid point.
memory
The memory to make available to Psi4 for computation
Returns
-------
The contents of the input file.
"""
# Compute the total formal charge on the molecule.
# Trust that it's in units of elementary charge.
formal_charge = sum(atom.formal_charge for atom in molecule.atoms).m
# Compute the spin multiplicity
total_atomic_number = sum(atom.atomic_number for atom in molecule.atoms)
spin_multiplicity = 1 if (formal_charge + total_atomic_number) % 2 == 0 else 2
# Store the atoms and coordinates in a jinja friendly dict.
conformer = conformer.to(unit.angstrom).m
atoms = [
{
"element": SYMBOLS[atom.atomic_number],
"x": conformer[index, 0],
"y": conformer[index, 1],
"z": conformer[index, 2],
}
for index, atom in enumerate(molecule.atoms)
]
# Format the jinja template
template_path = get_data_file_path(os.path.join("psi4", "input.dat"), "pympfit")
with open(template_path) as file:
template = jinja2.Template(file.read())
properties = []
if compute_mp:
properties.append("MULTIPOLE_MOMENT")
template_inputs = {
"charge": formal_charge,
"spin": spin_multiplicity,
"atoms": atoms,
"basis": settings.basis,
"method": settings.method,
"limit": settings.limit,
"multipole_units": settings.multipole_units,
"radius": settings.radius,
"switch": settings.switch,
"minimize": minimize,
"compute_mp": compute_mp,
"properties": str(properties),
"memory": f"{memory:~P}",
}
# Remove the white space after the for loop
return template.render(template_inputs).replace(" \n}", "}")
@classmethod
def _generate(
cls,
molecule: "Molecule",
conformer: Quantity,
settings: GDMASettings,
_directory: str,
minimize: bool,
compute_mp: bool,
n_threads: int,
memory: Quantity = 500 * unit.mebibytes,
) -> tuple[Quantity, Quantity | None, Quantity | None]:
with temporary_cd():
input_contents = cls._generate_input(
molecule,
conformer,
settings,
minimize,
compute_mp,
memory=memory,
)
with open("input.dat", "w") as file:
file.write(input_contents)
# Attempt to run the calculation
psi4_process = subprocess.Popen(
["psi4", "--nthread", str(n_threads), "input.dat", "output.dat"],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
std_output, std_error = psi4_process.communicate()
exit_code = psi4_process.returncode
if exit_code != 0:
raise Psi4Error(std_output.decode(), std_error.decode())
mp = None
if compute_mp:
mp = np.load("dma_distributed.npy")
with open("final-geometry.xyz") as file:
output_lines = file.read().splitlines(keepends=False)
final_coordinates = (
np.array(
[
[
float(coordinate)
for coordinate in coordinate_line.split()[1:]
]
for coordinate_line in output_lines[2:]
if len(coordinate_line) > 0
]
)
* unit.angstrom
)
return final_coordinates, mp