add test that is fixed by change

Author: orionarcher

tests/test_autobatching.py

@@ -107,6 +107,19 @@ def test_calculate_scaling_metric(si_sim_state: ts.SimState) -> None:
         calculate_memory_scaler(si_sim_state, "invalid_metric")
 
 
+def test_calculate_scaling_metric_non_periodic(benzene_sim_state: ts.SimState) -> None:
+    """Test calculation of scaling metrics for a non-periodic state."""
+    # Test that calculate passes
+    n_atoms_metric = calculate_memory_scaler(benzene_sim_state, "n_atoms")
+    assert n_atoms_metric == benzene_sim_state.n_atoms
+
+    # Test n_atoms_x_density metric works for non-periodic systems
+    n_atoms_x_density_metric = calculate_memory_scaler(
+        benzene_sim_state, "n_atoms_x_density"
+    )
+    assert n_atoms_x_density_metric > 0
+
+
 def test_split_state(si_double_sim_state: ts.SimState) -> None:
     """Test splitting a batched state into individual states."""
     split_states = si_double_sim_state.split()

torch_sim/autobatching.py

@@ -367,7 +367,7 @@ 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.prod() / 1000
+            volume = bbox.clamp(min=1.0).prod() / 1000  # min 1 Å for planar molecules
         number_density = state.n_atoms / volume.item()
         return state.n_atoms * number_density
     raise ValueError(