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.

curtischong · 2025-11-01T18:15:44Z

torch_sim/io.py

-    if not all(all(at.pbc) == all(phonopy_atoms_lst[0].pbc) for at in phonopy_atoms_lst):
-        raise ValueError("All systems must have the same periodic boundary conditions")
-    """
-


I'm unfamiliar with this "supercell" thing. but I deleted it because I think that phononpy always has periodic boundary conditions (https://github.com/phonopy/phonopy/blob/develop/phonopy/structure/atoms.py#L140)? Maybe I just revert this change and we can talk it out in a different PR

Phonopy only makes sense for periodic crystal, although I can imagine that some people may want to that for some kind of 2D crystals with a small z-axis thickness. Sounds reasonable to me to keep that.

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)

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

curtischong · 2025-11-01T21:24:52Z

torch_sim/transforms.py

+    wrapped_frac = frac_coords.clone()
+    wrapped_frac[:, pbc[0]] = frac_coords[:, pbc[0]] % 1.0
+    wrapped_frac[:, pbc[1]] = frac_coords[:, pbc[1]] % 1.0
+    wrapped_frac[:, pbc[2]] = frac_coords[:, pbc[2]] % 1.0


I tried wrapped_frac[:, pbc] = frac_coords[:, pbc] % 1.0 and it doesn't work :'(

Are you sure it doesn't work?
It's a deprecated method, we may not want to adapt it and use it.
If we do want to keep it, pbc needs to be a (d,) shape and so this method doesn't work (I'm actually surprise it would work here. It probably only does because you provide a (n_systems, 3) pbc tensor.

So yes actually the answer is my last observation. But as you don't have a system_idx, you probably need to squeeze to account for (1, d) shape and if n_systems>1 raise an error

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

thomasloux · 2025-11-06T09:09:54Z

torch_sim/transforms.py

+    wrapped_frac = frac_coords.clone()
+    wrapped_frac[:, pbc[0]] = frac_coords[:, pbc[0]] % 1.0
+    wrapped_frac[:, pbc[1]] = frac_coords[:, pbc[1]] % 1.0
+    wrapped_frac[:, pbc[2]] = frac_coords[:, pbc[2]] % 1.0


Are you sure it doesn't work?
It's a deprecated method, we may not want to adapt it and use it.
If we do want to keep it, pbc needs to be a (d,) shape and so this method doesn't work (I'm actually surprise it would work here. It probably only does because you provide a (n_systems, 3) pbc tensor.

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

torch_sim/transforms.py

+    wrapped_frac = frac_coords.clone()
+    wrapped_frac[:, pbc[0]] = frac_coords[:, pbc[0]] % 1.0
+    wrapped_frac[:, pbc[1]] = frac_coords[:, pbc[1]] % 1.0
+    wrapped_frac[:, pbc[2]] = frac_coords[:, pbc[2]] % 1.0


So yes actually the answer is my last observation. But as you don't have a system_idx, you probably need to squeeze to account for (1, d) shape and if n_systems>1 raise an error

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

thomasloux · 2025-11-06T09:20:15Z

torch_sim/transforms.py

+    if cell is None or not pbc.any():
        return dr

    # Convert to fractional coordinates


You need to adapt the rest of the function to account for mixed periodicity.
Chatgpt suggests this, which I think is correct:
dr_frac -= torch.where(pbc, torch.round(dr_frac), torch.zeros_like(dr_frac))

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

thomasloux · 2025-11-06T09:35:21Z

torch_sim/io.py

-    if not all(all(at.pbc) == all(phonopy_atoms_lst[0].pbc) for at in phonopy_atoms_lst):
-        raise ValueError("All systems must have the same periodic boundary conditions")
-    """
-


Phonopy only makes sense for periodic crystal, although I can imagine that some people may want to that for some kind of 2D crystals with a small z-axis thickness. Sounds reasonable to me to keep that.

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 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

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

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

Mixed Pbc Support #320

Are you sure you want to change the base?

Mixed Pbc Support #320

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!

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

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

curtischong commented Nov 4, 2025

Uh oh!

thomasloux left a comment

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

curtischong commented Oct 28, 2025 •

edited

Loading

curtischong Nov 1, 2025 •

edited

Loading