Update to AIMNet2 architecture for radial and vector embedding

[wiederm]

Update to AIMNet2 architecture for radial and vector embedding. Updated the wiki entry.

Summary of changes:

This PR introduces significant improvements to the AimNet2 architecture, specifically in handling radial and vector contributions. The original implementation was limited in capturing angular dependencies and omitted clear aggregation and transformations for certain atomic interactions. These changes aim to enhance model expressiveness and bring it closer to the (somewhat vague) description in the original paper.

Key Modifications:

• Refinement of Radial and Vector Features: gs and gv tensors are introduced to represent radial and vector (angular) symmetry functions, respectively. The radial basis (gs) captures pairwise distances, while gv encodes directional vectors for atom pairs.
• Vector Contribution Computation: using an Einstein summation operation (torch.einsum), gv and an agh transformation matrix capture the vector (angular) contributions. This setup adds expressiveness by handling directional information, aligning the model closer with angular-dependent symmetry functions. The operation structure: torch.einsum("pa, pdg, agh -> phd", a_j, gv, agh) ensures the features align over vector components.
• Enhanced Broadcasting and Aggregation: to match dimensions, radial and vector tensors undergo explicit broadcasting (via .unsqueeze()). This step avoids mismatches and clarifies the contribution structure.
• Aggregated Contributions: output radial and vector contributions for each pair are aggregated per atom using index_add_ to collect contributions correctly.
• Scaling with f Factor: a scaling factor f is applied to delta_q (charge updates), which is introduced to improve control over charge accumulation during message-passing steps. This approach allows initializing and incrementally updating charges in a more controlled manner.
• Unified Message Module: each message-passing module now utilizes an MLP with outputs split into delta_q, f, and delta_a, making the function of each component explicit. This organization supports independent updates for atomic embeddings (delta_a) and partial charges (delta_q), providing finer control over per-atom features and charge states.
• Structural Adjustments for Message Passing: the first interaction module updates only atomic embeddings, while subsequent modules also modify partial charges. The message composition in later modules combines features from both embeddings and charges, producing a richer atomic feature set.

Paper Vague on Equations
The original paper is vague regarding details, especially for the equations governing vector contributions and aggregation. This implementation assumes certain interpretations based on the partial information given, and are not easily to decode. This is therefore not a literal implementation of AimNet2, but only a interpretation.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 98.24561% with 1 line in your changes missing coverage. Please review.

Project coverage is 85.52%. Comparing base (65d2b43) to head (9115792).
Report is 9 commits behind head on main.

Additional details and impacted files