Skip to content

Implement augmentation routines in metatomic#264

Draft
ppegolo wants to merge 2 commits into
mainfrom
augmentation
Draft

Implement augmentation routines in metatomic#264
ppegolo wants to merge 2 commits into
mainfrom
augmentation

Conversation

@ppegolo

@ppegolo ppegolo commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Augmentation is currently implemented in metatrain, and it depends on spherical for augmentation of spherical targets.
In prepraration for the SymmetrizedModel PR (currently #119, that would depend on metatrain just for the augmentation routines) I re-implemented the relevant parts of spherical (MIT license) here, getting rid of unnecessary stuff and most dependencies, and all the necessary functions that handle augmentation.

Contributor (creator of pull-request) checklist

  • Tests updated (for new features and bugfixes)?
  • Documentation updated (for new features)?
  • Issue referenced (for PRs that solve an issue)?

Reviewer checklist

  • CHANGELOG updated with public API or any other important changes?

📚 Download documentation for this pull-request

@ppegolo ppegolo marked this pull request as draft June 17, 2026 11:13
ppegolo added 2 commits June 17, 2026 13:14
Adds batched O(3) augmentation for System objects and TensorMaps:

- _wigner.py: real Wigner-D matrices from ZYZ Euler angles via the
  Risbo/Trapani-Navaza recurrence, adapted from the spherical package
- _augmentation.py: _apply_augmentations applies per-system rotation
  and improper-rotation matrices to positions, cell vectors, neighbor
  lists, and scalar/Cartesian/spherical TensorMaps; public API exposes
  random_rotations (Haar-uniform SO(3)/O(3) sampling via quaternions)
  and apply_augmentations (validates inputs, derives Wigner-D matrices
  internally from ZYZ decomposition)
- __init__.py: export random_rotations and apply_augmentations
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant