Fix FixSymmetry constraint system_idx remapping on reordered state slicing#509
Merged
CompRhys merged 7 commits intoTorchSim:mainfrom Mar 17, 2026
Merged
Conversation
When _filter_attrs_by_index selects systems in non-ascending order (e.g. state[[3, 1, 4]]), the constraint's system_idx was remapped using only a boolean mask, which always produces indices in ascending order. This caused FixSymmetry's per-system data (rotations, symm_maps, reference_cells) to be applied to the wrong systems. BinningAutoBatcher triggers this because to_constant_volume_bins inserts items in descending-volume order, producing non-sorted index bins. Add system_idx remapping for constraints in _filter_attrs_by_index, analogous to the existing atom_idx remapping. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Member
Contributor
|
I tested the code on my use case, which required the fix in #457, and everything looks good to me. |
Member
|
Offline @falletta confirmed this fix is consistent with his fix for his workflow. Thanks for the fix! |
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.
Summary
Problem
When slicing a
SimStatewith out-of-order system indices (e.g.state[[3, 1, 4]]),FixSymmetry's per-system data (rotations,symm_maps,reference_cells) gets applied to the wrong systems.This is triggered during relaxation with
BinningAutoBatcher, which produces non-sorted system indices from its descending-volume bin-packing algorithm.Root cause
_filter_attrs_by_indexinstate.pyconvertssystem_indicesto a boolean mask before passing toselect_constraint. Boolean masks lose ordering information —torch.where(mask)always returns indices in ascending order. While constraintatom_idxremapping already existed to correct for this, the analogoussystem_idxremapping was missing.Fix
Build both
atom_remapandsystem_remaptables upfront, then remap bothatom_idxandsystem_idxon each constraint in a single pass afterselect_constraint. This also consolidates the previously separate atom/system remapping loops.Test
Added two tests for out-of-order slicing:
test_fix_symmetry_system_idx_remapped_on_reordered_slice— creates a 2-system state with distinctFixSymmetryrotations per system, slices withstate[[1, 0]], and verifies each output system is paired with the correct rotation and reference cell.test_fix_com_system_idx_remapped_on_reordered_slice— constrains only system 0 withFixCom, slices withstate[[1, 0]], and verifiessystem_idxis remapped to[1].Disclaimer: This is human-checked and verified LLM code :)
Checklist
Before a pull request can be merged, the following items must be checked: