Check lookup tables for values before ChemicalDomains

docs/CHANGELOG.md

@@ -26,6 +26,9 @@ The rules for this file:
 ### Changed
 - Removed unused, undocumented code paths, and updated docs (PR #132)
 
+### Fixed
+- Check lookup tables for allowed molecules before ChemicalDomain for forbidden ones (PR #145, Issue #144)
+
 
 ## v0.4.0 -- 2024-07-18
 

openff/nagl/nn/_models.py

@@ -217,6 +217,7 @@ def compute_properties(
         else:
             tensor = torch.empty
         for property_name, value in results[0].items():
+            print(as_numpy, value, value.dtype)
             combined_results[property_name] = tensor(
                 molecule.n_atoms,
                 dtype=value.dtype
@@ -278,24 +279,13 @@ def _compute_properties(
         -------
         result: Dict[str, torch.Tensor] or Dict[str, numpy.ndarray]
         """
-        if check_domains:
-            is_supported, error = self.chemical_domain.check_molecule(
-                molecule, return_error_message=True
-            )
-            if not is_supported:
-                if error_if_unsupported:
-                    raise ValueError(error)
-                else:
-                    warnings.warn(error)
 
-        try:
-            values = self._compute_properties_dgl(molecule)
-        except (MissingOptionalDependencyError, TypeError):
-            values = self._compute_properties_nagl(molecule)
-
+        values = {}
+
+        expected_value_keys = list(self.readout_modules.keys())
+
         if check_lookup_table and self.lookup_tables:
-            property_names = list(values)
-            for property_name in property_names:
+            for property_name in expected_value_keys:
                 try:
                     value = self._check_property_lookup_table(
                         molecule=molecule,
@@ -312,6 +302,29 @@ def _compute_properties(
                     )
                     values[property_name] = value
 
+
+        computed_value_keys = set(values.keys())
+        if computed_value_keys == set(expected_value_keys):
+            if as_numpy:
+                values = {k: v.detach().numpy().flatten() for k, v in values.items()}
+            return values
+
+        if check_domains:
+            is_supported, error = self.chemical_domain.check_molecule(
+                molecule, return_error_message=True
+            )
+            if not is_supported:
+                if error_if_unsupported:
+                    raise ValueError(error)
+                else:
+                    warnings.warn(error)
+
+        try:
+            values = self._compute_properties_dgl(molecule)
+        except (MissingOptionalDependencyError, TypeError):
+            values = self._compute_properties_nagl(molecule)
+
+
         if as_numpy:
             values = {k: v.detach().numpy().flatten() for k, v in values.items()}
         return values

openff/nagl/tests/nn/test_model.py

@@ -426,6 +426,12 @@ def test_compute_property(
 
         assert_allclose(charges, expected_charges, atol=1e-5)
 
+    def test_lookup_table_before_chemical_domain(self, model):
+        hcl = Molecule.from_mapped_smiles("[Cl:1][H:2]")
+        expected_charges = [-0.1680,  0.1680]
+        charges = model.compute_property(hcl, as_numpy=True, check_domains=True)
+        assert_allclose(charges, expected_charges, atol=1e-5)
+
     @pytest.mark.xfail(reason="Model does not include 0 bonds as feature")
     def test_assign_partial_charges_to_ion(self, model):
         mol = Molecule.from_smiles("[Cl-]")