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

typing and formatting updates

Browse Source
This commit is contained in:
Jan Petykiewicz 2022-10-04 14:32:40 -07:00
parent d42a625e5f
commit faecc79179
22 changed files with 621 additions and 486 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
meanas/eigensolvers.py
View File

@ -2,16 +2,18 @@
Solvers for eigenvalue / eigenvector problems
"""
from typing import Tuple, Callable, Optional, Union
import numpy # type: ignore
from numpy.linalg import norm # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm
from scipy import sparse # type: ignore
import scipy.sparse.linalg as spalg # type: ignore
def power_iteration(operator: sparse.spmatrix,
guess_vector: Optional[numpy.ndarray] = None,
def power_iteration(
operator: sparse.spmatrix,
guess_vector: Optional[NDArray[numpy.float64]] = None,
iterations: int = 20,
) -> Tuple[complex, numpy.ndarray]:
) -> Tuple[complex, NDArray[numpy.float64]]:
"""
Use power iteration to estimate the dominant eigenvector of a matrix.
@ -37,12 +39,13 @@ def power_iteration(operator: sparse.spmatrix,
return lm_eigval, v
def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOperator],
guess_vector: numpy.ndarray,
def rayleigh_quotient_iteration(
operator: Union[sparse.spmatrix, spalg.LinearOperator],
guess_vector: NDArray[numpy.float64],
iterations: int = 40,
tolerance: float = 1e-13,
solver: Optional[Callable[..., numpy.ndarray]] = None,
) -> Tuple[complex, numpy.ndarray]:
solver: Optional[Callable[..., NDArray[numpy.float64]]] = None,
) -> Tuple[complex, NDArray[numpy.float64]]:
"""
Use Rayleigh quotient iteration to refine an eigenvector guess.
@ -69,11 +72,13 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe
solver = spalg.spsolve
except TypeError:
def shift(eigval: float) -> spalg.LinearOperator:
return spalg.LinearOperator(shape=operator.shape,
return spalg.LinearOperator(
shape=operator.shape,
dtype=operator.dtype,
matvec=lambda v: eigval * v)
matvec=lambda v: eigval * v,
)
if solver is None:
def solver(A: spalg.LinearOperator, b: numpy.ndarray) -> numpy.ndarray:
def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.float64]:
return spalg.bicgstab(A, b)[0]
assert(solver is not None)
@ -90,10 +95,11 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe
return eigval, v
def signed_eigensolve(operator: Union[sparse.spmatrix, spalg.LinearOperator],
def signed_eigensolve(
operator: Union[sparse.spmatrix, spalg.LinearOperator],
how_many: int,
negative: bool = False,
) -> Tuple[numpy.ndarray, numpy.ndarray]:
) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
"""
Find the largest-magnitude positive-only (or negative-only) eigenvalues and
eigenvectors of the provided matrix.

20
meanas/fdfd/bloch.py
View File

@ -80,7 +80,7 @@ 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, Union
import logging
import numpy
from numpy import pi, real, trace
@ -433,11 +433,10 @@ def find_k(
`(k, actual_frequency)`
The found k-vector and its frequency.
"""
direction = numpy.array(direction) / norm(direction)
def get_f(k0_mag: float, band: int = 0) -> float:
k0 = direction * k0_mag
k0 = direction * k0_mag # type: ignore
n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
if solve_callback:
@ -482,6 +481,8 @@ def eigsolve(
`(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the
vector `eigenvectors[i, :]`
"""
k0 = numpy.array(k0, copy=False)
h_size = 2 * epsilon[0].size
kmag = norm(G_matrix @ k0)
@ -497,9 +498,9 @@ def eigsolve(
y_shape = (h_size, num_modes)
prev_E = 0
d_scale = 1
prev_traceGtKG = 0
prev_E = 0.0
d_scale = 1.0
prev_traceGtKG = 0.0
#prev_theta = 0.5
D = numpy.zeros(shape=y_shape, dtype=complex)
@ -545,7 +546,7 @@ def eigsolve(
if prev_traceGtKG == 0 or i % reset_iters == 0:
logger.info('CG reset')
gamma = 0
gamma = 0.0
else:
gamma = traceGtKG / prev_traceGtKG
@ -695,7 +696,10 @@ 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: NDArray[numpy.float64], B: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
def _rtrace_AtB(
A: NDArray[numpy.float64],
B: Union[NDArray[numpy.float64], float],
) -> float:
return real(numpy.sum(A.conj() * B))
def _symmetrize(A: NDArray[numpy.float64]) -> NDArray[numpy.float64]:

12
meanas/fdfd/farfield.py
View File

@ -2,14 +2,15 @@
Functions for performing near-to-farfield transformation (and the reverse).
"""
from typing import Dict, List, Any
import numpy # type: ignore
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift # type: ignore
from numpy import pi # type: ignore
import numpy
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
from numpy import pi
from ..fdmath import fdfield_t
def near_to_farfield(E_near: fdfield_t,
def near_to_farfield(
E_near: fdfield_t,
H_near: fdfield_t,
dx: float,
dy: float,
@ -120,7 +121,8 @@ def near_to_farfield(E_near: fdfield_t,
return outputs
def far_to_nearfield(E_far: fdfield_t,
def far_to_nearfield(
E_far: fdfield_t,
H_far: fdfield_t,
dkx: float,
dky: float,

25
meanas/fdfd/functional.py
View File

@ -5,8 +5,8 @@ Functional versions of many FDFD operators. These can be useful for performing
The functions generated here expect `fdfield_t` inputs with shape (3, X, Y, Z),
e.g. E = [E_x, E_y, E_z] where each component has shape (X, Y, Z)
"""
from typing import Callable, Tuple
import numpy # type: ignore
from typing import Callable, Tuple, Optional
import numpy
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
from ..fdmath.functional import curl_forward, curl_back
@ -15,7 +15,8 @@ from ..fdmath.functional import curl_forward, curl_back
__author__ = 'Jan Petykiewicz'
def e_full(omega: complex,
def e_full(
omega: complex,
dxes: dx_lists_t,
epsilon: fdfield_t,
mu: fdfield_t = None
@ -50,7 +51,8 @@ def e_full(omega: complex,
return op_mu
def eh_full(omega: complex,
def eh_full(
omega: complex,
dxes: dx_lists_t,
epsilon: fdfield_t,
mu: fdfield_t = None
@ -86,9 +88,10 @@ def eh_full(omega: complex,
return op_mu
def e2h(omega: complex,
def e2h(
omega: complex,
dxes: dx_lists_t,
mu: fdfield_t = None,
mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t:
"""
Utility operator for converting the `E` field into the `H` field.
@ -117,9 +120,10 @@ def e2h(omega: complex,
return e2h_mu
def m2j(omega: complex,
def m2j(
omega: complex,
dxes: dx_lists_t,
mu: fdfield_t = None,
mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t:
"""
Utility operator for converting magnetic current `M` distribution
@ -151,11 +155,12 @@ def m2j(omega: complex,
return m2j_mu
def e_tfsf_source(TF_region: fdfield_t,
def e_tfsf_source(
TF_region: fdfield_t,
omega: complex,
dxes: dx_lists_t,
epsilon: fdfield_t,
mu: fdfield_t = None,
mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t:
"""
Operator that turns an E-field distribution into a total-field/scattered-field

30
meanas/fdfd/operators.py
View File

@ -28,7 +28,7 @@ The following operators are included:
"""
from typing import Tuple, Optional
import numpy # type: ignore
import numpy
import scipy.sparse as sparse # type: ignore
from ..fdmath import vec, dx_lists_t, vfdfield_t
@ -38,7 +38,8 @@ from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl
__author__ = 'Jan Petykiewicz'
def e_full(omega: complex,
def e_full(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
@ -96,7 +97,8 @@ def e_full(omega: complex,
return op
def e_full_preconditioners(dxes: dx_lists_t
def e_full_preconditioners(
dxes: dx_lists_t,
) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
"""
Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator.
@ -122,7 +124,8 @@ def e_full_preconditioners(dxes: dx_lists_t
return P_left, P_right
def h_full(omega: complex,
def h_full(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
@ -178,12 +181,13 @@ def h_full(omega: complex,
return A
def eh_full(omega: complex,
def eh_full(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
pec: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None
pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
"""
Wave operator for `[E, H]` field representation. This operator implements Maxwell's
@ -247,7 +251,8 @@ def eh_full(omega: complex,
return A
def e2h(omega: complex,
def e2h(
omega: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None,
@ -278,9 +283,10 @@ def e2h(omega: complex,
return op
def m2j(omega: complex,
def m2j(
omega: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None
mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
"""
Operator for converting a magnetic current M into an electric current J.
@ -357,7 +363,8 @@ def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
return P
def e_tfsf_source(TF_region: vfdfield_t,
def e_tfsf_source(
TF_region: vfdfield_t,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
@ -387,7 +394,8 @@ def e_tfsf_source(TF_region: vfdfield_t,
return (A @ Q - Q @ A) / (-1j * omega)
def e_boundary_source(mask: vfdfield_t,
def e_boundary_source(
mask: vfdfield_t,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,

31
meanas/fdfd/scpml.py
View File

@ -3,7 +3,9 @@ Functions for creating stretched coordinate perfectly matched layer (PML) absorb
"""
from typing import Sequence, Union, Callable, Optional, List
import numpy # type: ignore
import numpy
from numpy.typing import ArrayLike, NDArray
__author__ = 'Jan Petykiewicz'
@ -13,7 +15,8 @@ s_function_t = Callable[[float], float]
"""Typedef for s-functions, see `prepare_s_function()`"""
def prepare_s_function(ln_R: float = -16,
def prepare_s_function(
ln_R: float = -16,
m: float = 4
) -> s_function_t:
"""
@ -29,18 +32,19 @@ def prepare_s_function(ln_R: float = -16,
of the cell width; needs to be divided by `sqrt(epilon_effective) * real(omega))`
before use.
"""
def s_factor(distance: numpy.ndarray) -> numpy.ndarray:
def s_factor(distance: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
s_max = (m + 1) * ln_R / 2 # / 2 because we assume periodic boundaries
return s_max * (distance ** m)
return s_factor
def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
thicknesses: Union[numpy.ndarray, Sequence[int]],
def uniform_grid_scpml(
shape: ArrayLike, # ints
thicknesses: ArrayLike, # ints
omega: float,
epsilon_effective: float = 1.0,
s_function: Optional[s_function_t] = None,
) -> List[List[numpy.ndarray]]:
) -> List[List[NDArray[numpy.float64]]]:
"""
Create dx arrays for a uniform grid with a cell width of 1 and a pml.
@ -67,7 +71,11 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
s_function = prepare_s_function()
# Normalized distance to nearest boundary
def ll(u: numpy.ndarray, n: numpy.ndarray, t: numpy.ndarray) -> numpy.ndarray:
def ll(
u: NDArray[numpy.float64],
n: NDArray[numpy.float64],
t: NDArray[numpy.float64],
) -> NDArray[numpy.float64]:
return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t
dx_a = [numpy.array(numpy.inf)] * 3
@ -88,14 +96,15 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
return [dx_a, dx_b]
def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
def stretch_with_scpml(
dxes: List[List[NDArray[numpy.float64]]],
axis: int,
polarity: int,
omega: float,
epsilon_effective: float = 1.0,
thickness: int = 10,
s_function: Optional[s_function_t] = None,
) -> List[List[numpy.ndarray]]:
) -> List[List[NDArray[numpy.float64]]]:
"""
Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis.
@ -132,7 +141,7 @@ def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
bound = pos[thickness]
d = bound - pos[0]
def l_d(x: numpy.ndarray) -> numpy.ndarray:
def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
return (bound - x) / (bound - pos[0])
slc = slice(thickness)
@ -142,7 +151,7 @@ def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
bound = pos[-thickness - 1]
d = pos[-1] - bound
def l_d(x: numpy.ndarray) -> numpy.ndarray:
def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
return (x - bound) / (pos[-1] - bound)
if thickness == 0:

29
meanas/fdfd/solvers.py
View File

@ -2,11 +2,12 @@
Solvers and solver interface for FDFD problems.
"""
from typing import Callable, Dict, Any
from typing import Callable, Dict, Any, Optional
import logging
import numpy # type: ignore
from numpy.linalg import norm # type: ignore
import numpy
from numpy.typing import ArrayLike, NDArray
from numpy.linalg import norm
import scipy.sparse.linalg # type: ignore
from ..fdmath import dx_lists_t, vfdfield_t
@ -16,10 +17,11 @@ from . import operators
logger = logging.getLogger(__name__)
def _scipy_qmr(A: scipy.sparse.csr_matrix,
b: numpy.ndarray,
def _scipy_qmr(
A: scipy.sparse.csr_matrix,
b: ArrayLike,
**kwargs: Any,
) -> numpy.ndarray:
) -> NDArray[numpy.float64]:
"""
Wrapper for scipy.sparse.linalg.qmr
@ -37,14 +39,14 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix,
'''
ii = 0
def log_residual(xk: numpy.ndarray) -> None:
def log_residual(xk: ArrayLike) -> None:
nonlocal ii
ii += 1
if ii % 100 == 0:
logger.info('Solver residual at iteration {} : {}'.format(ii, norm(A @ xk - b)))
if 'callback' in kwargs:
def augmented_callback(xk: numpy.ndarray) -> None:
def augmented_callback(xk: ArrayLike) -> None:
log_residual(xk)
kwargs['callback'](xk)
@ -60,7 +62,8 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix,
return x
def generic(omega: complex,
def generic(
omega: complex,
dxes: dx_lists_t,
J: vfdfield_t,
epsilon: vfdfield_t,
@ -68,8 +71,8 @@ def generic(omega: complex,
pec: vfdfield_t = None,
pmc: vfdfield_t = None,
adjoint: bool = False,
matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr,
matrix_solver_opts: Dict[str, Any] = None,
matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
matrix_solver_opts: Optional[Dict[str, Any]] = None,
) -> vfdfield_t:
"""
Conjugate gradient FDFD solver using CSR sparse matrices.
@ -90,8 +93,8 @@ def generic(omega: complex,
adjoint: If true, solves the adjoint problem.
matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`,
where `A`: `scipy.sparse.csr_matrix`;
`b`: `numpy.ndarray`;
`x`: `numpy.ndarray`;
`b`: `ArrayLike`;
`x`: `ArrayLike`;
Default is a wrapped version of `scipy.sparse.linalg.qmr()`
which doesn't return convergence info and logs the residual
every 100 iterations.

58
meanas/fdfd/waveguide_2d.py
View File

@ -179,8 +179,9 @@ to account for numerical dispersion if the result is introduced into a space wit
# TODO update module docs
from typing import List, Tuple, Optional, Any
import numpy # type: ignore
from numpy.linalg import norm # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm
import scipy.sparse as sparse # type: ignore
from ..fdmath.operators import deriv_forward, deriv_back, cross
@ -191,7 +192,8 @@ from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
__author__ = 'Jan Petykiewicz'
def operator_e(omega: complex,
def operator_e(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
@ -257,7 +259,8 @@ def operator_e(omega: complex,
return op
def operator_h(omega: complex,
def operator_h(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
@ -324,7 +327,8 @@ def operator_h(omega: complex,
return op
def normalized_fields_e(e_xy: numpy.ndarray,
def normalized_fields_e(
e_xy: ArrayLike,
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
@ -358,7 +362,8 @@ def normalized_fields_e(e_xy: numpy.ndarray,
return e_norm, h_norm
def normalized_fields_h(h_xy: numpy.ndarray,
def normalized_fields_h(
h_xy: ArrayLike,
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
@ -392,8 +397,9 @@ def normalized_fields_h(h_xy: numpy.ndarray,
return e_norm, h_norm
def _normalized_fields(e: numpy.ndarray,
h: numpy.ndarray,
def _normalized_fields(
e: ArrayLike,
h: ArrayLike,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
@ -434,7 +440,8 @@ def _normalized_fields(e: numpy.ndarray,
return e, h
def exy2h(wavenumber: complex,
def exy2h(
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
@ -459,7 +466,8 @@ def exy2h(wavenumber: complex,
return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)
def hxy2e(wavenumber: complex,
def hxy2e(
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
@ -484,7 +492,8 @@ def hxy2e(wavenumber: complex,
return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)
def hxy2h(wavenumber: complex,
def hxy2h(
wavenumber: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix:
@ -517,7 +526,8 @@ def hxy2h(wavenumber: complex,
return op
def exy2e(wavenumber: complex,
def exy2e(
wavenumber: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
) -> sparse.spmatrix:
@ -550,7 +560,8 @@ def exy2e(wavenumber: complex,
return op
def e2h(wavenumber: complex,
def e2h(
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None
@ -574,7 +585,8 @@ def e2h(wavenumber: complex,
return op
def h2e(wavenumber: complex,
def h2e(
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t
@ -636,12 +648,13 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
return cross([Dbx, Dby, Bz])
def h_err(h: vfdfield_t,
def h_err(
h: vfdfield_t,
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: vfdfield_t = None
mu: Optional[vfdfield_t] = None
) -> float:
"""
Calculates the relative error in the H field
@ -670,12 +683,13 @@ def h_err(h: vfdfield_t,
return norm(op) / norm(h)
def e_err(e: vfdfield_t,
def e_err(
e: vfdfield_t,
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: vfdfield_t = None
mu: vfdfield_t = Optional[None]
) -> float:
"""
Calculates the relative error in the E field
@ -703,13 +717,14 @@ def e_err(e: vfdfield_t,
return norm(op) / norm(e)
def solve_modes(mode_numbers: List[int],
def solve_modes(
mode_numbers: List[int],
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: vfdfield_t = None,
mode_margin: int = 2,
) -> Tuple[numpy.ndarray, List[complex]]:
) -> Tuple[NDArray[numpy.float64], List[complex]]:
"""
Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
@ -752,7 +767,8 @@ def solve_modes(mode_numbers: List[int],
return e_xys, wavenumbers
def solve_mode(mode_number: int,
def solve_mode(
mode_number: int,
*args: Any,
**kwargs: Any,
) -> Tuple[vfdfield_t, complex]:

27
meanas/fdfd/waveguide_3d.py
View File

@ -5,13 +5,15 @@ This module relies heavily on `waveguide_2d` and mostly just transforms
its parameters into 2D equivalents and expands the results back into 3D.
"""
from typing import Dict, Optional, Sequence, Union, Any
import numpy # type: ignore
import numpy
from numpy.typing import NDArray
from ..fdmath import vec, unvec, dx_lists_t, fdfield_t
from . import operators, waveguide_2d
def solve_mode(mode_number: int,
def solve_mode(
mode_number: int,
omega: complex,
dxes: dx_lists_t,
axis: int,
@ -19,7 +21,7 @@ def solve_mode(mode_number: int,
slices: Sequence[slice],
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
) -> Dict[str, Union[complex, numpy.ndarray]]:
) -> Dict[str, Union[complex, NDArray[numpy.float_]]]:
"""
Given a 3D grid, selects a slice from the grid and attempts to
solve for an eigenmode propagating through that slice.
@ -36,7 +38,13 @@ def solve_mode(mode_number: int,
mu: Magnetic permeability (default 1 everywhere)
Returns:
`{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}`
```
{
'E': List[NDArray[numpy.float_]],
'H': List[NDArray[numpy.float_]],
'wavenumber': complex,
}
```
"""
if mu is None:
mu = numpy.ones_like(epsilon)
@ -97,7 +105,8 @@ def solve_mode(mode_number: int,
return results
def compute_source(E: fdfield_t,
def compute_source(
E: fdfield_t,
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
@ -142,7 +151,8 @@ def compute_source(E: fdfield_t,
return J
def compute_overlap_e(E: fdfield_t,
def compute_overlap_e(
E: fdfield_t,
wavenumber: complex,
dxes: dx_lists_t,
axis: int,
@ -154,6 +164,8 @@ def compute_overlap_e(E: fdfield_t,
mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn)
[assumes reflection symmetry].
TODO: add reference
Args:
E: E-field of the mode
H: H-field of the mode (advanced by half of a Yee cell from E)
@ -187,7 +199,8 @@ def compute_overlap_e(E: fdfield_t,
return Etgt
def expand_e(E: fdfield_t,
def expand_e(
E: fdfield_t,
wavenumber: complex,
dxes: dx_lists_t,
axis: int,

16
meanas/fdfd/waveguide_cyl.py
View File

@ -9,7 +9,7 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
# TODO update module docs
from typing import Dict, Union
import numpy # type: ignore
import numpy
import scipy.sparse as sparse # type: ignore
from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t
@ -17,7 +17,8 @@ from ..fdmath.operators import deriv_forward, deriv_back
from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
def cylindrical_operator(omega: complex,
def cylindrical_operator(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
r0: float,
@ -78,7 +79,8 @@ def cylindrical_operator(omega: complex,
return op
def solve_mode(mode_number: int,
def solve_mode(
mode_number: int,
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
@ -99,7 +101,13 @@ def solve_mode(mode_number: int,
r within the simulation domain.
Returns:
`{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}`
```
{
'E': List[NDArray[numpy.float_]],
'H': List[NDArray[numpy.float_]],
'wavenumber': complex,
}
```
"""
'''

25
meanas/fdmath/functional.py
View File

@ -5,12 +5,14 @@ Basic discrete calculus etc.
"""
from typing import Sequence, Tuple, Optional, Callable
import numpy # type: ignore
import numpy
from numpy.typing import NDArray
from .types import fdfield_t, fdfield_updater_t
def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None
def deriv_forward(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
"""
Utility operators for taking discretized derivatives (backward variant).
@ -33,7 +35,8 @@ def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None
return derivs
def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None
def deriv_back(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
"""
Utility operators for taking discretized derivatives (forward variant).
@ -56,7 +59,9 @@ def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None
return derivs
def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t:
def curl_forward(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> fdfield_updater_t:
"""
Curl operator for use with the E field.
@ -83,7 +88,9 @@ def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_upda
return ce_fun
def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t:
def curl_back(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> fdfield_updater_t:
"""
Create a function which takes the backward curl of a field.
@ -110,7 +117,9 @@ def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater
return ch_fun
def curl_forward_parts(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> Callable:
def curl_forward_parts(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> Callable:
Dx, Dy, Dz = deriv_forward(dx_e)
def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:
@ -121,7 +130,9 @@ def curl_forward_parts(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> Callab
return mkparts_fwd
def curl_back_parts(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> Callable:
def curl_back_parts(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> Callable:
Dx, Dy, Dz = deriv_back(dx_e)
def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:

37
meanas/fdmath/operators.py
View File

@ -4,13 +4,18 @@ Matrix operators for finite difference simulations
Basic discrete calculus etc.
"""
from typing import Sequence, List
import numpy # type: ignore
import numpy
from numpy.typing import NDArray
import scipy.sparse as sparse # type: ignore
from .types import vfdfield_t
def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> sparse.spmatrix:
def shift_circ(
axis: int,
shape: Sequence[int],
shift_distance: int = 1,
) -> sparse.spmatrix:
"""
Utility operator for performing a circular shift along a specified axis by a
specified number of elements.
@ -46,7 +51,11 @@ def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> spar
return d
def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1) -> sparse.spmatrix:
def shift_with_mirror(
axis: int,
shape: Sequence[int],
shift_distance: int = 1,
) -> sparse.spmatrix:
"""
Utility operator for performing an n-element shift along a specified axis, with mirror
boundary conditions applied to the cells beyond the receding edge.
@ -67,7 +76,7 @@ def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1)
raise Exception('Shift ({}) is too large for axis {} of size {}'.format(
shift_distance, axis, shape[axis]))
def mirrored_range(n: int, s: int) -> numpy.ndarray:
def mirrored_range(n: int, s: int) -> NDArray[numpy.int_]:
v = numpy.arange(n) + s
v = numpy.where(v >= n, 2 * n - v - 1, v)
v = numpy.where(v < 0, - 1 - v, v)
@ -87,7 +96,9 @@ def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1)
return d
def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
def deriv_forward(
dx_e: Sequence[NDArray[numpy.float_]],
) -> List[sparse.spmatrix]:
"""
Utility operators for taking discretized derivatives (forward variant).
@ -112,7 +123,9 @@ def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
return Ds
def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
def deriv_back(
dx_h: Sequence[NDArray[numpy.float_]],
) -> List[sparse.spmatrix]:
"""
Utility operators for taking discretized derivatives (backward variant).
@ -137,7 +150,9 @@ def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
return Ds
def cross(B: Sequence[sparse.spmatrix]) -> sparse.spmatrix:
def cross(
B: Sequence[sparse.spmatrix],
) -> sparse.spmatrix:
"""
Cross product operator
@ -203,7 +218,9 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
return avg_forward(axis, shape).T
def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix:
def curl_forward(
dx_e: Sequence[NDArray[numpy.float_]],
) -> sparse.spmatrix:
"""
Curl operator for use with the E field.
@ -217,7 +234,9 @@ def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix:
return cross(deriv_forward(dx_e))
def curl_back(dx_h: Sequence[numpy.ndarray]) -> sparse.spmatrix:
def curl_back(
dx_h: Sequence[NDArray[numpy.float_]],
) -> sparse.spmatrix:
"""
Curl operator for use with the H field.

43
meanas/fdmath/types.py
View File

@ -2,31 +2,20 @@
Types shared across multiple submodules
"""
from typing import Sequence, Callable, MutableSequence
import numpy # type: ignore
import numpy
from numpy.typing import NDArray
# Field types
# TODO: figure out a better way to set the docstrings without creating actual subclasses?
# Probably not a big issue since they're only used for type hinting
class fdfield_t(numpy.ndarray):
fdfield_t = NDArray[numpy.float_]
"""Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
vfdfield_t = NDArray[numpy.float_]
"""Linearized vector field (single vector of length 3*X*Y*Z)"""
dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]]
"""
Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)
This is actually is just an unaltered `numpy.ndarray`
"""
pass
class vfdfield_t(numpy.ndarray):
"""
Linearized vector field (single vector of length 3*X*Y*Z)
This is actually just an unaltered `numpy.ndarray`
"""
pass
dx_lists_t = Sequence[Sequence[numpy.ndarray]]
'''
'dxes' datastructure which contains grid cell width information in the following format:
[[[dx_e[0], dx_e[1], ...], [dy_e[0], ...], [dz_e[0], ...]],
@ -34,15 +23,11 @@ dx_lists_t = Sequence[Sequence[numpy.ndarray]]
where `dx_e[0]` is the x-width of the `x=0` cells, as used when calculating dE/dx,
and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
'''
"""
dx_lists_mut = MutableSequence[MutableSequence[numpy.ndarray]]
'''
Mutable version of `dx_lists_t`
'''
dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]]
"""Mutable version of `dx_lists_t`"""
fdfield_updater_t = Callable[..., fdfield_t]
'''
Convenience type for functions which take and return an fdfield_t
'''
"""Convenience type for functions which take and return an fdfield_t"""

13
meanas/fdmath/vectorization.py
View File

@ -5,7 +5,8 @@ Vectorized versions of the field use row-major (ie., C-style) ordering.
"""
from typing import Optional, overload, Union, List
import numpy # type: ignore
import numpy
from numpy.typing import ArrayLike
from .types import fdfield_t, vfdfield_t
@ -15,10 +16,10 @@ def vec(f: None) -> None:
pass
@overload
def vec(f: Union[fdfield_t, List[numpy.ndarray]]) -> vfdfield_t:
def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t:
pass
def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield_t]:
def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]:
"""
Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
@ -37,14 +38,14 @@ def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield
@overload
def unvec(v: None, shape: numpy.ndarray) -> None:
def unvec(v: None, shape: ArrayLike) -> None:
pass
@overload
def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t:
def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t:
pass
def unvec(v: Optional[vfdfield_t], shape: numpy.ndarray) -> Optional[fdfield_t]:
def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> Optional[fdfield_t]:
"""
Perform the inverse of vec(): take a 1D ndarray and output a 3D field
of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional

3
meanas/fdtd/boundaries.py
View File

@ -9,7 +9,8 @@ from typing import Tuple, Any, List
from ..fdmath import fdfield_t, fdfield_updater_t
def conducting_boundary(direction: int,
def conducting_boundary(
direction: int,
polarity: int
) -> Tuple[fdfield_updater_t, fdfield_updater_t]:
dirs = [0, 1, 2]

28
meanas/fdtd/energy.py
View File

@ -1,5 +1,5 @@
from typing import Optional, Union
import numpy # type: ignore
import numpy
from ..fdmath import dx_lists_t, fdfield_t
from ..fdmath.functional import deriv_back
@ -8,7 +8,8 @@ from ..fdmath.functional import deriv_back
# TODO documentation
def poynting(e: fdfield_t,
def poynting(
e: fdfield_t,
h: fdfield_t,
dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
@ -87,7 +88,8 @@ def poynting(e: fdfield_t,
return s
def poynting_divergence(s: Optional[fdfield_t] = None,
def poynting_divergence(
s: Optional[fdfield_t] = None,
*,
e: Optional[fdfield_t] = None,
h: Optional[fdfield_t] = None,
@ -124,7 +126,8 @@ def poynting_divergence(s: Optional[fdfield_t] = None,
return ds
def energy_hstep(e0: fdfield_t,
def energy_hstep(
e0: fdfield_t,
h1: fdfield_t,
e2: fdfield_t,
epsilon: Optional[fdfield_t] = None,
@ -151,7 +154,8 @@ def energy_hstep(e0: fdfield_t,
return u
def energy_estep(h0: fdfield_t,
def energy_estep(
h0: fdfield_t,
e1: fdfield_t,
h2: fdfield_t,
epsilon: Optional[fdfield_t] = None,
@ -178,7 +182,8 @@ def energy_estep(h0: fdfield_t,
return u
def delta_energy_h2e(dt: float,
def delta_energy_h2e(
dt: float,
e0: fdfield_t,
h1: fdfield_t,
e2: fdfield_t,
@ -210,7 +215,8 @@ def delta_energy_h2e(dt: float,
return du
def delta_energy_e2h(dt: float,
def delta_energy_e2h(
dt: float,
h0: fdfield_t,
e1: fdfield_t,
h2: fdfield_t,
@ -242,7 +248,8 @@ def delta_energy_e2h(dt: float,
return du
def delta_energy_j(j0: fdfield_t,
def delta_energy_j(
j0: fdfield_t,
e1: fdfield_t,
dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
@ -264,11 +271,12 @@ def delta_energy_j(j0: fdfield_t,
return du
def dxmul(ee: fdfield_t,
def dxmul(
ee: fdfield_t,
hh: fdfield_t,
epsilon: Optional[Union[fdfield_t, float]] = None,
mu: Optional[Union[fdfield_t, float]] = None,
dxes: Optional[dx_lists_t] = None
dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
if epsilon is None:
epsilon = 1

5
meanas/fdtd/pml.py
View File

@ -8,7 +8,8 @@ PML implementations
# TODO retest pmls!
from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional
import numpy # type: ignore
import numpy
from typing import NDArray
from ..fdmath import fdfield_t, dx_lists_t
from ..fdmath.functional import deriv_forward, deriv_back
@ -61,7 +62,7 @@ def cpml_params(
expand_slice_l[axis] = slice(None)
expand_slice = tuple(expand_slice_l)
def par(x: numpy.ndarray) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
def par(x: NDArray[numpy.float64]) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
scaling = (x / thickness) ** m
sigma = scaling * sigma_max
kappa = 1 + scaling * (kappa_max - 1)

13
meanas/test/conftest.py
View File

@ -4,7 +4,8 @@ Test fixtures
"""
from typing import Tuple, Iterable, List, Any
import numpy # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
import pytest # type: ignore
from .utils import PRNG
@ -34,11 +35,12 @@ def epsilon_fg(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(scope='module', params=['center', '000', 'random'])
def epsilon(request: FixtureRequest,
def epsilon(
request: FixtureRequest,
shape: Tuple[int, ...],
epsilon_bg: float,
epsilon_fg: float,
) -> Iterable[numpy.ndarray]:
) -> Iterable[NDArray[numpy.float64]]:
is3d = (numpy.array(shape) == 1).sum() == 0
if is3d:
if request.param == '000':
@ -72,10 +74,11 @@ def dx(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(scope='module', params=['uniform', 'centerbig'])
def dxes(request: FixtureRequest,
def dxes(
request: FixtureRequest,
shape: Tuple[int, ...],
dx: float,
) -> Iterable[List[List[numpy.ndarray]]]:
) -> Iterable[List[List[NDArray[numpy.float64]]]]:
if request.param == 'uniform':
dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
elif request.param == 'centerbig':

37
meanas/test/test_fdfd.py
View File

@ -1,7 +1,8 @@
from typing import List, Tuple, Iterable, Optional
import dataclasses
import pytest # type: ignore
import numpy # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
#from numpy.testing import assert_allclose, assert_array_equal
from .. import fdfd
@ -59,12 +60,12 @@ def omega(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(params=[None])
def pec(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
yield request.param
@pytest.fixture(params=[None])
def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
yield request.param
@ -75,10 +76,11 @@ def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
@pytest.fixture(params=['diag']) # 'center'
def j_distribution(request: FixtureRequest,
def j_distribution(
request: FixtureRequest,
shape: Tuple[int, ...],
j_mag: float,
) -> Iterable[numpy.ndarray]:
) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape, dtype=complex)
center_mask = numpy.zeros(shape, dtype=bool)
center_mask[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = True
@ -94,24 +96,25 @@ def j_distribution(request: FixtureRequest,
@dataclasses.dataclass()
class FDResult:
shape: Tuple[int, ...]
dxes: List[List[numpy.ndarray]]
epsilon: numpy.ndarray
dxes: List[List[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64]
omega: complex
j: numpy.ndarray
e: numpy.ndarray
pmc: numpy.ndarray
pec: numpy.ndarray
j: NDArray[numpy.float64]
e: NDArray[numpy.float64]
pmc: Optional[NDArray[numpy.float64]]
pec: Optional[NDArray[numpy.float64]]
@pytest.fixture()
def sim(request: FixtureRequest,
def sim(
request: FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
dxes: List[List[numpy.ndarray]],
j_distribution: numpy.ndarray,
epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]],
j_distribution: NDArray[numpy.float64],
omega: float,
pec: Optional[numpy.ndarray],
pmc: Optional[numpy.ndarray],
pec: Optional[NDArray[numpy.float64]],
pmc: Optional[NDArray[numpy.float64]],
) -> FDResult:
"""
Build simulation from parts

57
meanas/test/test_fdfd_pml.py
View File

@ -1,7 +1,8 @@
from typing import Optional, Tuple, Iterable, List
import pytest # type: ignore
import numpy # type: ignore
from numpy.testing import assert_allclose # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.testing import assert_allclose
from .. import fdfd
from ..fdmath import vec, unvec, dx_lists_mut
@ -48,12 +49,12 @@ def omega(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(params=[None])
def pec(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
yield request.param
@pytest.fixture(params=[None])
def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
yield request.param
@ -70,13 +71,14 @@ def src_polarity(request: FixtureRequest) -> Iterable[int]:
@pytest.fixture()
def j_distribution(request: FixtureRequest,
def j_distribution(
request: FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut,
omega: float,
src_polarity: int,
) -> Iterable[numpy.ndarray]:
) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape, dtype=complex)
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
@ -108,47 +110,60 @@ def j_distribution(request: FixtureRequest,
@pytest.fixture()
def epsilon(request: FixtureRequest,
def epsilon(
request: FixtureRequest,
shape: Tuple[int, ...],
epsilon_bg: float,
epsilon_fg: float,
) -> Iterable[numpy.ndarray]:
) -> Iterable[NDArray[numpy.float64]]:
epsilon = numpy.full(shape, epsilon_fg, dtype=float)
yield epsilon
@pytest.fixture(params=['uniform'])
def dxes(request: FixtureRequest,
def dxes(
request: FixtureRequest,
shape: Tuple[int, ...],
dx: float,
omega: float,
epsilon_fg: float,
) -> Iterable[List[List[numpy.ndarray]]]:
) -> Iterable[List[List[NDArray[numpy.float64]]]]:
if request.param == 'uniform':
dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
for axis in (dim,):
for polarity in (-1, 1):
dxes = fdfd.scpml.stretch_with_scpml(dxes, axis=axis, polarity=polarity,
omega=omega, epsilon_effective=epsilon_fg,
thickness=10)
dxes = fdfd.scpml.stretch_with_scpml(
dxes,
axis=axis,
polarity=polarity,
omega=omega,
epsilon_effective=epsilon_fg,
thickness=10,
)
yield dxes
@pytest.fixture()
def sim(request: FixtureRequest,
def sim(
request: FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut,
j_distribution: numpy.ndarray,
j_distribution: NDArray[numpy.float64],
omega: float,
pec: Optional[numpy.ndarray],
pmc: Optional[numpy.ndarray],
pec: Optional[NDArray[numpy.float64]],
pmc: Optional[NDArray[numpy.float64]],
) -> FDResult:
j_vec = vec(j_distribution)
eps_vec = vec(epsilon)
e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec,
matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11})
e_vec = fdfd.solvers.generic(
J=j_vec,
omega=omega,
dxes=dxes,
epsilon=eps_vec,
matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11},
)
e = unvec(e_vec, shape[1:])
sim = FDResult(

57
meanas/test/test_fdtd.py
View File

@ -1,8 +1,9 @@
from typing import List, Tuple, Iterable
from typing import List, Tuple, Iterable, Any, Dict
import dataclasses
import pytest # type: ignore
import numpy # type: ignore
#from numpy.testing import assert_allclose, assert_array_equal # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
#from numpy.testing import assert_allclose, assert_array_equal
from .. import fdtd
from .utils import assert_close, assert_fields_close, PRNG
@ -32,8 +33,10 @@ def test_initial_energy(sim: 'TDResult') -> None:
dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0)
u0 = (j0 * j0.conj() / sim.epsilon * dV).sum(axis=0)
args = {'dxes': sim.dxes,
'epsilon': sim.epsilon}
args: Dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
# Make sure initial energy and E dot J are correct
energy0 = fdtd.energy_estep(h0=h0, e1=e0, h2=h1, **args)
@ -49,8 +52,10 @@ def test_energy_conservation(sim: 'TDResult') -> None:
e0 = sim.es[0]
j0 = sim.js[0]
u = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes).sum()
args = {'dxes': sim.dxes,
'epsilon': sim.epsilon}
args: Dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
for ii in range(1, 8):
u_hstep = fdtd.energy_hstep(e0=sim.es[ii - 1], h1=sim.hs[ii], e2=sim.es[ii], **args)
@ -65,8 +70,10 @@ def test_energy_conservation(sim: 'TDResult') -> None:
def test_poynting_divergence(sim: 'TDResult') -> None:
args = {'dxes': sim.dxes,
'epsilon': sim.epsilon}
args: Dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
u_eprev = None
for ii in range(1, 8):
@ -96,8 +103,10 @@ def test_poynting_planes(sim: 'TDResult') -> None:
if mask.sum() > 1:
pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum()))
args = {'dxes': sim.dxes,
'epsilon': sim.epsilon}
args: Dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
mx = numpy.roll(mask, -1, axis=0)
my = numpy.roll(mask, -1, axis=1)
@ -149,13 +158,13 @@ def dt(request: FixtureRequest) -> Iterable[float]:
class TDResult:
shape: Tuple[int, ...]
dt: float
dxes: List[List[numpy.ndarray]]
epsilon: numpy.ndarray
j_distribution: numpy.ndarray
dxes: List[List[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64]
j_distribution: NDArray[numpy.float64]
j_steps: Tuple[int, ...]
es: List[numpy.ndarray] = dataclasses.field(default_factory=list)
hs: List[numpy.ndarray] = dataclasses.field(default_factory=list)
js: List[numpy.ndarray] = dataclasses.field(default_factory=list)
es: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
hs: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
js: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
@pytest.fixture(params=[(0, 4, 8)]) # (0,)
@ -164,10 +173,11 @@ def j_steps(request: FixtureRequest) -> Iterable[Tuple[int, ...]]:
@pytest.fixture(params=['center', 'random'])
def j_distribution(request: FixtureRequest,
def j_distribution(
request: FixtureRequest,
shape: Tuple[int, ...],
j_mag: float,
) -> Iterable[numpy.ndarray]:
) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape)
if request.param == 'center':
j[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = j_mag
@ -179,12 +189,13 @@ def j_distribution(request: FixtureRequest,
@pytest.fixture()
def sim(request: FixtureRequest,
def sim(
request: FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
dxes: List[List[numpy.ndarray]],
epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]],
dt: float,
j_distribution: numpy.ndarray,
j_distribution: NDArray[numpy.float64],
j_steps: Tuple[int, ...],
) -> TDResult:
is3d = (numpy.array(shape) == 1).sum() == 0

13
meanas/test/utils.py
View File

@ -1,12 +1,14 @@
from typing import Any
import numpy # type: ignore
import numpy
from typing import ArrayLike
PRNG = numpy.random.RandomState(12345)
def assert_fields_close(x: numpy.ndarray,
y: numpy.ndarray,
def assert_fields_close(
x: ArrayLike,
y: ArrayLike,
*args: Any,
**kwargs: Any,
) -> None:
@ -15,8 +17,9 @@ def assert_fields_close(x: numpy.ndarray,
err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1),
numpy.rollaxis(y, -1)), *args, **kwargs)
def assert_close(x: numpy.ndarray,
y: numpy.ndarray,
def assert_close(
x: ArrayLike,
y: ArrayLike,
*args: Any,
**kwargs: Any,
) -> None: