search

choderalab / modelforge

Infrastructure to implement and train NNPs
https://modelforge.readthedocs.io/en/latest/
MIT License
9 stars 4 forks source link

Batching and pair list logic not applicable to different system sizes? #23

Closed yuanqing-wang closed 8 months ago

yuanqing-wang commented 9 months ago

There appear to be multiple design choices in the implementation of the base model for the machine learning potentials that would make batching difficult. @ArnNag can fill in more details.

ArnNag commented 9 months ago

I talked to @yuanqing-wang about this, and we think it is better to move batching operations outside the model and perform message passing operations on tensors that do not have a batch_size axis. The forward method of models that inherit from BaseNNP would take in a pairlist directly as opposed to a padded tensor with a batch_size axis. Node feature shapes would be (n_nodes_batched, n_node_features), and edge feature shapes would be (n_pairs_batched, n_edge_features).

In addition, the pairlist should keep track of which molecule each atom belongs to. This is necessary to read out the energy of each molecule. This can be an array of shape (n_molecules, ) where each element is the number of atoms in the molecule (along the lines of n_node here).

wiederm commented 8 months ago

I think this is now addressed in PR #27 . Issue can be closed as soon as this is merged.