dpa3 architecture

[littlepeachs]

Contributor (creator of pull-request) checklist

• Add your architecture to the experimental or stable folder. See the
  [docs/src/dev-docs/architecture-life-cycle.rst](Architecture life cycle)
  document for requirements. src/metatrain/experimental/<architecture_name>
• Add default hyperparameter file to
  src/metatrain/experimental/<architecture_name>/default-hypers.yml
• Add your architecture to the CI in .github/workflow/architecture-tests.yml
• Add a new dependencies entry in the optional-dependencies section in the
  pyproject.toml
• Add tests:
  • checking that the code is compatible with TorchScript
  • checking the basic functionality (invariance, fitting, prediction)
  • checking that the checkpoints are properly versionned (see the existing
    test_checkpoint.py in other architectures)
• Add maintainers as codeowners in CODEOWNERS
• Trigger a GPU test by asking a maintainer to comment "cscs-ci run".

Reviewer checklist

New experimental architectures

• Capability to fit at least a single quantity and predict it, verified through CI
  tests.
• Compatibility with JIT compilation using TorchScript <https://pytorch.org/docs/stable/jit.html>_.
• Provision of reasonable default hyperparameters.
• A contact person designated as the maintainer, mentioned in __maintainers__ and the CODEOWNERS file
• All external dependencies must be pip-installable. While not required to be on
  PyPI, a public git repository or another public URL with a repository is acceptable.

New stable architectures

• Provision of regression prediction tests with a small (not exported) checkpoint
  file.
• Comprehensive architecture documentation
• If an architecture has external dependencies, all must be publicly available on
  PyPI.
• Adherence to the standard output infrastructure of metatrain, including
  logging and model save locations.

---

📚 Documentation preview 📚: https://metatrain--725.org.readthedocs.build/en/725/

[ceriottm]

Thanks a lot for this contribution. We'll go through it and try to help as much as possible to merge it quickly and well - and to keep it working as we are still developing metatomic very actively, so things have a tendency to break ^_^

[frostedoyster]

Thank you very much for the contribution @littlepeachs!
A few small comments:

• At the moment, the default dtype for this model is torch.float64. Is this intended? Otherwise, you can set the default to torch.float32 by swapping the order of torch.float64 and torch.float32 in the supported dtypes of the architecture
• The architecture is using the OldCompositionModel. It would be ideal to switch it to CompositionModel right away to avoid more maintenance overhead later. The usage of CompositionModel can be taken exactly from the PET architecture
• We need at least a maintainer for this architecture. You can add maintainers for DPA3 to the CODEOWNERS file (top-level of the repository)

The tests look good to me and we would just need to add some documentation. I'd be happy to add that myself (with some help from you @littlepeachs) after we solve these few issues I raised. Otherwise, I think we're pretty close to being able to merge this. Thanks again!

[littlepeachs]

Thank you very much for the contribution @littlepeachs! A few small comments:

• At the moment, the default dtype for this model is torch.float64. Is this intended? Otherwise, you can set the default to torch.float32 by swapping the order of torch.float64 and torch.float32 in the supported devices of the architecture
• The architecture is using the OldCompositionModel. It would be ideal to switch it to CompositionModel right away to avoid more maintenance overhead later. The usage of CompositionModel can be taken exactly from the PET architecture
• We need at least a maintainer for this architecture. You can add maintainers for DPA3 to the CODEOWNERS file (top-level of the repository)

The tests look good to me and we would just need to add some documentation. I'd be happy to add that myself (with some help from you @littlepeachs) after we solve these few issues I raised. Otherwise, I think we're pretty close to being able to merge this. Thanks again!

Thank you very much for reviewing our code. We will further refine the modifications and submit.

[Luthaf]

Hey @littlepeachs, just checking out with you that everything is going well? Did you had time to look into the review comments?

[littlepeachs]

Hey @littlepeachs, just checking out with you that everything is going well? Did you had time to look into the review comments?

I am sorry to delay the code update. I have been busy with other work so far. I will update the pull request in quick.

[littlepeachs]

I have updated the code based on the reviewer's suggestion. Could some time be spared to review the code again when available? Thanks a lot!

[frostedoyster]

Thanks @littlepeachs for the work, I've finished a few things and now everything looks fine. Can I ask you to fill the section on "Tuning hyperparameters" in docs/src/architectures/dpa3.rst?

@Luthaf there are issues in mtt eval with extension loading. We should try to see what we can do about them

@Luthaf and @PicoCentauri let me know what you think about the PR overall

[PicoCentauri]

Maybe some hints of there or a link to an extrenal docs would be nice.

[PicoCentauri]

torch is required by all arches.

Suggested change

"torch>=2.7",

[Luthaf]

deepmd might require torch >= 2.7 specifically though? I could not find anything in the github repository, but I'm not sure about the wheels.

[PicoCentauri]

Ahh yes, might be possible.

[frostedoyster]

I see that deepmd-kit, as released on PyPI, is compiled with C++11 ABI compatibility and therefore fails to load with any torch version before 2.7. In that case, we need "torch>=2.7" here

[PicoCentauri]

Isn't this something that we decide internally based on the dataset info @frostedoyster ?

[frostedoyster]

Thanks for spotting this. @littlepeachs this one shouldn't be a hyperparameter, but it should be inferred from the atomic numbers that metatrain will give you (see the other architectures), which you can convert to chemical symbols internally with ase.data.chemical_symbols, if I remember correctly