From 54a8954298d1d118be85b99ed3b58a1c662b4c75 Mon Sep 17 00:00:00 2001
From: Omar Ashour <omar.ashour22@gmail.com>
Date: Mon, 6 May 2024 13:48:56 -0700
Subject: [PATCH] Fix failing tests

---
 darkmagic/calculator.py | 2 +-
 darkmagic/numerics.py   | 5 ++++-
 tests/test_material.py  | 2 +-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/darkmagic/calculator.py b/darkmagic/calculator.py
index 5bfbd75..efbc08d 100644
--- a/darkmagic/calculator.py
+++ b/darkmagic/calculator.py
@@ -252,7 +252,7 @@ def _evaluate_serial(
         for task in np.indices((nm, nv)).T.reshape(-1, 2):
             im, iv = task[0], task[1]
             print(
-                f"Rate calculation: {im * nv + iv + 1}/{(nv*nm)}.",
+                f"Rate calculation: {im + iv*nm + 1}/{(nv*nm)}.",
             )
             (
                 diff_rate[iv, im],
diff --git a/darkmagic/numerics.py b/darkmagic/numerics.py
index c6ffffe..7d6e92b 100644
--- a/darkmagic/numerics.py
+++ b/darkmagic/numerics.py
@@ -24,7 +24,7 @@ class MonkhorstPackGrid:
 
     def __init__(
         self,
-        N_grid: np.ndarray,
+        N_grid: ArrayLike,
         shift: bool = True,
     ):
         """
@@ -33,10 +33,13 @@ def __init__(
         N_grid (ArrayLike): The number of grid points along each reciprocal lattice vector.
         shift (bool): Whether to shift the grid. Defaults to a shifted grid to avoid singularities in DWF.
         """
+        N_grid = np.array(N_grid)
         s = np.array([1, 1, 1]) / 2 if shift else np.array([0, 0, 0])
+        # Generate grid and apply shifts
         grid_indices = np.meshgrid(*[np.arange(n) for n in N_grid], indexing="ij")
         k_list = np.stack(grid_indices, axis=-1).reshape(-1, 3) / N_grid
         k_list += ((N_grid + 1) % 2 - N_grid // 2 + s) / N_grid
+        # Sort the k-points by their coordinates
         self.k_frac = np.array(sorted(k_list, key=tuple), dtype=np.float64)
         self.weights = np.ones_like(self.k_frac[:, 0], dtype=np.int64)
 
diff --git a/tests/test_material.py b/tests/test_material.py
index 34f3b60..39c71e0 100644
--- a/tests/test_material.py
+++ b/tests/test_material.py
@@ -313,7 +313,7 @@ def test_phonon_material(
     assert material.m_atoms[0] == pytest.approx(m_atom1 * const.amu_to_eV)
     assert np.all(material.epsilon == eps)
     assert np.all(material.born[0] == bec1)
-    grid = MonkhorstPackGrid([1, 1, 1], material, shift=False)
+    grid = MonkhorstPackGrid([1, 1, 1], shift=False)
     freq, eigvec = material.get_eig(grid, with_eigenvectors=True)
     assert freq.shape == (1, n_atoms * 3)
     assert eigvec.shape == (1, n_atoms * 3, n_atoms, 3)