respond to rhys and curtis PRs

Author: orionarcher

torch_sim/autobatching.py

@@ -367,7 +367,11 @@ def calculate_memory_scaler(
             volume = torch.abs(torch.linalg.det(state.cell[0])) / 1000
         else:
             bbox = state.positions.max(dim=0).values - state.positions.min(dim=0).values
-            volume = bbox.clamp(min=2.0).prod() / 1000  # min 1 Å for planar molecules
+            # add 2 A in non-periodic directions to account for 2D systems and slabs
+            for i, periodic in enumerate(state.pbc):
+                if not periodic:
+                    bbox[i] += 2.0
+            volume = bbox.prod() / 1000  # convert A^3 to nm^3
         number_density = state.n_atoms / volume.item()
         return state.n_atoms * number_density
     raise ValueError(