Add SymmetrizedModel - #119
Draft
ppegolo wants to merge 15 commits into
Draft
Conversation
ppegolo
force-pushed
the
token_rotations
branch
from
November 21, 2025 13:42
c8f2b85 to
3147311
Compare
6 tasks
4 tasks
ppegolo
force-pushed
the
token_rotations
branch
5 times, most recently
from
July 9, 2026 15:08
715c0a8 to
0008113
Compare
Contributor
Author
|
There are a few remaining issues on if and how should this model be integrated into the SymmetrizedCalculator. In particular there are some inconsistencies on how the grid lambda is chosen. The choice made in the model are exact, while I think the calculator at this stage is only exact when you compute the mean (not the variance) |
…n in SymmetrizedModel
…nd SymmetrizedCalculator
…tion, naming and packaging
…lambda against the declared target
…rojection records, trim duplication
Make routing, improper parity, pole recovery, dtype/device handling, and all ten documented component axes explicit and tested. Construct real Wigner matrices directly and lazily, batch-validate internally generated random transformations, and enforce canonical Cartesian and spherical bases even for empty blocks. These changes preserve already-correct cases while removing repeated packing and per-matrix synchronization.
Preserve complete output schemas, reject gradients that would otherwise be lost, support raw and exported models plus selected atoms, and define conservative derivatives consistently with metatrain. Detached energies warn and return zeros; 3D stress is limited to finite-volume full PBC. Use reference-centered float64 moments, actual returned ranks, batched geometry transforms, trusted quadrature operations, and bounded reusable Wigner storage. Build CUDA Wigner batches on CPU and transfer completed stacks to avoid per-rotation synchronization. The public guide states the finite-grid, unit, derivative, memory, and convergence contracts.
Rotate every supported requested polar-vector input and reject unsupported axial-vector inversion. Propagate requested TensorMaps, watch requested ASE state for cache invalidation, and retain only closed space-group actions that preserve model inputs. Stream centered rotational moments by batch, skip scalar no-op group discovery, preserve caller-selected periodic images and translation-distinct permutations, and accelerate large periodic site matching with conservative indexed candidates plus exact bounded fallback. Stress follows the same full-PBC finite-volume contract as SymmetrizedModel.
Luthaf
reviewed
Jul 17, 2026
| ): | ||
| super().__init__() | ||
| self.base_model = base_model | ||
| self._base_is_atomistic = is_atomistic_model(base_model) |
Member
There was a problem hiding this comment.
you can get the inner module of an AtomisticModel and then simplify this to always use the inner module
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add
SymmetrizedModelwrapper to compute O(3) average and metrics for metatomic models.Contributor (creator of pull-request) checklist
[ ] Issue referenced (for PRs that solve an issue)?Reviewer checklist
📚 Download documentation for this pull-request
⚙️ Download Python wheels for this pull-request (you can install these with pip)