Mixed Pbc Support #320

curtischong · 2025-10-28T22:25:55Z

Summary

Updates torchsim to use the latest version of Vesin which supports mixed PBC neighborlist calculations.

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.

We highly recommended installing the prek hooks running in CI locally to speedup the development process. Simply run pip install prek && prek install to install the hooks which will check your code before each commit.

torch_sim/io.py

curtischong · 2025-11-01T18:16:09Z

torch_sim/models/soft_sphere.py

                cell=cell,
                pbc=self.pbc,
                cutoff=self.cutoff,
-                sorti=False,


yep. this is a spelling mistake

curtischong · 2025-11-01T18:16:59Z

examples/scripts/7_Others/7.3_Batched_neighbor_list.py


-# Fix: Ensure pbc has the correct shape [n_systems, 3]
-pbc_tensor = torch.tensor([[pbc] * 3] * len(atoms_list), dtype=torch.bool)
+pbc_tensor = torch.tensor(pbc).repeat(state.n_systems, 1)


this repeat is needed since torch_nl_linked_cell accepts (num_systems, 3) for the pbc. it doens't just accept a tensor of shape (3,)

curtischong · 2025-11-01T20:58:29Z

tests/test_trajectory.py

    assert trajectory.get_array("atomic_numbers").shape == (1, 10)
    assert trajectory.get_array("cell").shape == (1, 3, 3)
-    assert trajectory.get_array("pbc").shape == (1,)
+    assert trajectory.get_array("pbc").shape == (1, 3)


I'm not 100% sure this is the correct test. should it be shape (1,3) or just shape (3,)

I think pbc should be system_attributes so I would argue for (1, 3)

yeah I saw that we're also probably going to change trajectory reporter soon as well to make it less confusing

I don't think performance-wise making pbc a system attribute will do anything. Since I can see the case for torchsim taking in arbitrary states and manipulating them I am also leaning towards making pbc a system attribute

curtischong · 2025-11-01T21:12:54Z

torch_sim/models/metatomic.py

                    positions=system_positions,
                    types=system_atomic_numbers,
                    cell=system_cell,
-                    pbc=system_pbc,


the simstate.pbc is a torch.Tensor by default now! so we don't need this

curtischong · 2025-11-01T21:22:07Z

torch_sim/neighbors.py

 @torch.jit.script
 def primitive_neighbor_list(  # noqa: C901, PLR0915
    quantities: str,
-    pbc: tuple[bool, bool, bool],


this is a breaking change. but I'm fine with this change because this will make pbc match all of the other params. if this param was bool originally, I'd probably want to make it a torch.Tensor | bool instead.

Another reason why I don't want to do torch.Tensor | bool is because I think that torch.jit.script will be slightly confused by union types (and I'll need to manually specify overrides)

Seems reasonable

torch_sim/transforms.py

curtischong · 2025-11-04T15:09:47Z

there is an annoying bug with torch.jit that I'm trying to solve in a test (our nl function takes too long bc we are timing the compilation and the invocation of the function I believe) but the main API for this is pretty much done. So you can review if you want

wrap only the pbc axis for each atom support init simstate with list of bools better test that uses itertools update vesin version fix pbc check in integrators fix more pytests more fixing more fixes VesinNeighborListTorch is slow fix metatensor for pbc more fixes wip fix pbc trajectory trajectory runs bump metatomic version or vesin will complain fix errors wip fix more trajectory issues make trajectory pass fix pbc in diff_sim fix ase atoms to state conversion for pbc fix neighbors.py assert the pymatgen pbc is valid proper conversion between pbc for atoms, and pymatgen do not pass in pbc to phononpy rm warning and add doc to github make consistent with prev implementation fix io tests lint minor simplification simplify test more simplification changes fix some tests satisfy prek but ruff check errors :/ we'll fix this later more cleanup rename renamove more diffs more changes wip fix rm init in md add type checking and fix pbc type in dataclass def cleanup state loosen test for nl

thomasloux

Some modifications are needed in ts.transforms. Otherwise looks correct to me.
I would also suggest using this PR to port pbc to system_attributes

torch_sim/transforms.py

thomasloux · 2025-11-06T09:11:52Z

torch_sim/transforms.py

    # Wrap to reference cell [0,1) using modulo
-    wrapped_frac = frac_coords % 1.0
+    wrapped_frac = frac_coords.clone()
+    wrapped_frac[:, pbc[0]] = frac_coords[:, pbc[0]] % 1.0


Same remark, but now account for systems_idx:

wrapped_frac[pbc[system_idx]] should work

Since pbc is currently a global attribute I think the current change is okay.

torch_sim/transforms.py

thomasloux · 2025-11-06T09:27:56Z

torch_sim/neighbors.py

            A tensor containing the positions of atoms wrapped inside
            their respective unit cells.
-        cell (torch.Tensor [3*n_structure, 3]): Unit cell vectors according to
+        cell (torch.Tensor [3*num_systems, 3]): Unit cell vectors according to


I would rather use n_systems to be consistent with SimState

thomasloux · 2025-11-06T09:30:16Z

torch_sim/neighbors.py

 @torch.jit.script
 def primitive_neighbor_list(  # noqa: C901, PLR0915
    quantities: str,
-    pbc: tuple[bool, bool, bool],


Seems reasonable

torch_sim/io.py

thomasloux · 2025-11-06T09:38:51Z

tests/test_trajectory.py

    assert trajectory.get_array("atomic_numbers").shape == (1, 10)
    assert trajectory.get_array("cell").shape == (1, 3, 3)
-    assert trajectory.get_array("pbc").shape == (1,)
+    assert trajectory.get_array("pbc").shape == (1, 3)


I think pbc should be system_attributes so I would argue for (1, 3)

curtischong · 2025-11-07T02:29:00Z

@orionarcher or @abhijeetgangan do you recall why pbc is a global attribute? I'm actually fine with either option (it being a global/system attribute) because I think when ppl are doing simulations, they would choose to fix periodicity for all of the systems in their batch.

I do see the case for it being a system attribute. Just wanted your opinion!

abhijeetgangan · 2025-11-07T02:34:53Z

@curtischong In principle you could have different settings for each system in the batch, even mixed PBC for that matter. Only cases it would break would be where the simulations is meant for PBC systems only like phonons or elastic constants etc. We could certainly add error handling case by case but I think this could be future addition and we can keep and issue open for it.

IMO this PR is already a big feature compared to the latter as I see it having more applications than people running different PBC systems in a batch.

curtischong changed the title ~~Pbc~~ Mixed Pbc Oct 28, 2025

curtischong changed the title ~~Mixed Pbc~~ Mixed Pbc Support Oct 28, 2025

curtischong force-pushed the pbc branch from 33b6c2f to 4fb63c6 Compare November 1, 2025 02:19

curtischong commented Nov 1, 2025

View reviewed changes

torch_sim/io.py Show resolved Hide resolved

curtischong commented Nov 1, 2025

View reviewed changes

curtischong force-pushed the pbc branch 2 times, most recently from 94065f9 to 88d8036 Compare November 1, 2025 20:55

curtischong commented Nov 1, 2025

View reviewed changes

torch_sim/transforms.py Outdated Show resolved Hide resolved

curtischong force-pushed the pbc branch 3 times, most recently from 5aba8d7 to fcf1d1a Compare November 3, 2025 00:23

curtischong added the breaking Breaking changes label Nov 3, 2025

curtischong force-pushed the pbc branch from fcf1d1a to f0e7ac5 Compare November 3, 2025 14:40

curtischong added the feature Entirely new features, not improvements to existing ones label Nov 3, 2025

curtischong force-pushed the pbc branch from f0e7ac5 to b6208e0 Compare November 4, 2025 03:22

curtischong force-pushed the pbc branch 2 times, most recently from 19029e5 to 4f77996 Compare November 5, 2025 00:19

curtischong marked this pull request as ready for review November 5, 2025 00:19

curtischong requested review from orionarcher and thomasloux November 5, 2025 00:19

curtischong force-pushed the pbc branch from 4f77996 to 5097cef Compare November 5, 2025 00:19

thomasloux requested changes Nov 6, 2025

View reviewed changes

curtischong added 2 commits November 6, 2025 18:03

revert deleting a comment in io

eb2e5bb

revert pbc_wrap_general changes

278fd08

curtischong added 2 commits November 6, 2025 18:14

make naming consistent

58fc4b9

fix transform for mic

817ba1b

Mixed Pbc Support #320

Are you sure you want to change the base?

Mixed Pbc Support #320

Uh oh!

Conversation

curtischong commented Oct 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Checklist

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

curtischong Nov 7, 2025 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

curtischong commented Nov 4, 2025

Uh oh!

thomasloux left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

curtischong Nov 7, 2025 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

curtischong commented Nov 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

abhijeetgangan commented Nov 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

curtischong commented Oct 28, 2025 •

edited

Loading

curtischong Nov 7, 2025 •

edited

Loading

curtischong Nov 7, 2025 •

edited

Loading

curtischong commented Nov 7, 2025 •

edited

Loading

abhijeetgangan commented Nov 7, 2025 •

edited

Loading