jan/meanas
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

typing updates

Browse Source
This commit is contained in:
Jan Petykiewicz 2022-10-04 12:43:26 -07:00
parent eedcc7b919
commit d42a625e5f
2 changed files with 73 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
examples/fdfd.py
View File

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

134
meanas/fdfd/bloch.py
View File

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