From 240ed7448cc9fb1af7554284d48c0466aad73d7c Mon Sep 17 00:00:00 2001
From: praagnya <praagnya@gmail.com>
Date: Fri, 21 Nov 2025 22:01:12 -0700
Subject: [PATCH] Implement SimState.wrap_positions using pbc_wrap_batched

---
 torch_sim/state.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/torch_sim/state.py b/torch_sim/state.py
index c59c2efd..e5d5c8c5 100644
--- a/torch_sim/state.py
+++ b/torch_sim/state.py
@@ -167,8 +167,11 @@ def wrap_positions(self) -> torch.Tensor:
         """Atomic positions wrapped according to periodic boundary conditions if pbc=True,
         otherwise returns unwrapped positions with shape (n_atoms, 3).
         """
-        # TODO: implement a wrapping method
-        return self.positions
+        if not self.pbc:
+            return self.positions
+        return ts.transforms.pbc_wrap_batched(
+            self.positions, self.cell, self.system_idx, self.pbc
+        )
 
     @property
     def device(self) -> torch.device: