From 943dfc56c617520c6109b2bcb0c0063335b47c3b Mon Sep 17 00:00:00 2001 From: Matt Thompson Date: Tue, 8 Mar 2022 15:12:22 -0600 Subject: [PATCH] Cache generate_conformers with hash_molecule_args_and_kwargs approach (#251) * Cache generate_conformers with hash_molecule_args_and_kwargs approach * Ensure openff-toolkit 0.10.3 is installed --- src/smirnoff_hack.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/smirnoff_hack.py b/src/smirnoff_hack.py index 2734d91d4..4c17c99b3 100644 --- a/src/smirnoff_hack.py +++ b/src/smirnoff_hack.py @@ -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 @@ -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