Skip to content

Commit

Permalink
Cache generate_conformers with hash_molecule_args_and_kwargs approach (
Browse files Browse the repository at this point in the history
…#251)

* Cache generate_conformers with hash_molecule_args_and_kwargs approach

* Ensure openff-toolkit 0.10.3 is installed
  • Loading branch information
mattwthompson authored Mar 8, 2022
1 parent c655df6 commit 9f99464
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions src/smirnoff_hack.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,10 +125,10 @@ def at_cached_assign_partial_charges(self, molecule, *args, **kwargs):
# cache the OE generate_conformers function (save 15s)
OE_TOOLKIT_CACHE_molecule_conformers = {}
oe_original_generate_conformers = OpenEyeToolkitWrapper.generate_conformers
def oe_cached_generate_conformers(self, molecule, n_conformers=1, rms_cutoff=None, clear_existing=True):
cache_key = hash((hash_molecule(molecule), n_conformers, str(rms_cutoff), clear_existing))
def oe_cached_generate_conformers(self, molecule, *args, **kwargs):
cache_key = hash_molecule_args_and_kwargs(molecule, args, kwargs)
if cache_key not in OE_TOOLKIT_CACHE_molecule_conformers:
oe_original_generate_conformers(self, molecule, n_conformers=n_conformers, rms_cutoff=rms_cutoff, clear_existing=clear_existing)
oe_original_generate_conformers(self, molecule, *args, **kwargs)
OE_TOOLKIT_CACHE_molecule_conformers[cache_key] = molecule._conformers
molecule._conformers = OE_TOOLKIT_CACHE_molecule_conformers[cache_key]
OpenEyeToolkitWrapper.generate_conformers = oe_cached_generate_conformers
Expand All @@ -137,10 +137,10 @@ def oe_cached_generate_conformers(self, molecule, n_conformers=1, rms_cutoff=Non
# cache the RDKit generate_conformers function
RDK_TOOLKIT_CACHE_molecule_conformers = {}
rdk_original_generate_conformers = RDKitToolkitWrapper.generate_conformers
def rdk_cached_generate_conformers(self, molecule, n_conformers=1, rms_cutoff=None, clear_existing=True):
cache_key = hash((hash_molecule(molecule), n_conformers, str(rms_cutoff), clear_existing))
def rdk_cached_generate_conformers(self, molecule, *args, **kwargs):
cache_key = hash_molecule_args_and_kwargs(molecule, args, kwargs)
if cache_key not in RDK_TOOLKIT_CACHE_molecule_conformers:
rdk_original_generate_conformers(self, molecule, n_conformers=n_conformers, rms_cutoff=rms_cutoff, clear_existing=clear_existing)
rdk_original_generate_conformers(self, molecule, *args, **kwargs)
RDK_TOOLKIT_CACHE_molecule_conformers[cache_key] = molecule._conformers
molecule._conformers = RDK_TOOLKIT_CACHE_molecule_conformers[cache_key]
RDKitToolkitWrapper.generate_conformers = rdk_cached_generate_conformers
Expand Down

0 comments on commit 9f99464

Please sign in to comment.