openfe icon indicating copy to clipboard operation
openfe copied to clipboard

Un-vendor `toolkit_registry_manager`

Open mattwthompson opened this issue 6 months ago • 1 comments

Developers certificate of origin

I came across #711 and couldn't resist at least a minimal patch. This has the benefit of leaving other stuff in place and the downside of creating a function that doesn't strictly need to exist. The actual usage remains (I think) fairly clean-looking, i.e. as distinct from

diff --git a/openfe/protocols/openmm_afe/equil_solvation_afe_method.py b/openfe/protocols/openmm_afe/equil_solvation_afe_method.py
index 3122b3f..c301bd0 100644
--- a/openfe/protocols/openmm_afe/equil_solvation_afe_method.py
+++ b/openfe/protocols/openmm_afe/equil_solvation_afe_method.py
@@ -38,6 +38,12 @@ import itertools
 import numpy as np
 import numpy.typing as npt
 from openff.units import unit
+from openff.toolkit.utils.toolkits import (
+    toolkit_registry_manager,
+    AmberToolsToolkitWrapper,
+    RDKitToolkitWrapper,
+    ToolkitRegistry,
+)
 from openmmtools import multistate
 from typing import Optional, Union
 from typing import Any, Iterable
@@ -58,7 +64,7 @@ from openfe.protocols.openmm_afe.equil_afe_settings import (
 )
 from ..openmm_utils import system_validation, settings_validation
 from .base import BaseAbsoluteUnit
-from openfe.utils import without_oechem_backend, log_system_probe
+from openfe.utils import log_system_probe
 from openfe.due import due, Doi
 
 
@@ -778,7 +784,9 @@ class AbsoluteSolvationVacuumUnit(BaseAbsoluteUnit):
     ) -> dict[str, Any]:
         log_system_probe(logging.INFO, paths=[ctx.scratch])
 
-        with without_oechem_backend():
+        with toolkit_registry_manager(
+            ToolkitRegistry([RDKitToolkitWrapper, AmberToolsToolkitWrapper()])
+        ):
             outputs = self.run(scratch_basepath=ctx.scratch,
                                shared_basepath=ctx.shared)
 

mattwthompson avatar Feb 13 '24 17:02 mattwthompson