Feature: Add fix symmetry constraint by danielzuegner · Pull Request #438 · TorchSim/torch-sim

danielzuegner · 2026-02-03T17:26:37Z

Summary

This PR adds a FixSymmetry constraint inspired by the way ASE implements it. Disclaimer: I used an LLM to support this PR, and vetted the code quite thoroughly.

Checklist

Before a pull request can be merged, the following items must be checked:

Doc strings have been added in the Google docstring format.
Run ruff on your code.
Tests have been added for any new functionality or bug fixes.

orionarcher

Broadly this seems like a very nice and well-architected addition to the constraints API. Very nice work @danielzuegner. I'd tag in @thomasloux who initially implemented constraints but to me this looks good pretty much as is.

orionarcher · 2026-02-05T13:26:55Z

torch_sim/constraints.py

+    # Merge each group using the constraint's merge method
+    result = []
+    for constraint_type, (constraints, state_indices) in constraints_by_type.items():
+        merged = constraint_type.merge(constraints, state_indices, atom_offsets)


merge is a nice addition to the contraint API

thomasloux

Great work, I'm mainly dubious about the sys_idx_local, sys_idx_global when applying the constraint, but maybe I just misunderstood something.
I haven't really deeped dive the algorithm though, but the code seems overall correct.

pyproject.toml

torch_sim/optimizers/cell_filters.py

torch_sim/optimizers/fire.py

torch_sim/constraints.py

torch_sim/optimizers/cell_filters.py

torch_sim/constraints.py

torch_sim/symmetrize.py

…atcher When InFlightAutoBatcher concatenates states, it calls merge_constraints() which passes state_indices (the position of each source state in the concatenation list) to FixSymmetry.merge(). The old code used these state_indices as offsets: system_indices.append(offset + i). This breaks when a constraint covers multiple systems. E.g. merging constraint_A (5 systems) with constraint_B (3 systems) using state_indices=[0, 1] produces system_indices=[0,1,2,3,4, 1,2,3] — duplicates at 1,2,3 — triggering "Duplicate system indices found in SystemConstraint". Fix: use a cumulative running offset instead of state_indices, so the merged indices are always [0,1,2,...,N-1]. Co-authored-by: Cursor <cursoragent@cursor.com>

janosh · 2026-02-06T02:07:16Z

thanks a lot for this contribution @danielzuegner. super useful! just took it for a spin and encountered a bug: FixSymmetry.merge() produces duplicate system indices when used with InFlightAutoBatcher

ValueError: Duplicate system indices found in `SystemConstraint`.

FixSymmetry.merge() uses state_indices as offsets for system indices. works for single-system constraints but breaks for multi-system ones.

Example: merging constraint_A (5 systems) with constraint_B (3 systems) using state_indices=[0, 1] produces system_indices=[0,1,2,3,4, 1,2,3] — duplicates at 1, 2, 3.

c6a8d0f changes to a cumulative running offset instead of state_indices, so merged indices are always [0,1,...,N-1].

-        for constraint, offset in zip(constraints, state_indices, strict=False):
-            for i in range(len(constraint.rotations)):
-                system_indices.append(offset + i)
+        cumulative_offset = 0
+        for constraint in constraints:
+            for idx in range(len(constraint.rotations)):
+                system_indices.append(cumulative_offset + idx)
+            cumulative_offset += len(constraint.rotations)

i verified on 1, 100, and 10k structures with autobatching — all work now and symmetries go up, i did not encounter a case where symmetry broke

lbfgs_step was setting state.row_vector_cell directly, bypassing set_constrained_cell() which applies cell constraints like FixSymmetry. This meant FixSymmetry's adjust_cell (which symmetrizes the deformation gradient) was never called during LBFGS cell optimization. Fix: compute new_col_vector_cell and use set_constrained_cell() with scale_atoms=True, matching the pattern already used in fire_step. Tested on 1 and 100 structures with MACE+cueq LBFGS + FixSymmetry. Co-authored-by: Cursor <cursoragent@cursor.com>

janosh · 2026-02-06T02:17:14Z

found another issue (cc @abhijeetgangan) lbfgs_step (merged just yesterday #365) was still assigning state.row_vector_cell directly, bypassing set_constrained_cell() so adjust_cell was never called. Changed to set_constrained_cell(new_col_vector_cell, scale_atoms=True) matching fire_step.

@AbstractMethod

- test_merge_multi_system_constraints_no_duplicate_indices: merges a 3-system and 2-system FixSymmetry constraint and asserts sequential system_idx [0,1,2,3,4] — the old state_indices-based offset produced duplicates [0,1,2,1,2] which triggered SystemConstraint validation - test_lbfgs_cell_optimization_preserves_symmetry: runs LBFGS with cell_filter + FixSymmetry on a BCC structure with noisy forces and asserts spacegroup is preserved — the old lbfgs_step bypassed set_constrained_cell so FixSymmetry.adjust_cell was never called - make adjust_stress/adjust_cell default no-ops in Constraint base class instead of @AbstractMethod, removing 4 boilerplate overrides in FixAtoms/FixCom that did nothing - remove __all__ from symmetrize.py - remove unused OptimizationResult TypedDict, structure_with_spacegroup fixture, and unnecessary if-guards on constraint loops in cell_filters Co-authored-by: Cursor <cursoragent@cursor.com>

- replace all spglib usage with moyopy (moyo's Python bindings) for symmetry detection, refinement, and dataset queries - rewrite refine_symmetry using metric tensor polar decomposition + periodic-aware position averaging (no longer depends on spglib's standardized cell pipeline) - FixSymmetry.get_removed_dof returns 0 tensor instead of raising NotImplementedError so temperature calculations work in MD - rename sys_idx_local/sys_idx_global to constraint_idx/sys_idx for clarity (thomasloux review) - add cross-reference to transforms.get_fractional_coordinates in _get_scaled_positions docstring - document n_ops meaning and unbatched nature in symmetrize.py module docstring Co-authored-by: Cursor <cursoragent@cursor.com>

janosh · 2026-02-06T03:35:41Z

i think 73532ac addressed most (or all) of the outstanding PR comments

Reviewer	Comment	Action
CompRhys	Switch from spglib to moyopy	Rewrote `symmetrize.py` to use moyopy's `MoyoDataset` for all symmetry detection. `refine_symmetry` reimplemented using metric tensor polar decomposition + position orbit averaging (no longer depends on spglib internals).
thomasloux	`get_removed_dof` should return 0, not raise	Changed to return `torch.zeros(n_systems)` so temperature calculations work in MD. Updated test accordingly.
thomasloux	`.clone()` needed in cell_filters?	Yes — `adjust_cell` modifies tensors in-place, and `.mT` is a view. Without `.clone()`, `state.reference_cell` could be corrupted. Kept as-is.
thomasloux	Use `transforms.get_fractional_coordinates`	Added cross-reference in docstring. Kept local `_get_scaled_positions` since it's unbatched-only (simpler).
thomasloux	Add docs about n_ops and non-batched nature	Added to module docstring of `symmetrize.py`.
thomasloux	`sys_idx_local` naming confusion	Renamed to `constraint_idx` / `sys_idx` — makes clear that `constraint_idx` indexes into `self.rotations`/`self.symm_maps` while `sys_idx` indexes into the state.

danielzuegner · 2026-02-06T12:24:27Z

Thanks everyone! Let me know if there's anything left for me to do.

- Rewrite symmetrize.py: inline trivial helpers, extract shared moyopy logic into _extract_symmetry_ops and _refine_symmetry_impl - Add chunked fallback in build_symmetry_map for large systems (>200 atoms) to avoid O(n_ops * n_atoms^2) OOM - Add refine_and_prep_symmetry() to combine refinement + symmetry detection in a single moyopy call, eliminating redundant C-library invocation - Lazy-import symmetrize functions inside FixSymmetry methods so importing constraints.py for FixAtoms/FixCom doesn't load the symmetry module - Simplify FixSymmetry.adjust_cell by removing confusing .mT transpose dance - Revert unrelated _system_mapping renames in model files - Add missing tests: position symmetrization vs ASE, refine_symmetry correctness, moyopy to pyproject.toml symmetry extra Co-authored-by: Cursor <cursoragent@cursor.com>

janosh · 2026-02-06T16:50:47Z

while doing some cleanup/refactoring, i just noticed that FixCom.merge([c1, c2]) with multi-system constraints produces [0, 1, 1, 2] which hits the same duplicate validation bug already fixed for FixSymmetry in c6a8d0f. i changed FixCom to use cumulative offsets as well (i.e. same approach as FixSymmetry.merge).

…em states SystemConstraint.merge used the raw state enumeration index as the system offset, which is only correct when each state has exactly 1 system. For multi-system states (e.g. merging two FixCom([0,1]) constraints), this produced duplicate indices [0,1,1,2] instead of [0,1,2,3]. Fix: merge_constraints now computes cumulative system offsets from num_systems_per_state (passed from concatenate_states) and uses those as the state_indices for SystemConstraint.merge, so the offset correctly accounts for multi-system states and gaps from states without the constraint. Co-authored-by: Cursor <cursoragent@cursor.com>

thomasloux · 2026-02-06T16:54:09Z

@janosh Could you add some tests for that? Would be great, this is clearly something we don't want to mess up !

The old merge(constraints, state_indices, atom_offsets) conflated two concerns: shifting indices to global coordinates and concatenating constraints. The state_indices parameter meant different things for different constraint types, and FixSymmetry ignored it entirely. New design: - reindex(atom_offset, system_offset): returns a copy with indices shifted to global coordinates. Each subclass knows what to shift. - merge(constraints): just concatenates already-reindexed constraints. No offset logic needed. - merge_constraints() orchestrates: reindex first, then merge. This eliminates the ambiguous state_indices parameter and gives each method a single clear responsibility. Co-authored-by: Cursor <cursoragent@cursor.com>

janosh · 2026-02-06T16:58:38Z

seems like the current merge() conflates two distinct responsibilities: reindexing (shifting indices to global coordinates) and concatenating (combining multiple constraints into one). The state_indices parameter means different things for different constraint types, and FixSymmetry ignores it entirely to do its own thing. i split this into cea861d two distinct methods reindex() + merge() and line count in torch_sim/constraints.py went down by 45 lines as it does away with the brittle cumulative offsets counting i mentioned. usually a good sign for a method split if line count goes down

New tests covering key fixes: - reindex() preserves rotations/symm_maps while shifting system_idx - SystemConstraint.merge multi-system duplicate-indices regression - concatenate_states end-to-end with FixSymmetry and FixCom - FixSymmetry.__init__ validation, select returning None - adjust_positions/adjust_cell skip when disabled - build_symmetry_map chunked vs vectorized path equivalence - refine_symmetry recovers correct spacegroup after perturbation - position symmetrization matches ASE Simplifications: - Extract p6bar_both_constraints fixture (dedup 4 ASE comparison tests) - Merge 2 test classes into TestFixSymmetryMergeSelectReindex - Combine similar tests via parametrize and multi-assert - Unwrap single-test class, add CPU constant Co-authored-by: Cursor <cursoragent@cursor.com>

@AbstractMethod

- fix reindex() sharing mutable rotations/symm_maps lists between original and copy (shallow copy with list()) - validate adjust_positions/adjust_cell flag consistency in merge() - catch NaN in deformation guard via negated comparison (NaN > x is False) - clamp eigenvalues in _mat_sqrt to prevent NaN from float noise - use torch.linalg.solve instead of inv @ matmul for numerical stability - extract _cumsum_with_zero utility to deduplicate 4 identical patterns - skip stress clone in _get_constrained_stress when no constraints - make Constraint.merge @AbstractMethod (was non-abstract NotImplementedError) - use mask.nonzero() instead of manual list comprehension in select_constraint - add type annotations to NoisyModelWrapper test helper Co-authored-by: Cursor <cursoragent@cursor.com>

janosh

looks in good shape to me now. @thomasloux added bunch more tests. thanks again @danielzuegner 👍

Daniel Zuegner and others added 8 commits February 3, 2026 17:07

initial FixSymmetry

d0ab8e4

style

3d4880e

Merge branch 'main' into fix_symmetry_constraint

13b9795

style

0d08bba

style

29472af

fix

5d7b5b0

bump spglib

22f019d

fix

fb6e234

orionarcher approved these changes Feb 5, 2026

View reviewed changes

thomasloux requested changes Feb 5, 2026

View reviewed changes

abhijeetgangan and others added 2 commits February 5, 2026 16:52

Merge branch 'main' into fix_symmetry_constraint

6556453

janosh_per and others added 2 commits February 6, 2026 02:13

Merge branch 'main' into fix_symmetry_constraint

7bedb3a

janosh and others added 2 commits February 5, 2026 19:05

janosh and others added 2 commits February 6, 2026 09:09

janosh approved these changes Feb 6, 2026

View reviewed changes

janosh enabled auto-merge (squash) February 6, 2026 17:54

janosh requested a review from thomasloux February 6, 2026 17:58

janosh disabled auto-merge February 6, 2026 19:09

janosh merged commit b1d778c into TorchSim:main Feb 6, 2026
67 checks passed

janosh mentioned this pull request Feb 6, 2026

Add cumulative strain guard to FixSymmetry to prevent phase transitions #445

Open

Conversation

danielzuegner commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Checklist

Uh oh!

orionarcher left a comment

Choose a reason for hiding this comment

Uh oh!

orionarcher Feb 5, 2026

Choose a reason for hiding this comment

Uh oh!

thomasloux left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

janosh commented Feb 6, 2026

Uh oh!

janosh commented Feb 6, 2026

Uh oh!

janosh commented Feb 6, 2026

Uh oh!

danielzuegner commented Feb 6, 2026

Uh oh!

janosh commented Feb 6, 2026

Uh oh!

thomasloux commented Feb 6, 2026

Uh oh!

janosh commented Feb 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

janosh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

danielzuegner commented Feb 3, 2026 •

edited

Loading

thomasloux left a comment •

edited

Loading

janosh commented Feb 6, 2026 •

edited

Loading