Merge toolkit caching into topology biopolymer refactor

devtools/conda-envs/openeye.yaml

@@ -24,6 +24,7 @@ dependencies:
   - openmm
   - openff-forcefields
   - smirnoff99Frosst
+  - cachetools
   - pyyaml
   - toml
   - bson

devtools/conda-envs/rdkit.yaml

@@ -23,6 +23,7 @@ dependencies:
   - openmm
   - openff-forcefields
   - smirnoff99Frosst
+  - cachetools
   - pyyaml
   - toml
   - bson

devtools/conda-envs/test_env.yaml

@@ -25,6 +25,7 @@ dependencies:
   - openmm
   - openff-forcefields
   - smirnoff99Frosst
+  - cachetools
   - pyyaml
   - toml
   - bson

openff/toolkit/topology/molecule.py

@@ -38,6 +38,7 @@
 from collections import OrderedDict
 from copy import deepcopy
 from typing import Optional, Union
+from cached_property import cached_property
 
 import networkx as nx
 import numpy as np
@@ -459,7 +460,7 @@ def virtual_sites(self):
         # for vsite in self._vsites:
         #    yield vsite
 
-    @property
+    @cached_property
     def molecule_atom_index(self):
         """
         The index of this Atom within the the list of atoms in ``Molecules``.
@@ -2040,6 +2041,20 @@ def __hash__(self):
         """
         return hash(self.to_smiles())
 
+    #@cached_property
+    def ordered_connection_table_hash(self):
+        if self._ordered_connection_table_hash is not None:
+            return self._ordered_connection_table_hash
+
+        id = ''
+        for atom in self.atoms:
+            id += f'{atom.element.symbol}_{atom.formal_charge}_{atom.stereochemistry}__'
+        for bond in self.bonds:
+            id += f'{bond.bond_order}_{bond.stereochemistry}_{bond.atom1_index}_{bond.atom2_index}__'
+        #return hash(id)
+        self._ordered_connection_table_hash = hash(id)
+        return self._ordered_connection_table_hash
+
     @classmethod
     def from_dict(cls, molecule_dict):
         """
@@ -3032,6 +3047,10 @@ def _invalidate_cached_properties(self):
         self._cached_smiles = None
         # TODO: Clear fractional bond orders
         self._rings = None
+        self._ordered_connection_table_hash = None
+        for atom in self.atoms:
+            if 'molecule_atom_index' in atom.__dict__:
+                del atom.__dict__['molecule_atom_index']
 
     def to_networkx(self):
         """Generate a NetworkX undirected graph from the Molecule.
@@ -5062,8 +5081,9 @@ def remap(self, mapping_dict, current_to_new=True):
 
         return new_molecule
 
-    @OpenEyeToolkitWrapper.requires_toolkit()
-    def to_openeye(self, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
+    def to_openeye(self,
+                   toolkit_registry=GLOBAL_TOOLKIT_REGISTRY,
+                   aromaticity_model=DEFAULT_AROMATICITY_MODEL):
         """
         Create an OpenEye molecule
 
@@ -5093,8 +5113,8 @@ def to_openeye(self, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
         >>> oemol = molecule.to_openeye()
 
         """
-        toolkit = OpenEyeToolkitWrapper()
-        return toolkit.to_openeye(self, aromaticity_model=aromaticity_model)
+        #toolkit = OpenEyeToolkitWrapper()
+        return toolkit_registry.to_openeye(self, aromaticity_model=aromaticity_model)
 
     def _construct_angles(self):
         """

openff/toolkit/utils/toolkits.py

@@ -65,6 +65,7 @@
 from distutils.spawn import find_executable
 from functools import wraps
 from typing import TYPE_CHECKING, List, Optional, Tuple
+from cachetools import LRUCache, cached
 
 import numpy as np
 from simtk import unit
@@ -168,6 +169,10 @@ class AntechamberNotFoundError(MessageException):
 # UTILITY FUNCTIONS
 # =============================================================================================
 
+def mol_to_ctab_and_aro_key(self, molecule, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
+    return f"{molecule.ordered_connection_table_hash()}-{aromaticity_model}"
+
+
 # =============================================================================================
 # CHEMINFORMATICS TOOLKIT WRAPPERS
 # =============================================================================================
@@ -1482,8 +1487,11 @@ def describe_oeatom(oeatom):
 
         return molecule
 
-    @staticmethod
-    def to_openeye(molecule, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
+
+
+    to_openeye_cache = LRUCache(maxsize=4096)
+    @cached(to_openeye_cache, key=mol_to_ctab_and_aro_key)
+    def to_openeye(self, molecule, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
         """
         Create an OpenEye molecule using the specified aromaticity model
 
@@ -2653,6 +2661,7 @@ def _find_smarts_matches(
 
         # Make a copy of molecule so we don't influence original (probably safer than deepcopy per C Bayly)
         mol = oechem.OEMol(oemol)
+
         # Set up query
         qmol = oechem.OEQMol()
         if not oechem.OEParseSmarts(qmol, smarts):
@@ -2672,13 +2681,13 @@ def _find_smarts_matches(
 
         # OEPrepareSearch will clobber our desired aromaticity model if we don't sync up mol and qmol ahead of time
         # Prepare molecule
-        oechem.OEClearAromaticFlags(mol)
-        oechem.OEAssignAromaticFlags(mol, oearomodel)
+        #oechem.OEClearAromaticFlags(mol)
+        #oechem.OEAssignAromaticFlags(mol, oearomodel)
 
         # If aromaticity model was provided, prepare query molecule
         oechem.OEClearAromaticFlags(qmol)
         oechem.OEAssignAromaticFlags(qmol, oearomodel)
-        oechem.OEAssignHybridization(mol)
+        #oechem.OEAssignHybridization(mol)
         oechem.OEAssignHybridization(qmol)
 
         # Build list of matches
@@ -4175,8 +4184,9 @@ def from_rdkit(self, rdmol, allow_undefined_stereo=False, _cls=None):
             offmol.partial_charges = None
         return offmol
 
-    @classmethod
-    def to_rdkit(cls, molecule, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
+    to_rdkit_cache = LRUCache(maxsize=4096)
+    @cached(to_rdkit_cache, key=mol_to_ctab_and_aro_key)
+    def to_rdkit(self, molecule, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
         """
         Create an RDKit molecule
 
@@ -4341,7 +4351,7 @@ def to_rdkit(cls, molecule, aromaticity_model=DEFAULT_AROMATICITY_MODEL):
             raise RuntimeError(err_msg)
 
         # Copy bond stereo info from molecule to rdmol.
-        cls._assign_rdmol_bonds_stereo(molecule, rdmol)
+        self._assign_rdmol_bonds_stereo(molecule, rdmol)
 
         # Set coordinates if we have them
         if molecule._conformers:
@@ -4522,14 +4532,14 @@ def _find_smarts_matches(rdmol, smirks, aromaticity_model="OEAroModel_MDL"):
         from rdkit import Chem
 
         # Make a copy of the molecule
-        rdmol = Chem.Mol(rdmol)
+        #rdmol = Chem.Mol(rdmol)
         # Use designated aromaticity model
-        if aromaticity_model == "OEAroModel_MDL":
-            Chem.SanitizeMol(rdmol, Chem.SANITIZE_ALL ^ Chem.SANITIZE_SETAROMATICITY)
-            Chem.SetAromaticity(rdmol, Chem.AromaticityModel.AROMATICITY_MDL)
-        else:
-            # Only the OEAroModel_MDL is supported for now
-            raise ValueError("Unknown aromaticity model: {}".aromaticity_models)
+        #if aromaticity_model == "OEAroModel_MDL":
+        #    Chem.SanitizeMol(rdmol, Chem.SANITIZE_ALL ^ Chem.SANITIZE_SETAROMATICITY)
+        #    Chem.SetAromaticity(rdmol, Chem.AromaticityModel.AROMATICITY_MDL)
+        #else:
+        #    # Only the OEAroModel_MDL is supported for now
+        #    raise ValueError("Unknown aromaticity model: {}".aromaticity_models)
 
         # Set up query.
         qmol = Chem.MolFromSmarts(smirks)  # cannot catch the error