Add PhACE architecture by frostedoyster · Pull Request #921 · metatensor/metatrain

frostedoyster · 2025-11-18T14:21:18Z

Adds the PhACE architecture (#434 was closed)

Contributor (creator of pull-request) checklist

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.

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

… nanopet

…zation problems

frostedoyster · 2025-12-06T00:02:55Z

cscs-ci run

pfebrer

Sorry for more change requests :)

pfebrer · 2025-12-06T01:23:25Z

src/metatrain/experimental/phace/documentation.py

+class RadialBasisHypers(TypedDict):
+    """In some systems and datasets, enabling long-range Coulomb interactions
+    might be beneficial for the accuracy of the model and/or
+    its physical correctness."""


These are not the long range hypers :)

pfebrer · 2025-12-06T01:26:02Z

src/metatrain/experimental/phace/documentation.py

+    """Metric used to select the best model checkpoint."""
+
+    loss: str | dict[str, LossSpecification] = "mse"
+    """Loss function used for training."""


Can we copy the trainer hypers that are common to PET from the PET trainer, so that they have more extensive documentation?

This would also rename fixed_composition_weights to atomic_baseline, so the trainer and the checkpoint will require changes.

pfebrer · 2025-12-06T01:37:47Z

just dropping in to say that i have had some quite bad experiences with torch_geometric as a dependency and i would +1 avoiding it. a graph dataloader is indeed useful, but it's also very hard to make generic. i think it's not complex enough to warrant generalisation.

Ok, I just think that if there's no reason for the PhACE dataloader to be different from the MACE one (which is torch_geometric, they just copied the necessary code from the torch_geometric package), they should be the same/compatible. I can't think right now of an experiment where both models are trained simultaneously, but it could be something interesting in the future and sharing batches would make things more efficient.

Not that they need to be the same for the PR to be merged, but something that I think would be nice to keep in mind for the future.

jwa7

A couple of comments for now. Looking great though!

jwa7 · 2025-12-05T15:15:19Z

src/metatrain/experimental/phace/documentation.py

+    might be beneficial for the accuracy of the model and/or
+    its physical correctness."""
+
+    max_eigenvalue: float = 25.0


As this needs tuning with respect to the max L of the target, some useful heuristics on how to set this would be useful. For instance, with all other parameters as default, I found I had to set it to max_eigenvalue: 200.0 to reach L=8.

IIRC the laplacian eigenvalues paper uses nmax as the only parameter, and then finds lmax as a consequence. That is far more user-friendly then asking a mysterious eigenvalue. IDK if this applies also with the physical weighting, but it should

Thoughts on this one @frostedoyster ?

I don't think it's a good idea because there are many exact truncations that you can't get (or are ambiguous) if you move to n_max or l_max. I would rather add more documentation to the eigenvalue

jwa7 · 2025-12-08T16:17:46Z

src/metatrain/experimental/phace/modules/cg.py

Yes I second this - I think we need to have fast and more minimal CG code out of featomic. Paolo and I have temporarily included featomic as a dependency for PET-H as we need it for Blocks2Matrix but it seems to be a bit of a pain for torch versioning etc. Maybe we should discuss moving the CG subpackage to metatomic.

src/metatrain/experimental/phace/modules/base_model.py

jwa7 · 2025-12-08T16:24:59Z

src/metatrain/experimental/phace/documentation.py

+    warmup_fraction: float = 0.01
+    """Fraction of training steps for learning rate warmup."""
+
+    gradient_clipping: Optional[float] = None


As discussed - shall we set a default here?

Suggested change

gradient_clipping: Optional[float] = None

gradient_clipping: Optional[float] = 1.0

Yes, I think so. I'm running some more experiments to understand what a good default could be

frostedoyster · 2025-12-08T21:10:47Z

Just to be clear, we can't use any existing implementation of CG products here, because we use a different form of the CG iterations

stale

pfebrer · 2025-12-10T10:42:50Z

src/metatrain/experimental/phace/tests/test_basic.py

+        model(
+            [system],
+            model.outputs,
+        )


For the torchscript tests, I kept them as they were, but because small differences between models (like in your case deleting these attributes, or in MACE compiling with e3nn's function), I wonder if it is better to just test the exporting functionality instead of doing torch.jit.script explicitly, since every model will have its own needed things in model.export()

src/metatrain/experimental/phace/tests/test_basic.py

src/metatrain/experimental/phace/trainer.py

frostedoyster · 2025-12-14T21:32:59Z

cscs-ci run

frostedoyster and others added 30 commits September 4, 2024 11:17

Use covalent radii

770bde6

Revert some changes to the encoder

293c7a6

Remove random factor

94df82e

Merge branch 'main' into nanopet

9eff61f

Speed up a tiny bit

2fd9b8d

Merge branch 'main' into nanopet

6c8f1da

Fair comparison to PET

713cfbb

Calculate displacements independently + format

33e2c70

Fix pyproject.toml

4b8d8c1

Dependencies

3377b2d

Merge branch 'mae-logging' into nanopet

63bcfac

Automatic continuation

f79a36d

Call processing function before new outputs/ directory is created

3c136bf

Make it distributed-proof

c33bcf2

Merge branch 'auto-continue' into nanopet

4b2a877

Only save checkpoints from the main process

129c9b9

Merge branch 'main' into nanopet

d88cfcf

Get rid of nanopet-neighbors

a675bfe

Speed up

72a6502

Proper cutoff parameters

fc33c70

MLP for geometric embedding

a951638

Merge branch 'nanopet' of https://github.com/lab-cosmo/metatrain into…

92d2c80

… nanopet

Fix bug

45c2fc7

Add sample TensorMap to TargetInfo

b8cd9ba

Change TensorMaps inside TargetInfo to float64 to avoid seriali…

60b6b0d

…zation problems

Better documentation

4c614f7

Upgrade to metatensor-torch 0.6.0

172af0b

Upgrade metatensor-learn

0f3bfd7

Update strict NL

0e8e2f7

Fix PBCs

d632dc2

frostedoyster and others added 3 commits December 5, 2025 19:34

Merge branch 'main' into phace-hartmut

98b3289

Fix some tests

898bbd7

Fix some tests

4117ce3

frostedoyster force-pushed the phace-hartmut branch from 99e008c to 182ae3c Compare December 5, 2025 23:59

Fix tests

384210a

frostedoyster force-pushed the phace-hartmut branch from 182ae3c to 384210a Compare December 6, 2025 00:02

pfebrer previously requested changes Dec 6, 2025

View reviewed changes

Fix radial basis description

ad6945b

frostedoyster force-pushed the phace-hartmut branch from 384210a to ad6945b Compare December 6, 2025 07:58

frostedoyster and others added 3 commits December 6, 2025 08:58

Merge branch 'main' into phace-hartmut

398a797

Remove print

bd259b0

Fix gradient request bug

b4eedf0

jwa7 reviewed Dec 8, 2025

View reviewed changes

Lint

4938935

Provide a good error message if target L is too large

bfaeea2

pfebrer reviewed Dec 10, 2025

View reviewed changes

jwa7 and others added 3 commits December 14, 2025 20:04

Merge branch 'main' into phace-hartmut

34b1203

Use atomic baseline

c21e387

Merge branch 'main' into phace-hartmut

c75b088

frostedoyster force-pushed the phace-hartmut branch from b358cef to 740f819 Compare December 14, 2025 21:13

Fix tests

0ded1cc

frostedoyster force-pushed the phace-hartmut branch from 740f819 to 0ded1cc Compare December 14, 2025 21:17

jwa7 reviewed Dec 14, 2025

View reviewed changes

src/metatrain/experimental/phace/trainer.py Outdated Show resolved Hide resolved

Make sure all parameters are used

d7f1d70

Merge branch 'main' into phace-hartmut

4ca5b62

	gradient_clipping: Optional[float] = None
	gradient_clipping: Optional[float] = 1.0

Conversation

frostedoyster commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Contributor (creator of pull-request) checklist

Reviewer checklist

New experimental architectures

Uh oh!

frostedoyster commented Dec 6, 2025

Uh oh!

pfebrer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pfebrer commented Dec 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jwa7 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

frostedoyster commented Dec 8, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

frostedoyster commented Dec 14, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

frostedoyster commented Nov 18, 2025 •

edited

Loading

pfebrer commented Dec 6, 2025 •

edited

Loading