Implement NAGLChargesHandler

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",

openff/toolkit/_tests/test_parameters.py

@@ -2722,6 +2722,105 @@ 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"
+
+    def test_nagl_charges_handler_with_optional_fields(self):
+        from openff.toolkit.typing.engines.smirnoff import NAGLChargesHandler
+
+        # Test with model_file_hash
+        handler = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            model_file_hash="144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0",
+            skip_version_check=True
+        )
+        assert handler.model_file == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+        assert handler.model_file_hash == "144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0"
+        assert handler.digital_object_identifier is None
+
+        # Test with digital_object_identifier
+        handler = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            digital_object_identifier="10.5072/zenodo.203601",
+            skip_version_check=True
+        )
+        assert handler.model_file == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+        assert handler.model_file_hash is None
+        assert handler.digital_object_identifier == "10.5072/zenodo.203601"
+
+        # Test with both optional fields
+        handler = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            model_file_hash="144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0",
+            digital_object_identifier="10.5072/zenodo.203601",
+            skip_version_check=True
+        )
+        assert handler.model_file == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+        assert handler.model_file_hash == "144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0"
+        assert handler.digital_object_identifier == "10.5072/zenodo.203601"
+
+    def test_nagl_charges_handler_serialization_with_optional_fields(self):
+        from openff.toolkit.typing.engines.smirnoff import NAGLChargesHandler
+
+        # Test serialization with all fields
+        handler = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            model_file_hash="144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0",
+            digital_object_identifier="10.5072/zenodo.203601",
+            skip_version_check=True
+        )
+        handler_dict = handler.to_dict()
+        assert handler_dict["model_file"] == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+        assert handler_dict["model_file_hash"] == "144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0"
+        assert handler_dict["digital_object_identifier"] == "10.5072/zenodo.203601"
+
+        # Test deserialization via constructor
+        handler_from_dict = NAGLChargesHandler(**handler_dict)
+        assert handler_from_dict.model_file == "openff-gnn-am1bcc-0.1.0-rc.3.pt"
+        assert handler_from_dict.model_file_hash == "144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0"
+        assert handler_from_dict.digital_object_identifier == "10.5072/zenodo.203601"
+
+    def test_nagl_charges_handler_compatibility(self):
+        from openff.toolkit.typing.engines.smirnoff import NAGLChargesHandler
+        from openff.toolkit.utils.exceptions import IncompatibleParameterError
+
+        # Test compatible handlers (same model_file)
+        handler1 = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            skip_version_check=True
+        )
+        handler2 = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            model_file_hash="144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0",
+            skip_version_check=True
+        )
+        # Should not raise exception
+        handler1.check_handler_compatibility(handler2)
+
+        # Test incompatible handlers (different model_file)
+        handler3 = NAGLChargesHandler(
+            model_file="different-model-file.pt",
+            skip_version_check=True
+        )
+        with pytest.raises(IncompatibleParameterError, match="different model_files"):
+            handler1.check_handler_compatibility(handler3)
+
+    def test_nagl_charges_handler_defaults(self):
+        from openff.toolkit.typing.engines.smirnoff import NAGLChargesHandler
+
+        # Test that optional fields default to None
+        handler = NAGLChargesHandler(
+            model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+            skip_version_check=True
+        )
+        assert handler.model_file_hash is None
+        assert handler.digital_object_identifier is None
+
 
 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,102 @@ def find_matches(self, entity, unique=False):
             unique=unique,
         )
 
+class NAGLChargesHandler(_NonbondedHandler):
+    """ParameterHandler for applying partial charges from a pretrained NAGL model.
+
+    This handler processes the NAGLCharges section of SMIRNOFF force fields, which
+    specifies a pre-trained NAGL model for computing
+    partial charges on molecules.
+
+    Parameters
+    ----------
+    model_file : str
+        Path to the PyTorch model file (e.g., "openff-gnn-am1bcc-0.1.0-rc.3.pt").
+        This is the model that will be used for charge assignment.
+    model_file_hash : str, optional
+        SHA-256 hash of the model file for integrity verification. When provided,
+        the hash will be validated against the actual model file.
+    digital_object_identifier : str, optional
+        Zenodo DOI that can be used to retrieve the model file if it's not found
+        locally. Must point to a Zenodo record with an attached file matching
+        the model_file name.
+    version : str, optional
+        The version of the NAGLCharges section specification.
+    skip_version_check : bool, optional, default=False
+        If True, skips validation of the version parameter and sets it to the highest
+        supported version.
+    allow_cosmetic_attributes : bool, optional, default=False
+        If True, allows non-specification attributes to be present.
+
+    Examples
+    --------
+    Create a handler with just the model file:
+
+    >>> handler = NAGLChargesHandler(
+    ...     model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+    ...     skip_version_check=True
+    ... )
+
+    Create a handler with hash verification:
+
+    >>> handler = NAGLChargesHandler(
+    ...     model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+    ...     model_file_hash="144ed56e46c5b3ad80157b342c8c0f8f7340e4d382a678e30dd300c811646bd0",
+    ...     skip_version_check=True
+    ... )
+
+    Create a handler with DOI for model retrieval:
+
+    >>> handler = NAGLChargesHandler(
+    ...     model_file="openff-gnn-am1bcc-0.1.0-rc.3.pt",
+    ...     digital_object_identifier="10.5072/zenodo.203601",
+    ...     skip_version_check=True
+    ... )
+
+    Notes
+    -----
+    NAGLChargesHandler compatibility is determined solely by the model_file parameter. Two
+    handlers are compatible if and only if they specify the same model_file,
+    regardless of the values of model_file_hash or digital_object_identifier.
+
+    The actual model loading, hash verification, and DOI-based retrieval are
+    handled by the openff-nagl-models package, not by this handler directly.
+    """
+
+    _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)
+    model_file_hash = ParameterAttribute(default=None, converter=str)
+    digital_object_identifier = ParameterAttribute(default=None, 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 +3358,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 +3479,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,