Lazy grid evaluation and get rid of itertools

darkmagic/numerics.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import itertools
 from copy import deepcopy
 from typing import TYPE_CHECKING, Tuple
 
@@ -169,12 +168,8 @@ def _get_G_vectors(
             A tuple containing the Cartesian and fractional coordinates of the G vectors.
         """
         # Generate the 8 closest G-vectors to each q-point
-        G_frac = np.array(
-            [
-                list(itertools.product(*zip(np.floor(q), np.ceil(q))))
-                for q in self.q_frac
-            ]
-        )
+        offsets = np.indices((2,) * 3).reshape(3, -1).T
+        G_frac = np.floor(self.q_frac)[:, None, :] + offsets
         # TODO: double check that right multiplication works well
         G_cart = np.matmul(G_frac, recip_frac_to_cart)
         # Compute the distances between the q-point and each of the 8 G-vectors

darkmagic/rate.py

@@ -29,12 +29,24 @@ def __init__(
         numerics: Numerics,
     ):
         self.m_chi = m_chi
-        self.v_e = v_e  # self.compute_ve(time) if time is not None else v_e
+        self.v_e = v_e
         self.material = material
         self.model = model
         self.numerics = numerics
-        self.grid = numerics.get_grid(m_chi, self.v_e, material)
-        self.dwf_grid = numerics.get_DWF_grid(material)
+        self._grid = None
+        self._dwf_grid = None
+
+    @property
+    def grid(self):
+        if self._grid is None:
+            self._grid = self.numerics.get_grid(self.m_chi, self.v_e, self.material)
+        return self._grid
+
+    @property
+    def dwf_grid(self):
+        if self._dwf_grid is None:
+            self._dwf_grid = self.numerics.get_DWF_grid(self.material)
+        return self._dwf_grid
 
     @abstractmethod
     def calculate_sigma_q_nu(self, omegas, epsilons):