Fix constraint atom_idx not remapped on reordered slice#457
Merged
curtischong merged 4 commits intoTorchSim:mainfrom Feb 18, 2026
Merged
Fix constraint atom_idx not remapped on reordered slice#457curtischong merged 4 commits intoTorchSim:mainfrom
curtischong merged 4 commits intoTorchSim:mainfrom
Conversation
falletta
changed the title
curtischong
reviewed
Feb 18, 2026
| atom_idx so the constraint targets the same physical atoms in the new ordering. | ||
| """ | ||
| state = mixed_double_sim_state | ||
| n_sys0 = (state.system_idx == 0).sum().item() |
Collaborator
There was a problem hiding this comment.
maybe state.n_atoms_per_system()[0].item() is more clear
curtischong
reviewed
Feb 18, 2026
| constraint = sliced.constraints[0] | ||
| assert isinstance(constraint, FixAtoms) | ||
|
|
||
| assert constraint.atom_idx.item() == 1 |
Collaborator
There was a problem hiding this comment.
this line of code is slightly hard to follow because I need to look around and figure out what 1 represents. can we use like a variable to describe it? Maybe we also wouldn't need to use the (state.system_idx == 0).sum().item() above. maybe we can just use constants to say: "original_system1_idx = 1" or "system1_atom_idx=1"
Then we can go something like:
assert constraint.atom_idx.item() == system1_atom_idx
Contributor
Author
There was a problem hiding this comment.
Ok I defined expected_atom_idx = 1 which I then use for the asserts
curtischong
reviewed
Feb 18, 2026
falletta
force-pushed
the
fix-constraints
branch
from
3cfdb86 to
67c6efa
Compare
curtischong
reviewed
Feb 18, 2026
| # Remap constraint atom_idx to account for reordering by atom_indices | ||
| atom_remap = torch.empty(state.n_atoms, dtype=torch.long, device=state.device) | ||
| atom_remap[atom_indices] = torch.arange(len(atom_indices), device=state.device) | ||
| dense_to_reordered = atom_remap[torch.where(atom_mask)[0]] |
Collaborator
There was a problem hiding this comment.
It's hard for me to understand what dense_to_reordered means. Maybe rename to new_atom_idx
curtischong
approved these changes
Feb 18, 2026
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
2 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.
When slicing a batched SimState with a non-sequential system order (e.g.
state[[1, 0]]),select_constraintcorrectly filters constraint atom indices by the kept-atom mask, but the resulting indices were not remapped to reflect the new atom ordering. This causedFixAtoms(and anyAtomConstraint) to target the wrong atoms in the sliced state.This PR adds a remap step in
_filter_attrs_by_indexthat translates constraintatom_idxfrom dense-mask positions to their final reordered positions, so constraints remain physically consistent after any reordering slice.