Add PhACE architecture #434
Conversation
There was a problem hiding this comment.
This is lacking a documentation page explaining the architecture, especially since there is AFAIK no associated publication yet.
Also, there is quite a bit of duplicated code with other places in the ecosystem. While not a blocking point for this PR, I would really prefer that we don't have to maintain multiple copies of the same code. I think you should already be able to use things from metatensor-learn; and we should check that the CG code will be able to do everything required here!
| num_element_channels: 64 | ||
| radial_basis: | ||
| mlp: true | ||
| E_max: 50.0 |
There was a problem hiding this comment.
should this be defined in terms of max_radial instead? I don't think users will understand E_max
There was a problem hiding this comment.
The correspondence is not 1:1
There was a problem hiding this comment.
I will change it to max_eigenvalue, but using some max-radial-like argument removes information so I wouldn't do it
There was a problem hiding this comment.
fair enough, but I really doubt people will be able to modify this and understand what it does.
removes information
I'm not clear on what information is lost here?
| ) | ||
|
|
||
|
|
||
| class EquivariantLinear(torch.nn.Module): |
There was a problem hiding this comment.
can't you use the version in metatensor-learn?
There was a problem hiding this comment.
It needs the custom Linear that is defined above, which initializes the weights in a different way from pure torch
| def get_cartesian_vectors( | ||
| positions, cells, species, cell_shifts, pairs, structure_pairs, structure_offsets | ||
| ): |
There was a problem hiding this comment.
Why not use the metatensor NL directly? It already has everything required, except the species_center/species_neighbor metadata
There was a problem hiding this comment.
Because it's not as fast if there are multiple structures in the batch
The metadata needs to be changed a bit since successive structures need to have their indices changed. All the metadata indices here does come from the neighbor lists before being fed into this function
There was a problem hiding this comment.
ok, could you mention this in a comment? So the next person coming does not try to change it back without adding a bunch of benchmarks
| from .layers import EquivariantLinear | ||
|
|
||
|
|
||
| class EquivariantTensorAdd(torch.nn.Module): |
There was a problem hiding this comment.
Can you explain what this is doing? I could not quite understand from reading the code.
Contributor (creator of pull-request) checklist
[docs/src/dev-docs/architecture-life-cycle.rst](Architecture life cycle) document for
requirements.
src/metatrain/experimental/<architecture_name>src/metatrain/experimental/<architecture_name>/default-hypers.yml.ymlfile into github workflows.github/workflow/<architecture_name>.ymloptional-dependenciessection in thepyproject.tomlReviewer checklist
New experimental architectures
tests.
TorchScript <https://pytorch.org/docs/stable/jit.html>_.__maintainers__and theCODEOWNERSfilePyPI, a public git repository or another public URL with a repository is acceptable.
📚 Documentation preview 📚: https://metatrain--434.org.readthedocs.build/en/434/