Unify molecule graph data

Currently, I seem that there are two data classes, ConvMol and WeaveMol for a molecule graph and these classes are designed for each specific model (like ConvMol <- GraphConvModel)

I think that this situation is not good and we should unify one data class for molecule graph. Generally, all models receive the same graph data class like DGL or PyG. This design leads to make codes more readable and maintainable.

I recommend DeepChem has only one molecule graph data and all GNN/GCN models receive it. And, current graph data is too complicated, so I also recommend the graph data is simpler like PyG. The PyG design is from here.

For molecule graph, I imagine the data class which holds the following attributes.

• data.x: Node feature matrix with shape [num_nodes, num_node_features]
• data.edge_index: Graph connectivity in COO format with shape [2, num_edges] and type torch.long
• data.edge_attr: Edge feature matrix with shape [num_edges, num_edge_features]
• data.graph_attr:Graph feature with shape [1, num_graph_features]
• data.y: Target to train against graph-level targets of shape [1, *]

[peastman]

I’ve also wondered whether we really need different featurizers for them. Having a unified representation for all graph models would be very nice.

[rbharath]

I quite like this suggestion! Ideally our Graph class should be easily interoperable with PyG’s class (perhaps a Graph.to_pytorch_geometric_graph() method)

This will take a little bit of work to do right though. It would be great to field test this idea in Jaxchem and then upstream into DeepChem.