`unwrap_positions` gives wrong results for triclinic cells

Open Beginner friendly
  • Is something wrong?
#629 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

@lance-lii is already working on this.

Since Oct 31, 2026.

  • #632 by @lance-lii — open

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
85/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Open torch_sim/transforms.py and read unwrap_positions next to pbc_wrap_batched, which treats the columns of cell as lattice vectors. Transpose cells once at the top of unwrap_positions so both the constant-cell and the (T, n_systems, 3, 3) branch use the same convention. Run the reproducer in the issue to confirm the max error drops to ~0, then extend the existing cubic-cell test with a triclinic cell so the regression is covered.

Written by the indexing model from the issue text.

Description

torch_sim.transforms.unwrap_positions seems to assume the cell rows are the lattice vectors. It computes fractional coordinates as positions @ inv(cell) and converts back with dfrac @ cell. But SimState.cell, and pbc_wrap_batched, store the lattice vectors as columns. For orthorhombic cells the two are the same, so the current test, which uses a cubic cell, passes. For a triclinic cell the unwrapped trajectory comes out wrong. Both the constant-cell and the time-dependent (T, n_systems, 3, 3) paths are affected.

Repro (CPU, main @ 35306ec):

import torch
import torch_sim.transforms as tst

dtype = torch.float64
cell = torch.tensor([[6.0, 0.0, 0.0],
                     [1.5, 5.0, 0.0],
                     [1.0, 1.2, 4.5]], dtype=dtype).T.unsqueeze(0)  # columns = a, b, c
system_idx = torch.zeros(2, dtype=torch.long)
start = torch.tensor([[5.5, 4.5, 4.0], [1.0, 1.0, 1.0]], dtype=dtype)
step = torch.tensor([[0.3, 0.25, 0.2], [-0.2, 0.1, -0.15]], dtype=dtype)
true_pos = torch.stack([start + i * step for i in range(20)])
wrapped = torch.stack([tst.pbc_wrap_batched(p, cell, system_idx) for p in true_pos])

out = tst.unwrap_positions(wrapped, cell, system_idx)
print((out - true_pos).abs().max())  # expected ~0
print((tst.unwrap_positions(wrapped, cell.mT, system_idx) - true_pos).abs().max())

Expected: a max error of about 0.

Actual: a max error of 2.5. For example, the last frame of atom 0 should be [11.2, 9.25, 7.8] but comes out as [8.7, 9.55, 10.0]. Passing cell.mT gives about 1e-15, so it's only the convention.

Suggested fix: transpose the cell at the top of unwrap_positions (cells = cells.mT), which covers both branches, and add a triclinic case to the tests. I have a small patch with a regression test and can open a PR if that's welcome.

Disclosure: I used an AI assistant to help find and write up this bug, and I checked the repro myself.

Dominant language
Python
Stars
502
Forks
108
Avg merge
1d 1h
Merged PRs (30d)
12

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from TorchSim/torch-sim

All issues in TorchSim/torch-sim

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.