typing updates

examples/fdfd.py

@@ -156,7 +156,7 @@ def test1(solver=generic_solver):
     # grid.draw_cuboid(pmcg, center=[700, 0, 0], dimensions=[80, 1e8, 1e8], eps=1)
     # grid.visualize_isosurface(pmcg)
 
-    def pcolor(v):
+    def pcolor(v) -> None:
         vmax = numpy.max(numpy.abs(v))
         pyplot.pcolor(v, cmap='seismic', vmin=-vmax, vmax=vmax)
         pyplot.axis('equal')

meanas/fdfd/bloch.py

@@ -82,9 +82,10 @@ This module contains functions for generating and solving the
 
 from typing import Tuple, Callable, Any, List, Optional, cast
 import logging
-import numpy                            # type: ignore
+import numpy
-from numpy import pi, real, trace       # type: ignore
+from numpy import pi, real, trace
-from numpy.fft import fftfreq           # type: ignore
+from numpy.fft import fftfreq
+from numpy.typing import NDArray, ArrayLike
 import scipy                            # type: ignore
 import scipy.optimize                   # type: ignore
 from scipy.linalg import norm           # type: ignore
@@ -109,21 +110,22 @@ try:
         'planner_effort': 'FFTW_EXHAUSTIVE',
         }
 
-    def fftn(*args: Any, **kwargs: Any) -> numpy.ndarray:
+    def fftn(*args: Any, **kwargs: Any) -> NDArray[numpy.float64]:
         return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args)
 
-    def ifftn(*args: Any, **kwargs: Any) -> numpy.ndarray:
+    def ifftn(*args: Any, **kwargs: Any) -> NDArray[numpy.float64]:
         return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)
 
 except ImportError:
-    from numpy.fft import fftn, ifftn       # type: ignore
+    from numpy.fft import fftn, ifftn
     logger.info('Using numpy fft')
 
 
-def generate_kmn(k0: numpy.ndarray,
+def generate_kmn(
-                 G_matrix: numpy.ndarray,
+        k0: ArrayLike,
-                 shape: numpy.ndarray
+        G_matrix: ArrayLike,
-                 ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+        shape: ArrayLike,
+        ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
     """
     Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.
 
@@ -162,11 +164,12 @@ def generate_kmn(k0: numpy.ndarray,
     return k_mag, m, n
 
 
-def maxwell_operator(k0: numpy.ndarray,
+def maxwell_operator(
-                     G_matrix: numpy.ndarray,
+        k0: ArrayLike,
+        G_matrix: ArrayLike,
         epsilon: fdfield_t,
-                     mu: fdfield_t = None
+        mu: Optional[fdfield_t] = None
-                     ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+        ) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]:
     """
     Generate the Maxwell operator
 
@@ -199,7 +202,7 @@ def maxwell_operator(k0: numpy.ndarray,
     if mu is not None:
         mu = numpy.stack(mu, 3)
 
-    def operator(h: numpy.ndarray) -> numpy.ndarray:
+    def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
         """
         Maxwell operator for Bloch eigenmode simulation.
 
@@ -244,10 +247,11 @@ def maxwell_operator(k0: numpy.ndarray,
     return operator
 
 
-def hmn_2_exyz(k0: numpy.ndarray,
+def hmn_2_exyz(
-               G_matrix: numpy.ndarray,
+        k0: ArrayLike,
+        G_matrix: ArrayLike,
         epsilon: fdfield_t,
-               ) -> Callable[[numpy.ndarray], fdfield_t]:
+        ) -> Callable[[NDArray[numpy.float64]], fdfield_t]:
     """
     Generate an operator which converts a vectorized spatial-frequency-space
      `h_mn` into an E-field distribution, i.e.
@@ -272,7 +276,7 @@ def hmn_2_exyz(k0: numpy.ndarray,
 
     k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
-    def operator(h: numpy.ndarray) -> fdfield_t:
+    def operator(h: NDArray[numpy.float64]) -> fdfield_t:
         hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
         d_xyz = (n * hin_m
                - m * hin_n) * k_mag
@@ -283,10 +287,11 @@ def hmn_2_exyz(k0: numpy.ndarray,
     return operator
 
 
-def hmn_2_hxyz(k0: numpy.ndarray,
+def hmn_2_hxyz(
-               G_matrix: numpy.ndarray,
+        k0: ArrayLike,
+        G_matrix: ArrayLike,
         epsilon: fdfield_t
-               ) -> Callable[[numpy.ndarray], fdfield_t]:
+        ) -> Callable[[NDArray[numpy.float64]], fdfield_t]:
     """
     Generate an operator which converts a vectorized spatial-frequency-space
      `h_mn` into an H-field distribution, i.e.
@@ -309,7 +314,7 @@ def hmn_2_hxyz(k0: numpy.ndarray,
     shape = epsilon[0].shape + (1,)
     _k_mag, m, n = generate_kmn(k0, G_matrix, shape)
 
-    def operator(h: numpy.ndarray) -> fdfield_t:
+    def operator(h: NDArray[numpy.float64]) -> fdfield_t:
         hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
         h_xyz = (m * hin_m
                + n * hin_n)
@@ -318,11 +323,12 @@ def hmn_2_hxyz(k0: numpy.ndarray,
     return operator
 
 
-def inverse_maxwell_operator_approx(k0: numpy.ndarray,
+def inverse_maxwell_operator_approx(
-                                    G_matrix: numpy.ndarray,
+        k0: ArrayLike,
+        G_matrix: ArrayLike,
         epsilon: fdfield_t,
-                                    mu: fdfield_t = None
+        mu: Optional[fdfield_t] = None,
-                                    ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+        ) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]:
     """
     Generate an approximate inverse of the Maxwell operator,
 
@@ -351,7 +357,7 @@ def inverse_maxwell_operator_approx(k0: numpy.ndarray,
     if mu is not None:
         mu = numpy.stack(mu, 3)
 
-    def operator(h: numpy.ndarray) -> numpy.ndarray:
+    def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
         """
         Approximate inverse Maxwell operator for Bloch eigenmode simulation.
 
@@ -397,17 +403,18 @@ def inverse_maxwell_operator_approx(k0: numpy.ndarray,
     return operator
 
 
-def find_k(frequency: float,
+def find_k(
+        frequency: float,
         tolerance: float,
-           direction: numpy.ndarray,
+        direction: ArrayLike,
-           G_matrix: numpy.ndarray,
+        G_matrix: ArrayLike,
         epsilon: fdfield_t,
-           mu: fdfield_t = None,
+        mu: Optional[fdfield_t] = None,
         band: int = 0,
         k_min: float = 0,
         k_max: float = 0.5,
-           solve_callback: Callable = None
+        solve_callback: Optional[Callable] = None
-           ) -> Tuple[numpy.ndarray, float]:
+        ) -> Tuple[NDArray[numpy.float64], float]:
     """
     Search for a bloch vector that has a given frequency.
 
@@ -429,7 +436,7 @@ def find_k(frequency: float,
 
     direction = numpy.array(direction) / norm(direction)
 
-    def get_f(k0_mag: float, band: int = 0) -> numpy.ndarray:
+    def get_f(k0_mag: float, band: int = 0) -> float:
         k0 = direction * k0_mag
         n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
         f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
@@ -437,23 +444,26 @@ def find_k(frequency: float,
             solve_callback(k0_mag, n, v, f)
         return f
 
-    res = scipy.optimize.minimize_scalar(lambda x: abs(get_f(x, band) - frequency),
+    res = scipy.optimize.minimize_scalar(
+        lambda x: abs(get_f(x, band) - frequency),
         (k_min + k_max) / 2,
         method='Bounded',
         bounds=(k_min, k_max),
-                                         options={'xatol': abs(tolerance)})
+        options={'xatol': abs(tolerance)},
+        )
     return res.x * direction, res.fun + frequency
 
 
-def eigsolve(num_modes: int,
+def eigsolve(
-             k0: numpy.ndarray,
+        num_modes: int,
-             G_matrix: numpy.ndarray,
+        k0: ArrayLike,
+        G_matrix: ArrayLike,
         epsilon: fdfield_t,
-             mu: fdfield_t = None,
+        mu: Optional[fdfield_t] = None,
         tolerance: float = 1e-20,
         max_iters: int = 10000,
         reset_iters: int = 100,
-             ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+        ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
     """
     Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
      k0 of the specified structure.
@@ -625,7 +635,7 @@ def eigsolve(num_modes: int,
         theta = result.x
 
         improvement = numpy.abs(E - new_E) * 2 / numpy.abs(E + new_E)
-        logger.info('linmin improvement {}'.format(improvement))
+        logger.info(f'linmin improvement {improvement}')
         Z *= numpy.cos(theta)
         Z += D * numpy.sin(theta)
 
@@ -651,7 +661,7 @@ def eigsolve(num_modes: int,
         f = numpy.sqrt(-numpy.real(n))
         df = numpy.sqrt(-numpy.real(n + eigness))
         neff_err = kmag * (1 / df - 1 / f)
-        logger.info('eigness {}: {}\n neff_err: {}'.format(i, eigness, neff_err))
+        logger.info(f'eigness {i}: {eigness}\n neff_err: {neff_err}')
 
     order = numpy.argsort(numpy.abs(eigvals))
     return eigvals[order], eigvecs.T[order]
@@ -685,9 +695,9 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to
         return x, fx, dfx
 '''
 
-def _rtrace_AtB(A: numpy.ndarray, B: numpy.ndarray) -> numpy.ndarray:
+def _rtrace_AtB(A: NDArray[numpy.float64], B: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
     return real(numpy.sum(A.conj() * B))
 
-def _symmetrize(A: numpy.ndarray) -> numpy.ndarray:
+def _symmetrize(A: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
     return (A + A.conj().T) * 0.5