Skip to content

Review and Optimize SymmetrizedModel on top of #119 - #289

Merged
ppegolo merged 4 commits into
metatensor:token_rotationsfrom
MichelangeloDomina:review/symmetrized-model-pr119-final-series
Jul 15, 2026
Merged

Review and Optimize SymmetrizedModel on top of #119#289
ppegolo merged 4 commits into
metatensor:token_rotationsfrom
MichelangeloDomina:review/symmetrized-model-pr119-final-series

Conversation

@MichelangeloDomina

@MichelangeloDomina MichelangeloDomina commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

This is a stacked follow-up to #119. It intentionally targets the
token_rotations branch at 31cc945, so this pull request contains only the
four follow-up commits rather than repeating the complete parent PR.

Scope

The changes are limited to:

  • metatomic.torch.SymmetrizedModel;
  • the O(3) transformation utilities directly required by that wrapper; and
  • the existing ASE SymmetrizedCalculator, including the requested-input and
    space-group paths needed to construct its symmetry orbit.

No metatrain model, training procedure, loss, dataset, optimizer, or base-model
architecture is modified.

Motivation

The implementation in #119 establishes the main feature, but several
mathematical, numerical, API, and scalability contracts required additional
work before the feature could be considered internally consistent and suitable
for larger grids and datasets.

The follow-up therefore:

  • applies one proper/improper O(3) convention across every public path;
  • makes multi-System routing and spherical metadata validation explicit;
  • uses pole-safe Wigner construction;
  • resolves the Wigner-product space actually integrated by the finite grid;
  • accumulates centered moments to avoid catastrophic cancellation;
  • defines explicit derivative and three-dimensional stress behavior;
  • preserves requested ASE inputs and complete closed space-group actions; and
  • removes repeated representation work and full-grid prediction retention.

Intentional behavioral corrections

Some results intentionally differ from the original #119 implementation:

  • generated spherical outputs use canonical o3_lambda and o3_sigma keys;
  • centered accumulation replaces the unstable difference of large moments;
  • attached TensorMap gradients are rejected rather than silently discarded;
  • conservative forces and eligible fully periodic stress are derived through
    autograd when requested;
  • partial or nonperiodic Systems omit undefined volume-based 3D stress;
  • singular or non-finite fully periodic cells raise before stress division;
  • detached energies warn and produce exact-zero derivatives as constant
    energies;
  • unsupported MPS execution is rejected before evaluation without changing
    CPU or CUDA behavior; and
  • ASE space-group projection retains complete actions, including
    translation-distinct atom permutations.

These cases are tested against analytical, structural, or API targets.

Performance and scalability

The optimization are limited to:

  • direct and lazy real Wigner construction;
  • Wigner work sized from the actual returned rank;
  • a bounded persistent Wigner cache with an exact uncached fallback;
  • batched core-System transformations;
  • canonical storage of forward rotations without a duplicate inverse grid;
  • streamed ASE moment reduction whose retained state is independent of the
    number of quadrature operations; and
  • indexed periodic-site candidate matching with Cartesian minimum-image
    verification and an exact bounded fallback.

Controlled benchmarks showed material improvements on the affected paths.

Change composition

Relative to #119, this series changes 20 files:

  • production-source paths: +1,564 net lines, including reviewer-facing docstrings;
  • tests: +2,471 net lines;
  • public documentation, changelogs, and configuration: +436 net lines.

The total increase relative to #119 is 4,471 net lines. Tests account for
55.3% of this increase. Compared with the previous draft, the runtime
implementation is unchanged while repetitive tests and documentation were
consolidated, reducing the series by 221 net lines.

Validation

Validation on the exact four-commit tip:

  • complete local Torch suite: 372 passed, 7 skipped;
  • complete local ASE suite: 140 passed, 2 skipped;
  • CUDA O(3)/SymmetrizedModel suite: 267 passed, 2 MPS-only skips;
  • CUDA ASE symmetrization suite: 94 passed;
  • Torch coverage before and after compression: unchanged at 96% statements
    and 94% including branches;
  • ASE symmetrization coverage: unchanged at 90% including branches;
  • Ruff lint and formatting: passed;
  • git diff --check: passed;
  • Sphinx documentation build with warnings treated as errors: passed.

The three commits were also tested independently to preserve bisectability.

Review structure

The commits are intentionally divided into four review units:

  1. O(3) transformation correctness and Wigner construction;

  2. SymmetrizedModel numerical, derivative, schema, cache, and performance contracts;

  3. ASE requested-state preservation, streaming reduction, and space-group projection;

  4. reviewer-facing mathematical/API documentation and the direct Wigner-cache
    mutation regression.

    I have refreshed the draft series after a reviewability pass. The runtime implementation is unchanged from the previous draft. Repetitive tests were consolidated, non-obvious mathematical and API contracts were documented, and a direct quadrature-grid mutation regression was restored.

@ppegolo

ppegolo commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

@Luthaf don't worry about this one, we'll get to a decent final state and then request review

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.
@MichelangeloDomina
MichelangeloDomina force-pushed the review/symmetrized-model-pr119-final-series branch from 79d4581 to 9f4932c Compare July 15, 2026 15:16
@ppegolo
ppegolo marked this pull request as ready for review July 15, 2026 15:36
@ppegolo
ppegolo merged commit 4ca60ac into metatensor:token_rotations Jul 15, 2026
59 checks passed
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.

2 participants