Add ion capabilities (#146)

* add dgl ion capabilities

* move files around and update tests

* add test model and tidy files

* add ion tests

* remove commented out lines

* add pointer to apache 2 license

* update changelog [skip ci]

Author: lilyminium

docs/CHANGELOG.md

@@ -28,6 +28,7 @@ The rules for this file:
 
 ### Fixed
 - Check lookup tables for allowed molecules before ChemicalDomain for forbidden ones (PR #145, Issue #144)
+- Add support for single atoms (PR #146, Issue #138)
 
 
 ## v0.4.0 -- 2024-07-18

openff/nagl/molecule/_dgl/utils.py

@@ -1,6 +1,7 @@
 from typing import Dict, List, TYPE_CHECKING, Optional
 
 import torch
+import numpy as np
 from openff.utilities import requires_package
 
 from openff.nagl.features.atoms import AtomFeature
@@ -39,7 +40,8 @@ def openff_molecule_to_base_dgl_graph(
         {
             ("atom", forward, "atom"): (indices_a, indices_b),
             ("atom", reverse, "atom"): (indices_b, indices_a),
-        }
+        },
+        num_nodes_dict={"atom": molecule.n_atoms},
     )
     return molecule_graph
 
@@ -99,7 +101,7 @@ def openff_molecule_to_dgl_graph(
 
     for direction in (forward, reverse):
         n_bonds = len(molecule.bonds)
-        if bond_feature_tensor is not None:
+        if bond_feature_tensor is not None and n_bonds:
             bond_feature_tensor = bond_feature_tensor.reshape(n_bonds, -1)
         else:
             bond_feature_tensor = torch.zeros((n_bonds, 0))
@@ -108,13 +110,87 @@ def openff_molecule_to_dgl_graph(
 
     return molecule_graph
 
+@requires_package("dgl")
+def heterograph_to_homograph_no_edges(G: "dgl.DGLHeteroGraph", ndata=None, edata=None) -> "dgl.DGLGraph":
+    """
+    Copied and modified from dgl.python.dgl.convert.to_homogeneous,
+    but with the edges removed.
+
+    This part of the code is licensed under the Apache 2.0 license according
+    to the terms of DGL (https://github.com/dmlc/dgl?tab=Apache-2.0-1-ov-file).
+
+    Please see our third-party license file for more information
+    (https://github.com/openforcefield/openff-nagl/blob/main/LICENSE-3RD-PARTY)
+    """
+    import dgl
+    from dgl import backend as F
+    from dgl.base import EID, NID, ETYPE, NTYPE
+    from dgl.heterograph import combine_frames
+
+    # TODO: revisit in case DGL accounts for this in the future
+    num_nodes_per_ntype = [G.num_nodes(ntype) for ntype in G.ntypes]
+    offset_per_ntype = np.insert(np.cumsum(num_nodes_per_ntype), 0, 0)
+    srcs = []
+    dsts = []
+    nids = []
+    eids = []
+    ntype_ids = []
+    etype_ids = []
+    total_num_nodes = 0
+
+    for ntype_id, ntype in enumerate(G.ntypes):
+        num_nodes = G.num_nodes(ntype)
+        total_num_nodes += num_nodes
+        # Type ID is always in int64
+        ntype_ids.append(F.full_1d(num_nodes, ntype_id, F.int64, G.device))
+        nids.append(F.arange(0, num_nodes, G.idtype, G.device))
+
+    for etype_id, etype in enumerate(G.canonical_etypes):
+        srctype, _, dsttype = etype
+        src, dst = G.all_edges(etype=etype, order="eid")
+        num_edges = len(src)
+        srcs.append(src + int(offset_per_ntype[G.get_ntype_id(srctype)]))
+        dsts.append(dst + int(offset_per_ntype[G.get_ntype_id(dsttype)]))
+        etype_ids.append(F.full_1d(num_edges, etype_id, F.int64, G.device))
+        eids.append(F.arange(0, num_edges, G.idtype, G.device))
+
+    retg = dgl.graph(
+        (F.cat(srcs, 0), F.cat(dsts, 0)),
+        num_nodes=total_num_nodes,
+        idtype=G.idtype,
+        device=G.device,
+    )
+
+    # copy features
+    if ndata is None:
+        ndata = []
+    if edata is None:
+        edata = []
+    comb_nf = combine_frames(
+        G._node_frames, range(len(G.ntypes)), col_names=ndata
+    )
+    if comb_nf is not None:
+        retg.ndata.update(comb_nf)
+
+    retg.ndata[NID] = F.cat(nids, 0)
+    retg.edata[EID] = F.cat(eids, 0)
+    retg.ndata[NTYPE] = F.cat(ntype_ids, 0)
+    retg.edata[ETYPE] = F.cat(etype_ids, 0)
+
+    return retg
+
+
+
 
 @requires_package("dgl")
 def dgl_heterograph_to_homograph(graph: "dgl.DGLHeteroGraph") -> "dgl.DGLGraph":
     import dgl
 
     try:
         homo_graph = dgl.to_homogeneous(graph, ndata=[FEATURE], edata=[FEATURE])
+    except TypeError as e:
+        if graph.num_edges() == 0:
+            homo_graph = heterograph_to_homograph_no_edges(graph)
     except KeyError:
         # A nasty workaround to check when we don't have any atom / bond features as
         # DGL doesn't allow easy querying of features dicts for hetereographs with

openff/nagl/molecule/_graph/_graph.py

@@ -154,7 +154,13 @@ def in_edges(self, nodes, form="uv"):
             raise ValueError("Unknown form: {}".format(form))
 
     def _bond_indices(self):
-        u, v = map(list, zip(*self.graph.edges()))
+        try:
+            u, v = map(list, zip(*self.graph.edges()))
+        except ValueError as e:
+            # this may be due to there not being bonds
+            if not self.graph.edges():
+                return torch.tensor([], dtype=torch.long), torch.tensor([], dtype=torch.long)
+            raise e
         U = torch.tensor(u, dtype=torch.long)
         V = torch.tensor(v, dtype=torch.long)
         return U, V