Handle NaN values in the adjacency matrix.
@@ -45,12 +45,19 @@ class RSGCNPreprocessor(MolPreprocessor):
atom_array = construct_atomic_number_array(mol, out_size=self.out_size)
adj_array = construct_adj_matrix(mol, out_size=self.out_size)
- # adjust adjacent matrix
+ num_atoms = mol.GetNumAtoms()
degree_vec = numpy.sum(adj_array, axis=1)
- degree_sqrt_inv = 1. / numpy.sqrt(degree_vec)
+
+ # Add self-connections to avoid having NaN values in the matrix.
+ degree_sqrt_inv = numpy.zeros(degree_vec.shape)
+ mask = degree_vec > 0
+ degree_sqrt_inv[mask] = 1. / numpy.sqrt(degree_vec[mask])
+
+ # Adjust the adjacency matrix.
adj_array *= numpy.broadcast_to(degree_sqrt_inv[:, None],
adj_array.shape)
adj_array *= numpy.broadcast_to(degree_sqrt_inv[None, :],
adj_array.shape)
+ numpy.fill_diagonal(adj_array[num_atoms:, num_atoms:], 1)
return atom_array, adj_array
