def infer_mol_from_coordinates(
coordinates,
species,
smiles_ref=None,
coordinates_unit="angstrom",
):
# local import
from openeye import oechem
from simtk import unit
from simtk.unit import Quantity
if isinstance(coordinates_unit, str):
coordinates_unit = getattr(unit, coordinates_unit)
# make sure we have the coordinates
# in the unit system
coordinates = Quantity(coordinates, coordinates_unit).value_in_unit(
unit.angstrom # to make openeye happy
)
# initialize molecule
mol = oechem.OEGraphMol()
if all(isinstance(symbol, str) for symbol in species):
[
mol.NewAtom(getattr(oechem, "OEElemNo_" + symbol))
for symbol in species
]
elif all(isinstance(symbol, int) for symbol in species):
[
mol.NewAtom(
getattr(oechem,
"OEElemNo_" + oechem.OEGetAtomicSymbol(symbol)))
for symbol in species
]
else:
raise RuntimeError(
"The species can only be all strings or all integers.")
mol.SetCoords(coordinates.reshape([-1]))
mol.SetDimension(3)
oechem.OEDetermineConnectivity(mol)
oechem.OEFindRingAtomsAndBonds(mol)
oechem.OEPerceiveBondOrders(mol)
if smiles_ref is not None:
smiles_can = oechem.OECreateCanSmiString(mol)
ims = oechem.oemolistream()
ims.SetFormat(oechem.OEFormat_SMI)
ims.openstring(smiles_ref)
mol_ref = next(ims.GetOEMols())
smiles_ref = oechem.OECreateCanSmiString(mol_ref)
assert (smiles_ref == smiles_can
), "SMILES different. Input is %s, ref is %s" % (
smiles_can,
smiles_ref,
)
from openff.toolkit.topology import Molecule
_mol = Molecule.from_openeye(mol, allow_undefined_stereo=True)
g = esp.Graph(_mol)
return g
def evaluate(self, batch):
"""batch: (B, N*D) """
# make a list of positions
batch_array = assert_numpy(batch, arr_type=_OPENMM_FLOATING_TYPE)
# reshape to (B, N, D) and add physical units
from simtk.unit import Quantity
batch_array = Quantity(value=batch_array.reshape(
batch.shape[0], -1, _SPATIAL_DIM),
unit=self._length_scale)
if self.n_workers == 1:
energies_and_forces = _compute_energy_and_force_batch(batch_array)
else:
# multiprocessing Pool
from multiprocessing import Pool
pool = Pool(
self.n_workers, initialize_worker,
(self._openmm_system, self._openmm_integrator, self._platform))
# split list into equal parts
chunksize = batch.shape[0] // self.n_workers
batch_positions_ = [
batch_array[i:i + chunksize]
for i in range(0, batch.shape[0], chunksize)
]
energies_and_forces_ = pool.map(_compute_energy_and_force_batch,
batch_positions_)
# concat lists
energies_and_forces = []
for ef in energies_and_forces_:
energies_and_forces += ef
pool.close()
# remove units
energies = [self._reduce_units(ef[0]) for ef in energies_and_forces]
if not np.all(np.isfinite(energies)):
if _err_handling == "warning":
warnings.warn("Infinite energy.")
if _err_handling == "exception":
raise ValueError("Infinite energy.")
# remove units
forces = [
np.ravel(self._reduce_units(ef[1]) * self._length_scale)
for ef in energies_and_forces
]
# to PyTorch tensors
energies = torch.tensor(energies).to(batch).reshape(-1, 1)
forces = torch.tensor(forces).to(batch)
# store
self.last_energies = energies
self.last_forces = forces
return energies, forces
