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

typing and formatting updates

Browse Source
master
Jan Petykiewicz 2 years ago
parent d42a625e5f
commit faecc79179
22 changed files with 621 additions and 486 deletions
Show Stats Download Patch File Download Diff File

46
meanas/eigensolvers.py
Unescape Escape View File

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

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

32
meanas/fdfd/farfield.py
Unescape Escape View File

@ -2,19 +2,20 @@
Functions for performing near-to-farfield transformation (and the reverse). Functions for performing near-to-farfield transformation (and the reverse).
""" """
from typing import Dict, List, Any from typing import Dict, List, Any
import numpy # type: ignore import numpy
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift # type: ignore from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
from numpy import pi # type: ignore from numpy import pi
from ..fdmath import fdfield_t from ..fdmath import fdfield_t
def near_to_farfield(E_near: fdfield_t, def near_to_farfield(
H_near: fdfield_t, E_near: fdfield_t,
dx: float, H_near: fdfield_t,
dy: float, dx: float,
padded_size: List[int] = None dy: float,
) -> Dict[str, Any]: padded_size: List[int] = None
) -> Dict[str, Any]:
""" """
Compute the farfield, i.e. the distribution of the fields after propagation Compute the farfield, i.e. the distribution of the fields after propagation
through several wavelengths of uniform medium. through several wavelengths of uniform medium.
@ -120,12 +121,13 @@ def near_to_farfield(E_near: fdfield_t,
return outputs return outputs
def far_to_nearfield(E_far: fdfield_t, def far_to_nearfield(
H_far: fdfield_t, E_far: fdfield_t,
dkx: float, H_far: fdfield_t,
dky: float, dkx: float,
padded_size: List[int] = None dky: float,
) -> Dict[str, Any]: padded_size: List[int] = None
) -> Dict[str, Any]:
""" """
Compute the farfield, i.e. the distribution of the fields after propagation Compute the farfield, i.e. the distribution of the fields after propagation
through several wavelengths of uniform medium. through several wavelengths of uniform medium.

49
meanas/fdfd/functional.py
Unescape Escape 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), 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) e.g. E = [E_x, E_y, E_z] where each component has shape (X, Y, Z)
""" """
from typing import Callable, Tuple from typing import Callable, Tuple, Optional
import numpy # type: ignore import numpy
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
from ..fdmath.functional import curl_forward, curl_back from ..fdmath.functional import curl_forward, curl_back
@ -15,11 +15,12 @@ from ..fdmath.functional import curl_forward, curl_back
__author__ = 'Jan Petykiewicz' __author__ = 'Jan Petykiewicz'
def e_full(omega: complex, def e_full(
dxes: dx_lists_t, omega: complex,
epsilon: fdfield_t, dxes: dx_lists_t,
mu: fdfield_t = None epsilon: fdfield_t,
) -> fdfield_updater_t: mu: fdfield_t = None
) -> fdfield_updater_t:
""" """
Wave operator for use with E-field. See `operators.e_full` for details. Wave operator for use with E-field. See `operators.e_full` for details.
@ -50,11 +51,12 @@ def e_full(omega: complex,
return op_mu return op_mu
def eh_full(omega: complex, def eh_full(
dxes: dx_lists_t, omega: complex,
epsilon: fdfield_t, dxes: dx_lists_t,
mu: fdfield_t = None epsilon: fdfield_t,
) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]: mu: fdfield_t = None
) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]:
""" """
Wave operator for full (both E and H) field representation. Wave operator for full (both E and H) field representation.
See `operators.eh_full`. See `operators.eh_full`.
@ -86,9 +88,10 @@ def eh_full(omega: complex,
return op_mu return op_mu
def e2h(omega: complex, def e2h(
omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t: ) -> fdfield_updater_t:
""" """
Utility operator for converting the `E` field into the `H` field. Utility operator for converting the `E` field into the `H` field.
@ -117,9 +120,10 @@ def e2h(omega: complex,
return e2h_mu return e2h_mu
def m2j(omega: complex, def m2j(
omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t: ) -> fdfield_updater_t:
""" """
Utility operator for converting magnetic current `M` distribution Utility operator for converting magnetic current `M` distribution
@ -151,12 +155,13 @@ def m2j(omega: complex,
return m2j_mu return m2j_mu
def e_tfsf_source(TF_region: fdfield_t, def e_tfsf_source(
omega: complex, TF_region: fdfield_t,
dxes: dx_lists_t, omega: complex,
epsilon: fdfield_t, dxes: dx_lists_t,
mu: fdfield_t = None, epsilon: fdfield_t,
) -> fdfield_updater_t: mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t:
""" """
Operator that turns an E-field distribution into a total-field/scattered-field Operator that turns an E-field distribution into a total-field/scattered-field
(TFSF) source. (TFSF) source.

88
meanas/fdfd/operators.py
Unescape Escape View File

@ -28,7 +28,7 @@ The following operators are included:
""" """
from typing import Tuple, Optional from typing import Tuple, Optional
import numpy # type: ignore import numpy
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
from ..fdmath import vec, dx_lists_t, vfdfield_t from ..fdmath import vec, dx_lists_t, vfdfield_t
@ -38,13 +38,14 @@ from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl
__author__ = 'Jan Petykiewicz' __author__ = 'Jan Petykiewicz'
def e_full(omega: complex, def e_full(
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
pec: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None, pec: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
""" """
Wave operator Wave operator
$$ \\nabla \\times (\\frac{1}{\\mu} \\nabla \\times) - \\Omega^2 \\epsilon $$ $$ \\nabla \\times (\\frac{1}{\\mu} \\nabla \\times) - \\Omega^2 \\epsilon $$
@ -96,8 +97,9 @@ def e_full(omega: complex,
return op return op
def e_full_preconditioners(dxes: dx_lists_t def e_full_preconditioners(
) -> Tuple[sparse.spmatrix, sparse.spmatrix]: dxes: dx_lists_t,
) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
""" """
Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator. Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator.
@ -122,13 +124,14 @@ def e_full_preconditioners(dxes: dx_lists_t
return P_left, P_right return P_left, P_right
def h_full(omega: complex, def h_full(
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
pec: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None, pec: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
""" """
Wave operator Wave operator
$$ \\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu $$ $$ \\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu $$
@ -178,13 +181,14 @@ def h_full(omega: complex,
return A return A
def eh_full(omega: complex, def eh_full(
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
pec: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None pec: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
""" """
Wave operator for `[E, H]` field representation. This operator implements Maxwell's Wave operator for `[E, H]` field representation. This operator implements Maxwell's
equations without cancelling out either E or H. The operator is equations without cancelling out either E or H. The operator is
@ -247,7 +251,8 @@ def eh_full(omega: complex,
return A return A
def e2h(omega: complex, def e2h(
omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None, pmc: Optional[vfdfield_t] = None,
@ -278,9 +283,10 @@ def e2h(omega: complex,
return op return op
def m2j(omega: complex, def m2j(
omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Operator for converting a magnetic current M into an electric current J. Operator for converting a magnetic current M into an electric current J.
@ -357,12 +363,13 @@ def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
return P return P
def e_tfsf_source(TF_region: vfdfield_t, def e_tfsf_source(
omega: complex, TF_region: vfdfield_t,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
) -> sparse.spmatrix: mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
""" """
Operator that turns a desired E-field distribution into a Operator that turns a desired E-field distribution into a
total-field/scattered-field (TFSF) source. total-field/scattered-field (TFSF) source.
@ -387,13 +394,14 @@ def e_tfsf_source(TF_region: vfdfield_t,
return (A @ Q - Q @ A) / (-1j * omega) return (A @ Q - Q @ A) / (-1j * omega)
def e_boundary_source(mask: vfdfield_t, def e_boundary_source(
omega: complex, mask: vfdfield_t,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
periodic_mask_edges: bool = False, mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: periodic_mask_edges: bool = False,
) -> sparse.spmatrix:
""" """
Operator that turns an E-field distrubtion into a current (J) distribution Operator that turns an E-field distrubtion into a current (J) distribution
along the edges (external and internal) of the provided mask. This is just an along the edges (external and internal) of the provided mask. This is just an

53
meanas/fdfd/scpml.py
Unescape Escape View File

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

47
meanas/fdfd/solvers.py
Unescape Escape View File

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

182
meanas/fdfd/waveguide_2d.py
Unescape Escape View File

@ -179,8 +179,9 @@ to account for numerical dispersion if the result is introduced into a space wit
# TODO update module docs # TODO update module docs
from typing import List, Tuple, Optional, Any from typing import List, Tuple, Optional, Any
import numpy # type: ignore import numpy
from numpy.linalg import norm # type: ignore from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
from ..fdmath.operators import deriv_forward, deriv_back, cross from ..fdmath.operators import deriv_forward, deriv_back, cross
@ -191,11 +192,12 @@ from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
__author__ = 'Jan Petykiewicz' __author__ = 'Jan Petykiewicz'
def operator_e(omega: complex, def operator_e(
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
) -> sparse.spmatrix: mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
""" """
Waveguide operator of the form Waveguide operator of the form
@ -257,11 +259,12 @@ def operator_e(omega: complex,
return op return op
def operator_h(omega: complex, def operator_h(
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
) -> sparse.spmatrix: mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix:
""" """
Waveguide operator of the form Waveguide operator of the form
@ -324,14 +327,15 @@ def operator_h(omega: complex,
return op return op
def normalized_fields_e(e_xy: numpy.ndarray, def normalized_fields_e(
wavenumber: complex, e_xy: ArrayLike,
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
prop_phase: float = 0, mu: Optional[vfdfield_t] = None,
) -> Tuple[vfdfield_t, vfdfield_t]: prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]:
""" """
Given a vector `e_xy` containing the vectorized E_x and E_y fields, Given a vector `e_xy` containing the vectorized E_x and E_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -358,14 +362,15 @@ def normalized_fields_e(e_xy: numpy.ndarray,
return e_norm, h_norm return e_norm, h_norm
def normalized_fields_h(h_xy: numpy.ndarray, def normalized_fields_h(
wavenumber: complex, h_xy: ArrayLike,
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
prop_phase: float = 0, mu: Optional[vfdfield_t] = None,
) -> Tuple[vfdfield_t, vfdfield_t]: prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]:
""" """
Given a vector `h_xy` containing the vectorized H_x and H_y fields, Given a vector `h_xy` containing the vectorized H_x and H_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -392,14 +397,15 @@ def normalized_fields_h(h_xy: numpy.ndarray,
return e_norm, h_norm return e_norm, h_norm
def _normalized_fields(e: numpy.ndarray, def _normalized_fields(
h: numpy.ndarray, e: ArrayLike,
omega: complex, h: ArrayLike,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None, epsilon: vfdfield_t,
prop_phase: float = 0, mu: Optional[vfdfield_t] = None,
) -> Tuple[vfdfield_t, vfdfield_t]: prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]:
# TODO documentation # TODO documentation
shape = [s.size for s in dxes[0]] shape = [s.size for s in dxes[0]]
dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)] dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)]
@ -434,12 +440,13 @@ def _normalized_fields(e: numpy.ndarray,
return e, h return e, h
def exy2h(wavenumber: complex, def exy2h(
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None epsilon: vfdfield_t,
) -> sparse.spmatrix: mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix:
""" """
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
into a vectorized H containing all three H components into a vectorized H containing all three H components
@ -459,12 +466,13 @@ def exy2h(wavenumber: complex,
return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon) return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)
def hxy2e(wavenumber: complex, def hxy2e(
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None epsilon: vfdfield_t,
) -> sparse.spmatrix: mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix:
""" """
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
into a vectorized E containing all three E components into a vectorized E containing all three E components
@ -484,10 +492,11 @@ def hxy2e(wavenumber: complex,
return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu) return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)
def hxy2h(wavenumber: complex, def hxy2h(
dxes: dx_lists_t, wavenumber: complex,
mu: Optional[vfdfield_t] = None dxes: dx_lists_t,
) -> sparse.spmatrix: mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix:
""" """
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
into a vectorized H containing all three H components into a vectorized H containing all three H components
@ -517,10 +526,11 @@ def hxy2h(wavenumber: complex,
return op return op
def exy2e(wavenumber: complex, def exy2e(
dxes: dx_lists_t, wavenumber: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
) -> sparse.spmatrix: epsilon: vfdfield_t,
) -> sparse.spmatrix:
""" """
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
into a vectorized E containing all three E components into a vectorized E containing all three E components
@ -550,7 +560,8 @@ def exy2e(wavenumber: complex,
return op return op
def e2h(wavenumber: complex, def e2h(
wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None mu: Optional[vfdfield_t] = None
@ -574,7 +585,8 @@ def e2h(wavenumber: complex,
return op return op
def h2e(wavenumber: complex, def h2e(
wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t epsilon: vfdfield_t
@ -636,13 +648,14 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
return cross([Dbx, Dby, Bz]) return cross([Dbx, Dby, Bz])
def h_err(h: vfdfield_t, def h_err(
wavenumber: complex, h: vfdfield_t,
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: vfdfield_t = None epsilon: vfdfield_t,
) -> float: mu: Optional[vfdfield_t] = None
) -> float:
""" """
Calculates the relative error in the H field Calculates the relative error in the H field
@ -670,13 +683,14 @@ def h_err(h: vfdfield_t,
return norm(op) / norm(h) return norm(op) / norm(h)
def e_err(e: vfdfield_t, def e_err(
wavenumber: complex, e: vfdfield_t,
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: vfdfield_t = None epsilon: vfdfield_t,
) -> float: mu: vfdfield_t = Optional[None]
) -> float:
""" """
Calculates the relative error in the E field Calculates the relative error in the E field
@ -703,13 +717,14 @@ def e_err(e: vfdfield_t,
return norm(op) / norm(e) return norm(op) / norm(e)
def solve_modes(mode_numbers: List[int], def solve_modes(
omega: complex, mode_numbers: List[int],
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
mu: vfdfield_t = None, epsilon: vfdfield_t,
mode_margin: int = 2, mu: vfdfield_t = None,
) -> Tuple[numpy.ndarray, List[complex]]: mode_margin: int = 2,
) -> Tuple[NDArray[numpy.float64], List[complex]]:
""" """
Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers. Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
@ -752,10 +767,11 @@ def solve_modes(mode_numbers: List[int],
return e_xys, wavenumbers return e_xys, wavenumbers
def solve_mode(mode_number: int, def solve_mode(
*args: Any, mode_number: int,
**kwargs: Any, *args: Any,
) -> Tuple[vfdfield_t, complex]: **kwargs: Any,
) -> Tuple[vfdfield_t, complex]:
""" """
Wrapper around `solve_modes()` that solves for a single mode. Wrapper around `solve_modes()` that solves for a single mode.

83
meanas/fdfd/waveguide_3d.py
Unescape Escape View File

@ -5,21 +5,23 @@ This module relies heavily on `waveguide_2d` and mostly just transforms
its parameters into 2D equivalents and expands the results back into 3D. its parameters into 2D equivalents and expands the results back into 3D.
""" """
from typing import Dict, Optional, Sequence, Union, Any 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 ..fdmath import vec, unvec, dx_lists_t, fdfield_t
from . import operators, waveguide_2d from . import operators, waveguide_2d
def solve_mode(mode_number: int, def solve_mode(
omega: complex, mode_number: int,
dxes: dx_lists_t, omega: complex,
axis: int, dxes: dx_lists_t,
polarity: int, axis: int,
slices: Sequence[slice], polarity: int,
epsilon: fdfield_t, slices: Sequence[slice],
mu: Optional[fdfield_t] = None, epsilon: fdfield_t,
) -> Dict[str, Union[complex, numpy.ndarray]]: mu: Optional[fdfield_t] = None,
) -> Dict[str, Union[complex, NDArray[numpy.float_]]]:
""" """
Given a 3D grid, selects a slice from the grid and attempts to Given a 3D grid, selects a slice from the grid and attempts to
solve for an eigenmode propagating through that slice. solve for an eigenmode propagating through that slice.
@ -36,7 +38,13 @@ def solve_mode(mode_number: int,
mu: Magnetic permeability (default 1 everywhere) mu: Magnetic permeability (default 1 everywhere)
Returns: 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: if mu is None:
mu = numpy.ones_like(epsilon) mu = numpy.ones_like(epsilon)
@ -97,16 +105,17 @@ def solve_mode(mode_number: int,
return results return results
def compute_source(E: fdfield_t, def compute_source(
wavenumber: complex, E: fdfield_t,
omega: complex, wavenumber: complex,
dxes: dx_lists_t, omega: complex,
axis: int, dxes: dx_lists_t,
polarity: int, axis: int,
slices: Sequence[slice], polarity: int,
epsilon: fdfield_t, slices: Sequence[slice],
mu: Optional[fdfield_t] = None, epsilon: fdfield_t,
) -> fdfield_t: mu: Optional[fdfield_t] = None,
) -> fdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, returns the current source distribution Given an eigenmode obtained by `solve_mode`, returns the current source distribution
necessary to position a unidirectional source at the slice location. necessary to position a unidirectional source at the slice location.
@ -142,18 +151,21 @@ def compute_source(E: fdfield_t,
return J return J
def compute_overlap_e(E: fdfield_t, def compute_overlap_e(
wavenumber: complex, E: fdfield_t,
dxes: dx_lists_t, wavenumber: complex,
axis: int, dxes: dx_lists_t,
polarity: int, axis: int,
slices: Sequence[slice], polarity: int,
) -> fdfield_t: # TODO DOCS slices: Sequence[slice],
) -> fdfield_t: # TODO DOCS
""" """
Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the
mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn) mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn)
[assumes reflection symmetry]. [assumes reflection symmetry].
TODO: add reference
Args: Args:
E: E-field of the mode E: E-field of the mode
H: H-field of the mode (advanced by half of a Yee cell from E) H: H-field of the mode (advanced by half of a Yee cell from E)
@ -187,13 +199,14 @@ def compute_overlap_e(E: fdfield_t,
return Etgt return Etgt
def expand_e(E: fdfield_t, def expand_e(
wavenumber: complex, E: fdfield_t,
dxes: dx_lists_t, wavenumber: complex,
axis: int, dxes: dx_lists_t,
polarity: int, axis: int,
slices: Sequence[slice], polarity: int,
) -> fdfield_t: slices: Sequence[slice],
) -> fdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D
slice where the mode was calculated to the entire domain (along the propagation slice where the mode was calculated to the entire domain (along the propagation

34
meanas/fdfd/waveguide_cyl.py
Unescape Escape 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 # TODO update module docs
from typing import Dict, Union from typing import Dict, Union
import numpy # type: ignore import numpy
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t
@ -17,11 +17,12 @@ from ..fdmath.operators import deriv_forward, deriv_back
from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
def cylindrical_operator(omega: complex, def cylindrical_operator(
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
r0: float, epsilon: vfdfield_t,
) -> sparse.spmatrix: r0: float,
) -> sparse.spmatrix:
""" """
Cylindrical coordinate waveguide operator of the form Cylindrical coordinate waveguide operator of the form
@ -78,12 +79,13 @@ def cylindrical_operator(omega: complex,
return op return op
def solve_mode(mode_number: int, def solve_mode(
omega: complex, mode_number: int,
dxes: dx_lists_t, omega: complex,
epsilon: vfdfield_t, dxes: dx_lists_t,
r0: float, epsilon: vfdfield_t,
) -> Dict[str, Union[complex, fdfield_t]]: r0: float,
) -> Dict[str, Union[complex, fdfield_t]]:
""" """
TODO: fixup TODO: fixup
Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@ -99,7 +101,13 @@ def solve_mode(mode_number: int,
r within the simulation domain. r within the simulation domain.
Returns: Returns:
`{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}` ```
{
'E': List[NDArray[numpy.float_]],
'H': List[NDArray[numpy.float_]],
'wavenumber': complex,
}
```
""" """
''' '''

29
meanas/fdmath/functional.py
Unescape Escape View File

@ -5,13 +5,15 @@ Basic discrete calculus etc.
""" """
from typing import Sequence, Tuple, Optional, Callable 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 from .types import fdfield_t, fdfield_updater_t
def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None def deriv_forward(
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: 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). Utility operators for taking discretized derivatives (backward variant).
@ -33,8 +35,9 @@ def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None
return derivs return derivs
def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None def deriv_back(
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: 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). Utility operators for taking discretized derivatives (forward variant).
@ -56,7 +59,9 @@ def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None
return derivs 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. 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 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. 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 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) Dx, Dy, Dz = deriv_forward(dx_e)
def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]: 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 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) Dx, Dy, Dz = deriv_back(dx_e)
def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]: def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]:

37
meanas/fdmath/operators.py
Unescape Escape View File

@ -4,13 +4,18 @@ Matrix operators for finite difference simulations
Basic discrete calculus etc. Basic discrete calculus etc.
""" """
from typing import Sequence, List from typing import Sequence, List
import numpy # type: ignore import numpy
from numpy.typing import NDArray
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
from .types import vfdfield_t 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 Utility operator for performing a circular shift along a specified axis by a
specified number of elements. specified number of elements.
@ -46,7 +51,11 @@ def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> spar
return d 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 Utility operator for performing an n-element shift along a specified axis, with mirror
boundary conditions applied to the cells beyond the receding edge. 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( raise Exception('Shift ({}) is too large for axis {} of size {}'.format(
shift_distance, axis, shape[axis])) 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.arange(n) + s
v = numpy.where(v >= n, 2 * n - v - 1, v) v = numpy.where(v >= n, 2 * n - v - 1, v)
v = numpy.where(v < 0, - 1 - v, 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 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). 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 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). 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 return Ds
def cross(B: Sequence[sparse.spmatrix]) -> sparse.spmatrix: def cross(
B: Sequence[sparse.spmatrix],
) -> sparse.spmatrix:
""" """
Cross product operator Cross product operator
@ -203,7 +218,9 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
return avg_forward(axis, shape).T 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. 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)) 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. Curl operator for use with the H field.

39
meanas/fdmath/types.py
Unescape Escape View File

@ -2,31 +2,20 @@
Types shared across multiple submodules Types shared across multiple submodules
""" """
from typing import Sequence, Callable, MutableSequence from typing import Sequence, Callable, MutableSequence
import numpy # type: ignore import numpy
from numpy.typing import NDArray
# Field types # Field types
# TODO: figure out a better way to set the docstrings without creating actual subclasses? fdfield_t = NDArray[numpy.float_]
# Probably not a big issue since they're only used for type hinting """Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
class fdfield_t(numpy.ndarray):
"""
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` vfdfield_t = NDArray[numpy.float_]
""" """Linearized vector field (single vector of length 3*X*Y*Z)"""
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]] dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]]
''' """
'dxes' datastructure which contains grid cell width information in the following format: 'dxes' datastructure which contains grid cell width information in the following format:
[[[dx_e[0], dx_e[1], ...], [dy_e[0], ...], [dz_e[0], ...]], [[[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, 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. 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]] dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]]
''' """Mutable version of `dx_lists_t`"""
Mutable version of `dx_lists_t`
'''
fdfield_updater_t = Callable[..., fdfield_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
Unescape Escape 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 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 from .types import fdfield_t, vfdfield_t
@ -15,10 +16,10 @@ def vec(f: None) -> None:
pass pass
@overload @overload
def vec(f: Union[fdfield_t, List[numpy.ndarray]]) -> vfdfield_t: def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t:
pass 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. 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 @overload
def unvec(v: None, shape: numpy.ndarray) -> None: def unvec(v: None, shape: ArrayLike) -> None:
pass pass
@overload @overload
def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t: def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t:
pass 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 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 of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional

7
meanas/fdtd/boundaries.py
Unescape Escape View File

@ -9,9 +9,10 @@ from typing import Tuple, Any, List
from ..fdmath import fdfield_t, fdfield_updater_t from ..fdmath import fdfield_t, fdfield_updater_t
def conducting_boundary(direction: int, def conducting_boundary(
polarity: int direction: int,
) -> Tuple[fdfield_updater_t, fdfield_updater_t]: polarity: int
) -> Tuple[fdfield_updater_t, fdfield_updater_t]:
dirs = [0, 1, 2] dirs = [0, 1, 2]
if direction not in dirs: if direction not in dirs:
raise Exception('Invalid direction: {}'.format(direction)) raise Exception('Invalid direction: {}'.format(direction))

114
meanas/fdtd/energy.py
Unescape Escape View File

@ -1,5 +1,5 @@
from typing import Optional, Union from typing import Optional, Union
import numpy # type: ignore import numpy
from ..fdmath import dx_lists_t, fdfield_t from ..fdmath import dx_lists_t, fdfield_t
from ..fdmath.functional import deriv_back from ..fdmath.functional import deriv_back
@ -8,10 +8,11 @@ from ..fdmath.functional import deriv_back
# TODO documentation # TODO documentation
def poynting(e: fdfield_t, def poynting(
h: fdfield_t, e: fdfield_t,
dxes: Optional[dx_lists_t] = None, h: fdfield_t,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Calculate the poynting vector `S` ($S$). Calculate the poynting vector `S` ($S$).
@ -87,12 +88,13 @@ def poynting(e: fdfield_t,
return s 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, e: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None, h: Optional[fdfield_t] = None,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Calculate the divergence of the poynting vector. Calculate the divergence of the poynting vector.
@ -124,13 +126,14 @@ def poynting_divergence(s: Optional[fdfield_t] = None,
return ds return ds
def energy_hstep(e0: fdfield_t, def energy_hstep(
h1: fdfield_t, e0: fdfield_t,
e2: fdfield_t, h1: fdfield_t,
epsilon: Optional[fdfield_t] = None, e2: fdfield_t,
mu: Optional[fdfield_t] = None, epsilon: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Calculate energy `U` at the time of the provided H-field `h1`. Calculate energy `U` at the time of the provided H-field `h1`.
@ -151,13 +154,14 @@ def energy_hstep(e0: fdfield_t,
return u return u
def energy_estep(h0: fdfield_t, def energy_estep(
e1: fdfield_t, h0: fdfield_t,
h2: fdfield_t, e1: fdfield_t,
epsilon: Optional[fdfield_t] = None, h2: fdfield_t,
mu: Optional[fdfield_t] = None, epsilon: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Calculate energy `U` at the time of the provided E-field `e1`. Calculate energy `U` at the time of the provided E-field `e1`.
@ -178,15 +182,16 @@ def energy_estep(h0: fdfield_t,
return u return u
def delta_energy_h2e(dt: float, def delta_energy_h2e(
e0: fdfield_t, dt: float,
h1: fdfield_t, e0: fdfield_t,
e2: fdfield_t, h1: fdfield_t,
h3: fdfield_t, e2: fdfield_t,
epsilon: Optional[fdfield_t] = None, h3: fdfield_t,
mu: Optional[fdfield_t] = None, epsilon: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Change in energy during the half-step from `h1` to `e2`. Change in energy during the half-step from `h1` to `e2`.
@ -210,15 +215,16 @@ def delta_energy_h2e(dt: float,
return du return du
def delta_energy_e2h(dt: float, def delta_energy_e2h(
h0: fdfield_t, dt: float,
e1: fdfield_t, h0: fdfield_t,
h2: fdfield_t, e1: fdfield_t,
e3: fdfield_t, h2: fdfield_t,
epsilon: Optional[fdfield_t] = None, e3: fdfield_t,
mu: Optional[fdfield_t] = None, epsilon: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Change in energy during the half-step from `e1` to `h2`. Change in energy during the half-step from `e1` to `h2`.
@ -242,10 +248,11 @@ def delta_energy_e2h(dt: float,
return du return du
def delta_energy_j(j0: fdfield_t, def delta_energy_j(
e1: fdfield_t, j0: fdfield_t,
dxes: Optional[dx_lists_t] = None, e1: fdfield_t,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
""" """
Calculate Calculate
@ -264,12 +271,13 @@ def delta_energy_j(j0: fdfield_t,
return du return du
def dxmul(ee: fdfield_t, def dxmul(
hh: fdfield_t, ee: fdfield_t,
epsilon: Optional[Union[fdfield_t, float]] = None, hh: fdfield_t,
mu: Optional[Union[fdfield_t, float]] = None, epsilon: Optional[Union[fdfield_t, float]] = None,
dxes: Optional[dx_lists_t] = None mu: Optional[Union[fdfield_t, float]] = None,
) -> fdfield_t: dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
if epsilon is None: if epsilon is None:
epsilon = 1 epsilon = 1
if mu is None: if mu is None:

5
meanas/fdtd/pml.py
Unescape Escape View File

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

23
meanas/test/conftest.py
Unescape Escape View File

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

41
meanas/test/test_fdfd.py
Unescape Escape View File

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

79
meanas/test/test_fdfd_pml.py
Unescape Escape View File

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

61
meanas/test/test_fdtd.py
Unescape Escape View File

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

25
meanas/test/utils.py
Unescape Escape View File

@ -1,24 +1,27 @@
from typing import Any from typing import Any
import numpy # type: ignore import numpy
from typing import ArrayLike
PRNG = numpy.random.RandomState(12345) PRNG = numpy.random.RandomState(12345)
def assert_fields_close(x: numpy.ndarray, def assert_fields_close(
y: numpy.ndarray, x: ArrayLike,
*args: Any, y: ArrayLike,
**kwargs: Any, *args: Any,
) -> None: **kwargs: Any,
) -> None:
numpy.testing.assert_allclose( numpy.testing.assert_allclose(
x, y, verbose=False, x, y, verbose=False,
err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1), err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1),
numpy.rollaxis(y, -1)), *args, **kwargs) numpy.rollaxis(y, -1)), *args, **kwargs)
def assert_close(x: numpy.ndarray, def assert_close(
y: numpy.ndarray, x: ArrayLike,
*args: Any, y: ArrayLike,
**kwargs: Any, *args: Any,
) -> None: **kwargs: Any,
) -> None:
numpy.testing.assert_allclose(x, y, *args, **kwargs) numpy.testing.assert_allclose(x, y, *args, **kwargs)

Write Preview
Loading…
Cancel
Save