make convert_to_8bit_clip_percentile_normalization more robust

wip trajectory

fix trajectory

rm print

Author: curtischong

tests/test_trajectory.py

@@ -33,7 +33,7 @@ def random_state() -> MDState:
         cell=torch.unsqueeze(torch.eye(3) * 10.0, 0),
         atomic_numbers=torch.ones(10, dtype=torch.int32),
         system_idx=torch.zeros(10, dtype=torch.int32),
-        pbc=True,
+        pbc=[True, True, False],
     )
 
 
@@ -473,6 +473,7 @@ def test_get_state(trajectory: TorchSimTrajectory, random_state: MDState) -> Non
         assert state.positions.dtype == expected_dtype
         assert state.cell.dtype == expected_dtype
         assert state.atomic_numbers.dtype == torch.int  # Should always be int
+        assert state.pbc.dtype == torch.bool  # Should always be bool
 
         # Test values (convert to CPU for comparison)
         np.testing.assert_allclose(state.positions, random_state.positions)
@@ -509,7 +510,7 @@ def test_write_ase_trajectory(
         np.testing.assert_allclose(
             atoms.get_atomic_numbers(), random_state.atomic_numbers.numpy()
         )
-        np.testing.assert_array_equal(atoms.pbc, random_state.pbc.numpy()[0])
+        np.testing.assert_array_equal(atoms.pbc, random_state.pbc.numpy())
 
     # Clean up
     ase_traj.close()

torch_sim/trajectory.py

@@ -518,6 +518,20 @@ def write_arrays(
 
         self.flush()
 
+    def write_global_array(self, name: str, array: np.ndarray | torch.Tensor) -> None:
+        """Write a global array to the trajectory file.
+
+        This function is used to write a global array to the trajectory file.
+        """
+        if isinstance(array, torch.Tensor):
+            array = array.cpu().detach().numpy()
+
+        steps = [0]
+        if name not in self.array_registry:
+            self._initialize_array(name, array)
+        self._validate_array(name, array, steps)
+        self._serialize_array(name, array, steps)
+
     def _initialize_array(self, name: str, array: np.ndarray) -> None:
         """Initialize a single array and add it to the registry.
 
@@ -643,15 +657,10 @@ def get_array(
         if name not in self.array_registry:
             raise ValueError(f"Array {name} not found in registry")
 
-        data = self._file.root.data.__getitem__(name).read(
+        return 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,
@@ -788,7 +797,8 @@ def write_state(  # noqa: C901
             self.write_arrays({"atomic_numbers": state[0].atomic_numbers}, 0)
 
         if "pbc" not in self.array_registry:
-            self.write_arrays({"pbc": state[0].pbc}, 0)
+            print("not in array registry")
+            self.write_global_array("pbc", state[0].pbc)
 
         # Write all arrays to file
         self.write_arrays(data, steps)
@@ -830,15 +840,17 @@ def _get_state_arrays(self, frame: int) -> dict[str, np.ndarray]:
         arrays["positions"] = self.get_array("positions", start=frame, stop=frame + 1)[0]
 
         def return_prop(self: Self, prop: str, frame: int) -> np.ndarray:
+            if prop == "pbc":
+                return self.get_array(prop, start=0, stop=3)
             if getattr(self._file.root.data, prop).shape[0] > 1:  # Variable prop
                 start, stop = frame, frame + 1
             else:  # Static prop
                 start, stop = 0, 1
-            return self.get_array(prop, start=start, stop=stop)
+            return self.get_array(prop, start=start, stop=stop)[0]
 
-        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["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["pbc"] = return_prop(self, "pbc", frame)
 
         return arrays