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