update lookups

openff/nagl/lookups.py

@@ -14,9 +14,22 @@
 if typing.TYPE_CHECKING:
     from openff.toolkit.topology import Molecule
 
+
+class PropertyProvenance(ImmutableModel):
+    description: str = Field(
+        description="A description of the provenance"
+    )
+    versions: dict[str, str] = Field(
+        default_factory=dict,
+        description="The versions of the relevant software used to compute the property"
+    )
+
 class BasePropertiesLookupTableEntry(ImmutableModel):
-    inchi_key: str = Field(
-        description="The InChI key of the molecule"
+    inchi: str = Field(
+        description="The InChI of the molecule"
+    )
+    provenance: PropertyProvenance = Field(
+        description="The provenance of the property value"
     )
 
 class AtomPropertiesLookupTableEntry(BasePropertiesLookupTableEntry):
@@ -80,7 +93,7 @@ def _convert_property_lookup_table(cls, v):
             raise ValueError("All entries must be AtomPropertiesLookupTableEntry instances")
 
         return types.MappingProxyType({
-            entry.inchi_key: entry
+            entry.inchi: entry
             for entry in v
         })
 
@@ -118,7 +131,7 @@ def lookup(self, molecule: "Molecule") -> torch.Tensor:
         try:
             entry = self.properties[inchi_key]
         except KeyError:
-            raise KeyError(f"Could not find property value for molecule with InChI key {inchi_key}")
+            raise KeyError(f"Could not find property value for molecule with InChI {inchi_key}")
 
         assert len(entry) == molecule.n_atoms
 

openff/nagl/nn/_models.py

@@ -57,6 +57,22 @@ def _forward_unpostprocessed(self, molecule: "DGLMoleculeOrBatch"):
 
 
 class GNNModel(BaseGNNModel):
+    """
+    A GNN model for predicting properties of molecules.
+
+    Parameters
+    ----------
+    config: ModelConfig or dict
+        The configuration for the model.
+
+    chemical_domain: ChemicalDomain or dict
+        The applicable chemical domain for the model.
+
+    lookup_tables: dict
+        A dictionary of lookup tables for properties.
+        The keys should be the property names, and the values
+        should be instances of :class:`~openff.nagl.lookups.BaseLookupTable`.
+    """
     def __init__(
         self,
         config: ModelConfig,
@@ -107,12 +123,16 @@ def __init__(
             readout_modules=readout_modules,
         )
 
+        lookup_tables_dict = {}
+        for k, v in valid_lookup_tables.items():
+            v_ = v.dict()
+            v_["properties"] = dict(v_["properties"])
+            lookup_tables_dict[k] = v_
+
         self.save_hyperparameters({
             "config": config.dict(),
             "chemical_domain": chemical_domain.dict(),
-            "lookup_tables": {
-                k: v.dict() for k, v in valid_lookup_tables.items()
-            },
+            "lookup_tables": lookup_tables_dict,
         })
         self.config = config
         self.chemical_domain = chemical_domain

openff/nagl/tests/test_lookups.py

@@ -15,18 +15,20 @@
 @pytest.fixture(scope="function")
 def CNO2_entry():
     return AtomPropertiesLookupTableEntry(
-        inchi_key="InChI=1/CH3NO2/c1-2(3)4/h1H3",
+        inchi="InChI=1/CH3NO2/c1-2(3)4/h1H3",
         mapped_smiles="[H:5][C:1]([H:6])([H:7])[N+:2](=[O:3])[O-:4]",
         property_value=[-0.103, 0.234, -0.209, -0.209, 0.096, 0.096, 0.096],
+        provenance={"description": "test"}
     )
 
 
 @pytest.fixture(scope="function")
 def SH2_entry():
     return AtomPropertiesLookupTableEntry(
-        inchi_key="InChI=1/H2S/h1H2",
+        inchi="InChI=1/H2S/h1H2",
         mapped_smiles="[H:2][S:1][H:3]",
         property_value=[-0.441, 0.22, 0.22],
+        provenance={"description": "test"}
     )
 
 
@@ -37,8 +39,8 @@ def lookup_table(self, CNO2_entry, SH2_entry):
         return AtomPropertiesLookupTable(
             property_name="test",
             properties={
-                CNO2_entry.inchi_key: CNO2_entry,
-                SH2_entry.inchi_key: SH2_entry,
+                CNO2_entry.inchi: CNO2_entry,
+                SH2_entry.inchi: SH2_entry,
             }
         )