[WIP] Handle NAGLCharges

devtools/conda-envs/openeye-examples.yaml

@@ -42,3 +42,5 @@ dependencies:
   - pdbfixer
   - openmmforcefields >=0.11.2
   - gromacs >=2023.3
+  - pip:
+      - git+https://github.com/openforcefield/openff-interchange.git@naglcharges-handler

devtools/conda-envs/openeye.yaml

@@ -51,3 +51,4 @@ dependencies:
     - types-xmltodict
     - types-cachetools
     - mongo-types
+    - git+https://github.com/openforcefield/openff-interchange.git@naglcharges-handler

devtools/conda-envs/rdkit-examples.yaml

@@ -43,3 +43,5 @@ dependencies:
   - pdbfixer
   - openmmforcefields >=0.11.2
   - gromacs >=2023.3
+  - pip:
+      - git+https://github.com/openforcefield/openff-interchange.git@naglcharges-handler

devtools/conda-envs/rdkit.yaml

@@ -40,3 +40,5 @@ dependencies:
   - qcportal >=0.50
   - qcengine
   - nglview
+  - pip:
+      - git+https://github.com/openforcefield/openff-interchange.git@naglcharges-handler

devtools/conda-envs/test_env.yaml

@@ -49,3 +49,5 @@ dependencies:
   - nbval
   # No idea why this is necessary, see https://github.com/openforcefield/openff-toolkit/pull/1821
   - nomkl
+  - pip:
+      - git+https://github.com/openforcefield/openff-interchange.git@naglcharges-handler

openff/toolkit/_tests/test_nagl.py

@@ -14,6 +14,7 @@
     create_reversed_ethanol,
 )
 from openff.toolkit._tests.utils import requires_openeye
+from openff.toolkit.utils import GLOBAL_TOOLKIT_REGISTRY
 from openff.toolkit.utils.exceptions import (
     ChargeMethodUnavailableError,
     ToolkitUnavailableException,
@@ -38,6 +39,10 @@ def test_version(self):
 
         assert parsed_version == NAGLToolkitWrapper()._toolkit_version
 
+    def test_nagl_in_global_toolkit_registry(self):
+        assert "NAGL" in GLOBAL_TOOLKIT_REGISTRY.__repr__()
+
+
     @requires_openeye
     @pytest.mark.parametrize(
         "molecule_function",
@@ -66,7 +71,7 @@ def test_assign_partial_charges_basic(self, molecule_function, nagl_model):
 
         molecule.assign_partial_charges(
             partial_charge_method=nagl_model,
-            toolkit_registry=NAGLToolkitWrapper(),
+            #toolkit_registry=NAGLToolkitWrapper(),
         )
 
         assert molecule.partial_charges is not None

openff/toolkit/_tests/test_parameters.py

@@ -2722,6 +2722,16 @@ def test_charge_increment_one_ci_missing(self):
             ],
         )
 
+class TestNAGLChargesHandler:
+    def test_nagl_charges_handler_serialization(self):
+        from openff.toolkit.typing.engines.smirnoff import NAGLChargesHandler
+        handler = NAGLChargesHandler(model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt", skip_version_check=True)
+        assert handler.model_file == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+        handler_dict = handler.to_dict()
+        assert handler_dict["model_file"] == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+
+        # TODO: test_nagl_charges_handler_are_compatible
+
 
 class TestGBSAHandler:
     def test_create_default_gbsahandler(self):

openff/toolkit/typing/engines/smirnoff/__init__.py

@@ -18,6 +18,7 @@
     IndexedParameterAttribute,
     LibraryChargeHandler,
     MappedParameterAttribute,
+    NAGLChargesHandler,
     ParameterAttribute,
     ParameterHandler,
     ParameterList,

openff/toolkit/typing/engines/smirnoff/parameters.py

@@ -29,6 +29,7 @@
     "LibraryChargeHandler",
     "LibraryChargeType",
     "MappedParameterAttribute",
+    "NAGLChargesHandler",
     "NotEnoughPointsForInterpolationError",
     "ParameterAttribute",
     "ParameterHandler",
@@ -3253,6 +3254,41 @@ def find_matches(self, entity, unique=False):
             unique=unique,
         )
 
+class NAGLChargesHandler(_NonbondedHandler):
+    """ParameterHandler for applying partial charges from a pretrained NAGL model."""
+
+    _TAGNAME = "NAGLCharges"
+    _DEPENDENCIES = [vdWHandler, ElectrostaticsHandler, LibraryChargeHandler]
+    _INFOTYPE = None  # No separate parameter types; just a model path
+    _MAX_SUPPORTED_SECTION_VERSION = Version("0.3")
+    model_file = ParameterAttribute(converter=str)
+
+    def check_handler_compatibility(
+        self,
+        other_handler: "NAGLChargesHandler",
+        assume_missing_is_default: bool = True,
+    ):
+        """
+        Checks whether this ParameterHandler encodes compatible physics as another ParameterHandler. This is
+        called if a second handler is attempted to be initialized for the same tag.
+
+        Parameters
+        ----------
+        other_handler
+            The handler to compare to.
+        assume_missing_is_default
+
+        Raises
+        ------
+        IncompatibleParameterError if handler_kwargs are incompatible with existing parameters.
+        """
+        if self.model_file != other_handler.model_file:
+            raise IncompatibleParameterError("Attempted to initialize two NAGLCharges sections with different "
+                                             "model_files: "
+                                             f"{self.model_file=} is not identical to {other_handler.model_file=}")
+
+
+
 
 class ToolkitAM1BCCHandler(_NonbondedHandler):
     """Handle SMIRNOFF ``<ToolkitAM1BCC>`` tags
@@ -3261,7 +3297,7 @@ class ToolkitAM1BCCHandler(_NonbondedHandler):
     """
 
     _TAGNAME = "ToolkitAM1BCC"  # SMIRNOFF tag name to process
-    _DEPENDENCIES = [vdWHandler, ElectrostaticsHandler, LibraryChargeHandler]
+    _DEPENDENCIES = [vdWHandler, ElectrostaticsHandler, LibraryChargeHandler, NAGLChargesHandler]
     _KWARGS = ["toolkit_registry"]  # Kwargs to catch when create_force is called
 
     def check_handler_compatibility(
@@ -3382,6 +3418,7 @@ def find_matches(self, entity, unique=False):
         return matches
 
 
+
 class GBSAHandler(ParameterHandler):
     """Handle SMIRNOFF ``<GBSA>`` tags
 

openff/toolkit/utils/nagl_wrapper.py

@@ -135,7 +135,7 @@ def assign_partial_charges(
         except FileNotFoundError as error:
             raise ChargeMethodUnavailableError(
                 f"Charge model {partial_charge_method} not supported by "
-                f"{self.__class__.__name__}."
+                f"{self.__class__.__name__}, or model file can not be found."
             ) from error
 
         model = GNNModel.load(model_path, eval_mode=True)

openff/toolkit/utils/toolkits.py

@@ -101,6 +101,7 @@
 # Create global toolkit registry, where all available toolkits are registered
 GLOBAL_TOOLKIT_REGISTRY = ToolkitRegistry(
     toolkit_precedence=[
+        NAGLToolkitWrapper,
         OpenEyeToolkitWrapper,
         RDKitToolkitWrapper,
         AmberToolsToolkitWrapper,