Add an easier way to integrate an MLIP-only architecture #876
Conversation
|
@copilot I've added the base class for the model. Now we would need the same, but for the trainer. The only thing that the trainer needs to know from the derived class is whether to do data augmentation or not. To write the base trainer, you can take inspiration from those of SOAP-BPNN (no rotational augmentation) and PET (does rotational augmentation) |
|
@frostedoyster I've opened a new pull request, #877, to work on those changes. Once the pull request is ready, I'll request review from you. |
|
@copilot Would you be able to
|
|
@frostedoyster I've opened a new pull request, #880, to work on those changes. Once the pull request is ready, I'll request review from you. |
|
@copilot can you fix the tests and docs? |
|
@frostedoyster I've opened a new pull request, #884, to work on those changes. Once the pull request is ready, I'll request review from you. |
|
@copilot Could you address the review comments, please? |
|
@frostedoyster I've opened a new pull request, #885, to work on those changes. Once the pull request is ready, I'll request review from you. |
|
@copilot Can you implement the review comment? |
|
@frostedoyster I've opened a new pull request, #887, to work on those changes. Once the pull request is ready, I'll request review from you. |
frostedoyster
requested review from
E-Rum,
Luthaf,
johannes-spies and
sirmarcel
| Example: Creating an MLIP-only architecture | ||
| -------------------------------------------- | ||
|
|
||
| To demonstrate how easy it is to add a new MLIP-only architecture using the base |
There was a problem hiding this comment.
| To demonstrate how easy it is to add a new MLIP-only architecture using the base | |
| To demonstrate how to add a new MLIP-only architecture using the base |
| __authors__ = [ | ||
| ("GitHub Copilot <[email protected]>", "@copilot"), | ||
| ] | ||
|
|
||
| __maintainers__ = [ | ||
| ("GitHub Copilot <[email protected]>", "@copilot"), | ||
| ] |
There was a problem hiding this comment.
I'm not sure we want example code like this to live inside src/metatrain. It is not something we want to distribute to everyone
| edge_vectors: torch.Tensor, | ||
| species: torch.Tensor, | ||
| centers: torch.Tensor, | ||
| neighbors: torch.Tensor, | ||
| system_indices: torch.Tensor, |
There was a problem hiding this comment.
I don't love that the API here is so different from the metatomic interface. Ideally this class should help people move toward a deeper integration if they need it down the line.
Maybe something like edge_vector => pair_vectors/pair_distances, centers, neighbors => pair_indices (not sure about this one), species => atomic_types.
Also, should the model have access to the raw positions/cell? Or is this limited to GNN style models?
| edge_vectors: torch.Tensor, | |
| species: torch.Tensor, | |
| centers: torch.Tensor, | |
| neighbors: torch.Tensor, | |
| system_indices: torch.Tensor, | |
| system_indices: torch.Tensor, | |
| atomic_types: torch.Tensor, | |
| centers: torch.Tensor, | |
| neighbors: torch.Tensor, | |
| pair_vectors: torch.Tensor, |
or
| edge_vectors: torch.Tensor, | |
| species: torch.Tensor, | |
| centers: torch.Tensor, | |
| neighbors: torch.Tensor, | |
| system_indices: torch.Tensor, | |
| system_indices: torch.Tensor, | |
| atomic_types: torch.Tensor, | |
| pair_indices: torch.Tensor, | |
| pair_vectors: torch.Tensor, |
| return AtomisticModel(self.eval(), metadata, capabilities) | ||
|
|
||
| @classmethod | ||
| def upgrade_checkpoint(cls, checkpoint: Dict[str, Any]) -> Dict[str, Any]: |
There was a problem hiding this comment.
I think this should be possible to override from the child class, since MLIPModel has no idea about what the checkpoint will look like.
Addresses #571
Contributor (creator of pull-request) checklist
Maintainer/Reviewer checklist
📚 Documentation preview 📚: https://metatrain--876.org.readthedocs.build/en/876/