Modularity features: (inter)dependence of feats, edges, coors

[mireklzicar]

I am curious about whether it is possible to make the model more modular. I haven't gone through the code and don't know what is possible to do without major changes in the architecture, but I wanted to start the discussion.

Current status:

• Training on atoms and coors, edges are voluntary
• Outputting atoms and coors (no edges)

Proposed modularity features:

• Possibility to output edges
• Possibility to train on any pair of inputs: a) atoms and coors (already implemented), b) atoms and edges c) coors and edges
• Possibility to train on any single input

Reasons why & Use cases:

• Outputting edges might be super-useful for predicting bonds, protein-ligand binding sites etc.
• Any pair: atoms-edges makes sense when we don't have 3D structures and we are training on "2D" molecular or protein graphs with only atoms and connectivity info. Might be used for de novo prediction of coordinates too.
• Single input: a) On atoms that would collapse into simple transformer b) on coordinates it could be used for training on 3D point-clouds (and benchmarked with PointNet etc.) c) on edges it could be used for training on graphs.