Merge pull request #8 from ecrl/dev

Dev
ecrl · Jun 15, 2021 · 94361ee · 94361ee
2 parents bac5dda + 7ff945d
commit 94361ee
Show file tree

Hide file tree

Showing 14 changed files with 1,014 additions and 131 deletions.
diff --git a/examples/model.enc b/examples/model.enc
diff --git a/examples/model.pt b/examples/model.pt
diff --git a/examples/model_cn.enc b/examples/model_cn.enc
diff --git a/examples/model_cn.pt b/examples/model_cn.pt
diff --git a/examples/model_mon.enc b/examples/model_mon.enc
diff --git a/examples/model_mon.pt b/examples/model_mon.pt
diff --git a/examples/model_ron.enc b/examples/model_ron.enc
diff --git a/examples/model_ron.pt b/examples/model_ron.pt
diff --git a/examples/predict_cn.ipynb b/examples/predict_cn.ipynb
diff --git a/examples/predict_mon.ipynb b/examples/predict_mon.ipynb
diff --git a/examples/predict_ron.ipynb b/examples/predict_ron.ipynb
diff --git a/graphchem/models/gcn.py b/graphchem/models/gcn.py
@@ -39,13 +39,15 @@ def __init__(self, node_dim: int, edge_dim: int, output_dim: int,
         self._n_messages = n_messages
 
         self.lin0 = nn.Linear(node_dim, node_dim)
+        self.lin0_edge = nn.Linear(edge_dim, edge_dim)
 
         # Construct message passing layers for node, edge networks
         self.node_conv = pyg_nn.MFConv(node_dim, node_dim)
         self.edge_conv = pyg_nn.EdgeConv(nn.Sequential(
             nn.Linear(2 * edge_dim, edge_dim)
         ))
         self.node_gru = nn.GRU(node_dim, node_dim)
+        self.edge_gru = nn.GRU(edge_dim, edge_dim)
 
         # Construct post-message passing layers
         self.post_conv = nn.ModuleList()
@@ -82,7 +84,9 @@ def forward(self, data: 'torch_geometric.data.Data') -> Tuple[
             x = torch.ones(data.num_nodes, 1)
 
         out = F.relu(self.lin0(x))
+        out_edge = F.relu(self.lin0_edge(edge_attr))
         h = out.unsqueeze(0)
+        h_edge = out_edge.unsqueeze(0)
 
         # Feed forward, node and edge messages
         for i in range(self._n_messages):
@@ -91,14 +95,15 @@ def forward(self, data: 'torch_geometric.data.Data') -> Tuple[
             m = F.dropout(m, p=self._dropout, training=self.training)
             out, h = self.node_gru(m.unsqueeze(0), h)
             out = out.squeeze(0)
-            edge_attr = self.edge_conv(edge_attr, edge_index)
-            emb_edge = edge_attr
-            edge_attr = F.relu(edge_attr)
-            edge_attr = F.dropout(edge_attr, p=self._dropout,
-                                  training=self.training)
+
+            m_edge = F.relu(self.edge_conv(out_edge, edge_index))
+            emb_edge = m_edge
+            m_edge = F.dropout(m_edge, p=self._dropout, training=self.training)
+            out_edge, h_edge = self.edge_gru(m_edge.unsqueeze(0), h_edge)
+            out_edge = out_edge.squeeze(0)
 
         # Concatenate node network and edge network output tensors
-        out = torch.cat([out[row], edge_attr[col]], dim=1)
+        out = torch.cat([out[row], out_edge[col]], dim=1)
 
         # Perform scatter add, reshape to original node dimensionality
         out = scatter_add(out, col, dim=0, dim_size=x.size(0))

diff --git a/graphchem/preprocessing/features.py b/graphchem/preprocessing/features.py
@@ -136,9 +136,15 @@ def encode(self, smiles: str) -> Tuple['torch.tensor', 'torch.tensor',
         connectivity = np.zeros((2, 2 * mol.GetNumBonds()))
         bond_index = 0
         for atom in mol.GetAtoms():
+            start_index = atom.GetIdx()
             for bond in atom.GetBonds():
-                connectivity[0, bond_index] = bond.GetBeginAtomIdx()
-                connectivity[1, bond_index] = bond.GetEndAtomIdx()
+                rev = bond.GetBeginAtomIdx() != start_index
+                if not rev:
+                    connectivity[0, bond_index] = bond.GetBeginAtomIdx()
+                    connectivity[1, bond_index] = bond.GetEndAtomIdx()
+                else:
+                    connectivity[0, bond_index] = bond.GetEndAtomIdx()
+                    connectivity[1, bond_index] = bond.GetBeginAtomIdx()
                 bond_index += 1
         connectivity = torch.from_numpy(connectivity).type(torch.long)
 

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='graphchem',
-    version='1.1.0',
+    version='1.2.0',
     description='Graph-based machine learning for chemical property prediction',
     url='https://github.com/ecrl/graphchem',
     author='Travis Kessler',