Indexing error for two-atom molecules

[lilyminium]

Noticed while using Espaloma to assign charges to a batch of molecules -- as far as I can tell, any 2-atom molecule results in an indexing error when trying to create a esp.Graph. Example input:

In [1]: %load mve.py

In [2]: # %load mve.py
   ...: import espaloma as esp
   ...: from openff.toolkit import Molecule
   ...:
   ...: molecule = Molecule.from_mapped_smiles("[N:1]#[N:2]")
   ...:
   ...: # create an Espaloma Graph object to represent the molecule of interest
   ...: molecule_graph = esp.Graph(molecule)
   ...:
   ...: # load pretrained model
   ...: espaloma_model = esp.get_model("latest")
   ...:
   ...: # apply a trained espaloma model to assign parameters
   ...: espaloma_model(molecule_graph.heterograph)
   ...:
   ...: # create an OpenMM System for the specified molecule
   ...: openmm_system = esp.graphs.deploy.openmm_system_from_graph(molecule_graph)
   ...:
LICENSE: Could not open license file "oe_license.txt" in local directory
LICENSE: N.B. OE_LICENSE environment variable is not set
LICENSE: N.B. OE_DIR environment variable is not set
LICENSE: No product keys!
LICENSE: No product keys!
LICENSE: No product keys!
LICENSE: No product keys!
The OpenEye Toolkits are found to be installed but not licensed and therefore will not be used.
The OpenEye Toolkits require a (free for academics) license, see https://docs.eyesopen.com/toolkits/python/quickstart-python/license.html
/lila/home/lilywang/micromamba/envs/espaloma-0.3.2/lib/python3.11/site-packages/dgl/heterograph.py:92: DGLWarning: Recommend creating graphs by `dgl.graph(data)` instead of `dgl.DGLGraph(data)`.
  dgl_warning(
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[2], line 8
      5 molecule = Molecule.from_mapped_smiles("[N:1]#[N:2]")
      7 # create an Espaloma Graph object to represent the molecule of interest
----> 8 molecule_graph = esp.Graph(molecule)
     10 # load pretrained model
     11 espaloma_model = esp.get_model("latest")

File /lila/home/lilywang/micromamba/envs/espaloma-0.3.2/lib/python3.11/site-packages/espaloma/graphs/graph.py:63, in Graph.__init__(self, mol, homograph, heterograph)
     60     homograph = self.get_homograph_from_mol(mol)
     62 if homograph is not None and heterograph is None:
---> 63     heterograph = self.get_heterograph_from_graph_and_mol(
     64         homograph, mol
     65     )
     67 self.mol = mol
     68 self.homograph = homograph

File /lila/home/lilywang/micromamba/envs/espaloma-0.3.2/lib/python3.11/site-packages/espaloma/graphs/graph.py:137, in Graph.get_heterograph_from_graph_and_mol(graph, mol)
    131 import dgl
    133 assert isinstance(
    134     graph, dgl.DGLGraph
    135 ), "graph can only be dgl Graph object."
--> 137 heterograph = esp.graphs.utils.read_heterogeneous_graph.from_homogeneous_and_mol(
    138     graph, mol
    139 )
    141 return heterograph

File /lila/home/lilywang/micromamba/envs/espaloma-0.3.2/lib/python3.11/site-packages/espaloma/graphs/utils/read_heterogeneous_graph.py:150, in from_homogeneous_and_mol(g, offmol)
    136     for big_idx in range(small_idx + 1, 5):
    137         for pos_idx in range(big_idx - small_idx + 1):
    139             hg[
    140                 (
    141                     "n%s" % small_idx,
    142                     "n%s_as_%s_in_n%s" % (small_idx, pos_idx, big_idx),
    143                     "n%s" % big_idx,
    144                 )
    145             ] = np.stack(
    146                 [
    147                     np.array(
    148                         [
    149                             idxs_to_ordering["n%s" % small_idx][tuple(x)]
--> 150                             for x in idxs["n%s" % big_idx][
    151                                 :, pos_idx : pos_idx + small_idx
    152                             ]
    153                         ]
    154                     ),
    155                     np.arange(idxs["n%s" % big_idx].shape[0]),
    156                 ],
    157                 axis=1,
    158             )
    160             hg[
    161                 (
    162                     "n%s" % big_idx,
   (...)
    178                 axis=1,
    179             )
    181 # ======================================
    182 # nonbonded terms
    183 # ======================================
   (...)
    187
    188 # make dense

IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed

Software versions
Espaloma 0.3.2

[mikemhenry]

@lilyminium thanks for the report! @yuanqing-wang Can you take a look at this?

[yuanqing-wang]

Sorry for the delay, @lilyminium . Yes a hydrogen would indeed trigger an error, which was the designed behavior, since we assume that there are always some angles in the system. Maybe we should consider flagging this earlier.