From 704d3597008419586f71874dcc9ce8086295d8ce Mon Sep 17 00:00:00 2001
From: Davis Rens <166565828+davrens101@users.noreply.github.com>
Date: Sun, 8 Nov 2026 23:18:18 -0800
Subject: [PATCH] fix(utils): correctly round cube root in auto_grid_detection

Signed-off-by: Davis Rens <166565828+davrens101@users.noreply.github.com>
---
 src/mattersim/utils/supercell_utils.py |  6 ++--
 tests/utils/__init__.py                |  0
 tests/utils/test_supercell_utils.py    | 39 ++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 tests/utils/__init__.py
 create mode 100644 tests/utils/test_supercell_utils.py

diff --git a/src/mattersim/utils/supercell_utils.py b/src/mattersim/utils/supercell_utils.py
index 7d061fb5..bdc93854 100644
--- a/src/mattersim/utils/supercell_utils.py
+++ b/src/mattersim/utils/supercell_utils.py
@@ -38,7 +38,7 @@ def auto_grid_detection(
         == lattice_vector_lengths[1]
         == lattice_vector_lengths[2]
     ):
-        number_of_replicas = int(np.round(max_atoms / (len(atom))) ** (1 / 3))
+        number_of_replicas = int(np.round((max_atoms / len(atom)) ** (1 / 3)))
         number_of_replicas = max(number_of_replicas, 1)
         max_replication = (
             number_of_replicas,
@@ -62,7 +62,7 @@ def auto_grid_detection(
         # Variation in cell length with in the tolerance we still consider it
         # as an N-N-N replications
         if ratios[0] <= ratio_tolerance and ratios[1] <= ratio_tolerance:
-            number_of_replicas = int(np.round(max_atoms / (len(atom))) ** (1 / 3))
+            number_of_replicas = int(np.round((max_atoms / len(atom)) ** (1 / 3)))
             number_of_replicas = max(number_of_replicas, 1)
             max_replication = (
                 number_of_replicas,
@@ -74,7 +74,7 @@ def auto_grid_detection(
         else:
             # Compute the replica along the most asymmetric direction
             asymmetric_replica = int(
-                (max_atoms / len(atom) / np.prod(ratios)) ** (1 / 3)
+                np.round((max_atoms / len(atom) / np.prod(ratios)) ** (1 / 3))
             )
             asymmetric_replica = max(asymmetric_replica, 1)
 
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/utils/test_supercell_utils.py b/tests/utils/test_supercell_utils.py
new file mode 100644
index 00000000..fa52fa45
--- /dev/null
+++ b/tests/utils/test_supercell_utils.py
@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+import unittest
+
+from ase.build import bulk
+
+from mattersim.utils.supercell_utils import auto_grid_detection
+
+
+class SupercellUtilsTestCase(unittest.TestCase):
+    def test_auto_grid_detection_exact_cubes(self):
+        # 1-atom primitive fcc Cu
+        atoms = bulk("Cu", "fcc", a=3.6)
+        # Perfect cube roots: 64 = 4^3, 216 = 6^3
+        grid_64 = auto_grid_detection(
+            atoms, max_atoms=64, is_santity_check=False, is_verbose=False
+        )
+        self.assertEqual(grid_64, (4, 4, 4))
+
+        grid_216 = auto_grid_detection(
+            atoms, max_atoms=216, is_santity_check=False, is_verbose=False
+        )
+        self.assertEqual(grid_216, (6, 6, 6))
+
+    def test_auto_grid_detection_approximate_lengths(self):
+        # Cell lengths within ratio_tolerance
+        atoms = bulk("Cu", "fcc", a=3.6)
+        atoms.cell[0] *= 1.05
+        grid_64 = auto_grid_detection(
+            atoms,
+            max_atoms=64,
+            ratio_tolerance=1.1,
+            is_santity_check=False,
+            is_verbose=False,
+        )
+        self.assertEqual(grid_64, (4, 4, 4))
+
+
+if __name__ == "__main__":
+    unittest.main()