Small fixes for newer versions of numpy and matplotlib

fancompute · Feb 5, 2023 · 20dedea · 20dedea
1 parent 793bdf7
commit 20dedea
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 17 deletions.
diff --git a/legume/gme/gme.py b/legume/gme/gme.py
@@ -668,25 +668,22 @@ def run(self, kpoints: np.ndarray = np.array([[0], [0]]), **kwargs):
 
         # Store an array of the effective permittivity for every layer
         #(including claddings)
-        eps_array = bd.array(list(
-            getattr(layer, layer_eps) for layer in [self.phc.claddings[0]] +
-            self.phc.layers + [self.phc.claddings[1]]),
-                             dtype=bd.float).ravel()
+        eps_array = bd.array([
+            bd.array(getattr(layer, layer_eps), dtype=bd.float).ravel() for layer in
+            [self.phc.claddings[0]] + self.phc.layers + [self.phc.claddings[1]]]).ravel()
         # A separate array where the values are converted from ArrayBox to numpy
         # array, if using the 'autograd' backend.
-        eps_array_val = np.array(list(
-            get_value(getattr(layer, layer_eps))
-            for layer in [self.phc.claddings[0]] + self.phc.layers +
-            [self.phc.claddings[1]]),
-                                 dtype=np.float).ravel()
+        eps_array_val = np.array([
+            np.float64(get_value(getattr(layer, layer_eps))) for layer in
+            [self.phc.claddings[0]] + self.phc.layers + [self.phc.claddings[1]]]).ravel()
 
         # Store an array of thickness of every layer (not including claddings)
         d_array = bd.array(list(layer.d for layer in \
             self.phc.layers), dtype=bd.float).ravel()
         # A separate array where the values are converted from ArrayBox to numpy
         # array, if using the 'autograd' backend.
         d_array_val = np.array(list(get_value(layer.d) for layer in \
-            self.phc.layers), dtype=np.float).ravel()
+            self.phc.layers), dtype=np.float64).ravel()
 
         (self.eps_array_val, self.eps_array, self.d_array_val, self.d_array) = \
                                 (eps_array_val, eps_array, d_array_val, d_array)

diff --git a/legume/gme/slab_modes.py b/legume/gme/slab_modes.py
@@ -324,11 +324,11 @@ def S_T_prod(mats, omegas, g, eps1, eps2, d):
         T11 = np.exp(1j * chis1 * d)
         T22 = np.exp(-1j * chis1 * d)
 
-        T_dot_mats = np.zeros(mats.shape, dtype=np.complex)
+        T_dot_mats = np.zeros(mats.shape, dtype=np.complex128)
         T_dot_mats[0::2, :] = mats[0::2, :] * T11[:, np.newaxis]
         T_dot_mats[1::2, :] = mats[1::2, :] * T22[:, np.newaxis]
 
-        S_dot_T = np.zeros(mats.shape, dtype=np.complex)
+        S_dot_T = np.zeros(mats.shape, dtype=np.complex128)
         S_dot_T[0::2,
                 0] = S11 * T_dot_mats[0::2, 0] + S12 * T_dot_mats[1::2, 0]
         S_dot_T[0::2,
@@ -366,7 +366,7 @@ def S_T_prod(mats, omegas, g, eps1, eps2, d):
         elif pol.lower() == 'tm':
             (S11, S12, S21, S22) = S_TM(eps1, eps2, chis1, chis2)
 
-        mats = np.zeros((2 * N_oms, 2), dtype=np.complex)
+        mats = np.zeros((2 * N_oms, 2), dtype=np.complex128)
         mats[0::2, 0] = S11
         mats[1::2, 0] = S21
         mats[0::2, 1] = S12

diff --git a/legume/viz.py b/legume/viz.py
@@ -703,9 +703,9 @@ def eps_ft(struct,
                                        width_ratios=[1 - cbwidth, cbwidth])
     elif gridspec is not None and fig is not None:
         if cbar == False:
-            gs = mpl.gridspec.GridSpecFromSubplotSpec(len(ars), 1, gridspec)
+            gs = mpl.gridspec.GridSpecFromSubplotSpec(len(all_layers), 1, gridspec)
         else:
-            gs = mpl.gridspec.GridSpecFromSubplotSpec(len(ars), 2, gridspec)
+            gs = mpl.gridspec.GridSpecFromSubplotSpec(len(all_layers), 2, gridspec)
     else:
         raise ValueError(
             "Parameters gridspec and fig should be both specified "
@@ -849,8 +849,8 @@ def field(struct,
 
         pl, o, v = 'xy', 'z', zval
         if periodic == False:
-            kenv = np.exp(1j * grid1 * struct.kpoints[0, kind] +
-                          1j * grid2 * struct.kpoints[1, kind])
+            kenv = np.exp(1j * grid1[None, :] * struct.kpoints[0, kind] +
+                          1j * grid2[:, None] * struct.kpoints[1, kind])
 
     elif x is not None and z is None and y is None:
         if str_type == 'pwe':