From 468aaec947b2ff461cdc87699c5defd54c133d3f Mon Sep 17 00:00:00 2001
From: Omar Ashour <omar.ashour22@gmail.com>
Date: Mon, 6 May 2024 11:31:39 -0700
Subject: [PATCH] Lazy grid evaluation and get rid of itertools

---
 darkmagic/numerics.py |  9 ++-------
 darkmagic/rate.py     | 18 +++++++++++++++---
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/darkmagic/numerics.py b/darkmagic/numerics.py
index 63f20d0..2b9cb57 100644
--- a/darkmagic/numerics.py
+++ b/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
diff --git a/darkmagic/rate.py b/darkmagic/rate.py
index 9d1790d..faad46e 100644
--- a/darkmagic/rate.py
+++ b/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):