Thomas' review

Author: curtischong

tests/test_trajectory.py

@@ -93,7 +93,7 @@ def test_write_state_single(
     assert trajectory.get_array("positions").shape == (1, 10, 3)
     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, 3)
+    assert trajectory.get_array("pbc").shape == (3,)
 
 
 def test_write_state_multiple(
@@ -106,7 +106,7 @@ def test_write_state_multiple(
     assert trajectory.get_array("positions").shape == (2, 10, 3)
     assert trajectory.get_array("atomic_numbers").shape == (1, 10)
     assert trajectory.get_array("cell").shape == (2, 3, 3)
-    assert trajectory.get_array("pbc").shape == (1, 3)
+    assert trajectory.get_array("pbc").shape == (3,)
 
 
 def test_optional_arrays(trajectory: TorchSimTrajectory, random_state: MDState) -> None:

tests/test_transforms.py

@@ -894,10 +894,20 @@ def test_get_fractional_coordinates_batched() -> None:
             True,
             [[0.2, 0.0, 0.0], [0.0, 0.2, 0.0], [0.0, 0.0, 0.2]],
         ),
+        (
+            [[2.2, 0.0, 0.0], [0.0, 2.2, 0.0], [0.0, 0.0, 2.2]],
+            torch.eye(3, dtype=DTYPE) * 2.0,
+            torch.tensor([True, False, True], dtype=torch.bool),
+            [[0.2, 0.0, 0.0], [0.0, 2.2, 0.0], [0.0, 0.0, 0.2]],
+        ),
     ],
 )
 def test_minimum_image_displacement(
-    *, dr: list[list[float]], cell: torch.Tensor, pbc: bool, expected: list[list[float]]
+    *,
+    dr: list[list[float]],
+    cell: torch.Tensor,
+    pbc: bool | torch.Tensor,
+    expected: list[list[float]],
 ) -> None:
     """Test minimum_image_displacement with various inputs.
 

torch_sim/io.py

@@ -369,11 +369,18 @@ def phonopy_to_state(
         torch.arange(len(phonopy_atoms_list), device=device), atoms_per_system
     )
 
+    """
+    NOTE: PhonopyAtoms does not have pbc attribute for Supercells assume True
+    Verify consistent pbc
+    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")
+    """
+
     return ts.SimState(
         positions=positions,
         masses=masses,
         cell=cell,
-        pbc=True,  # phononpy always assumes periodic boundary conditions https://github.com/phonopy/phonopy/blob/develop/phonopy/structure/atoms.py#L140
+        pbc=True,
         atomic_numbers=atomic_numbers,
         system_idx=system_idx,
     )

torch_sim/neighbors.py

@@ -783,9 +783,9 @@ def torch_nl_linked_cell(
         positions (torch.Tensor [n_atom, 3]):
             A tensor containing the positions of atoms wrapped inside
             their respective unit cells.
-        cell (torch.Tensor [3*num_systems, 3]): Unit cell vectors according to
+        cell (torch.Tensor [3*n_systems, 3]): Unit cell vectors according to
             the row vector convention, i.e. `[[a1, a2, a3], [b1, b2, b3], [c1, c2, c3]]`.
-        pbc (torch.Tensor [num_systems, 3] bool):
+        pbc (torch.Tensor [n_systems, 3] bool):
             A tensor indicating the periodic boundary conditions to apply.
             Partial PBC are not supported yet.
         system_idx (torch.Tensor [n_atom,] torch.long):

torch_sim/trajectory.py

@@ -643,10 +643,15 @@ def get_array(
         if name not in self.array_registry:
             raise ValueError(f"Array {name} not found in registry")
 
-        return self._file.root.data.__getitem__(name).read(
+        data = self._file.root.data.__getitem__(name).read(
             start=start, stop=stop, step=step
         )
 
+        if name == "pbc":
+            return np.squeeze(data, axis=0)
+
+        return data
+
     def get_steps(
         self,
         name: str,
@@ -829,11 +834,11 @@ def return_prop(self: Self, prop: str, frame: int) -> np.ndarray:
                 start, stop = frame, frame + 1
             else:  # Static prop
                 start, stop = 0, 1
-            return self.get_array(prop, start=start, stop=stop)[0]
+            return self.get_array(prop, start=start, stop=stop)
 
-        arrays["cell"] = np.expand_dims(return_prop(self, "cell", frame), axis=0)
-        arrays["atomic_numbers"] = return_prop(self, "atomic_numbers", frame)
-        arrays["masses"] = return_prop(self, "masses", frame)
+        arrays["cell"] = np.expand_dims(return_prop(self, "cell", frame), axis=0)[0]
+        arrays["atomic_numbers"] = return_prop(self, "atomic_numbers", frame)[0]
+        arrays["masses"] = return_prop(self, "masses", frame)[0]
         arrays["pbc"] = return_prop(self, "pbc", frame)
 
         return arrays
@@ -926,14 +931,11 @@ def get_state(
         arrays = self._get_state_arrays(frame)
 
         # Create SimState with required attributes
-        pbc_tensor = torch.tensor(
-            arrays["pbc"], device=device, dtype=torch.bool
-        ).squeeze()
         return SimState(
             positions=torch.tensor(arrays["positions"], device=device, dtype=dtype),
             masses=torch.tensor(arrays.get("masses", None), device=device, dtype=dtype),
             cell=torch.tensor(arrays["cell"], device=device, dtype=dtype),
-            pbc=pbc_tensor,
+            pbc=torch.tensor(arrays["pbc"], device=device, dtype=torch.bool),
             atomic_numbers=torch.tensor(
                 arrays["atomic_numbers"], device=device, dtype=torch.int
             ),

torch_sim/transforms.py

@@ -113,9 +113,7 @@ def inverse_box(box: torch.Tensor) -> torch.Tensor:
 
 @deprecated("Use wrap_positions instead")
 def pbc_wrap_general(
-    positions: torch.Tensor,
-    lattice_vectors: torch.Tensor,
-    pbc: torch.Tensor | bool = True,  # noqa: FBT002
+    positions: torch.Tensor, lattice_vectors: torch.Tensor
 ) -> torch.Tensor:
     """Apply periodic boundary conditions using lattice
         vector transformation method.
@@ -131,16 +129,10 @@ def pbc_wrap_general(
             containing particle positions in real space.
         lattice_vectors (torch.Tensor): Tensor of shape (d, d) containing
             lattice vectors as columns (A matrix in the equations).
-        pbc (torch.Tensor | bool): Boolean tensor of shape (3,) or boolean indicating
-            whether periodic boundary conditions are applied in each dimension.
-            If a boolean is provided, all axes are assumed to have the same periodic
-            boundary conditions.
 
     Returns:
         torch.Tensor: Wrapped positions in real space with same shape as input positions.
     """
-    if isinstance(pbc, bool):
-        pbc = torch.tensor([pbc] * 3)
     # Validate inputs
     if not torch.is_floating_point(positions) or not torch.is_floating_point(
         lattice_vectors
@@ -157,10 +149,7 @@ def pbc_wrap_general(
     frac_coords = positions @ torch.linalg.inv(lattice_vectors).T
 
     # Wrap to reference cell [0,1) using modulo
-    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
+    wrapped_frac = frac_coords % 1.0
 
     # Transform back to real space: r_row_wrapped = wrapped_f_row @ M_row
     return wrapped_frac @ lattice_vectors.T
@@ -223,9 +212,7 @@ def pbc_wrap_batched(
 
     # Wrap to reference cell [0,1) using modulo
     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
+    wrapped_frac[:, pbc] = frac_coords[:, pbc] % 1.0
 
     # Transform back to real space: r = A·f
     # Get the cell for each atom based on its system index
@@ -262,7 +249,7 @@ def minimum_image_displacement(
     dr_frac = torch.einsum("ij,...j->...i", cell_inv, dr)
 
     # Apply minimum image convention
-    dr_frac -= torch.round(dr_frac)
+    dr_frac -= torch.where(pbc, torch.round(dr_frac), torch.zeros_like(dr_frac))
 
     # Convert back to cartesian
     return torch.einsum("ij,...j->...i", cell, dr_frac)