add more type hints

This commit is contained in:
Jan Petykiewicz 2020-10-16 21:46:04 -07:00
parent d13a3796a9
commit bc428f5e8e
17 changed files with 170 additions and 106 deletions

View File

@ -73,8 +73,9 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe
dtype=operator.dtype, dtype=operator.dtype,
matvec=lambda v: eigval * v) matvec=lambda v: eigval * v)
if solver is None: if solver is None:
def solver(A, b): def solver(A: spalg.LinearOperator, b: numpy.ndarray) -> numpy.ndarray:
return spalg.bicgstab(A, b)[0] return spalg.bicgstab(A, b)[0]
assert(solver is not None)
v = numpy.squeeze(guess_vector) v = numpy.squeeze(guess_vector)
v /= norm(v) v /= norm(v)

View File

@ -80,7 +80,7 @@ This module contains functions for generating and solving the
''' '''
from typing import Tuple, Callable from typing import Tuple, Callable, Any, List, Optional, cast
import logging import logging
import numpy # type: ignore import numpy # type: ignore
from numpy import pi, real, trace # type: ignore from numpy import pi, real, trace # type: ignore
@ -109,10 +109,10 @@ try:
'planner_effort': 'FFTW_EXHAUSTIVE', 'planner_effort': 'FFTW_EXHAUSTIVE',
} }
def fftn(*args, **kwargs): def fftn(*args: Any, **kwargs: Any) -> numpy.ndarray:
return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args) return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args)
def ifftn(*args, **kwargs): def ifftn(*args: Any, **kwargs: Any) -> numpy.ndarray:
return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args) return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)
except ImportError: except ImportError:
@ -199,7 +199,7 @@ def maxwell_operator(k0: numpy.ndarray,
if mu is not None: if mu is not None:
mu = numpy.stack(mu, 3) mu = numpy.stack(mu, 3)
def operator(h: numpy.ndarray): def operator(h: numpy.ndarray) -> numpy.ndarray:
""" """
Maxwell operator for Bloch eigenmode simulation. Maxwell operator for Bloch eigenmode simulation.
@ -309,11 +309,11 @@ def hmn_2_hxyz(k0: numpy.ndarray,
shape = epsilon[0].shape + (1,) shape = epsilon[0].shape + (1,)
_k_mag, m, n = generate_kmn(k0, G_matrix, shape) _k_mag, m, n = generate_kmn(k0, G_matrix, shape)
def operator(h: numpy.ndarray): def operator(h: numpy.ndarray) -> fdfield_t:
hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)] hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
h_xyz = (m * hin_m h_xyz = (m * hin_m
+ n * hin_n) + n * hin_n)
return [ifftn(hi) for hi in numpy.rollaxis(h_xyz, 3)] return numpy.array([ifftn(hi) for hi in numpy.rollaxis(h_xyz, 3)])
return operator return operator
@ -351,7 +351,7 @@ def inverse_maxwell_operator_approx(k0: numpy.ndarray,
if mu is not None: if mu is not None:
mu = numpy.stack(mu, 3) mu = numpy.stack(mu, 3)
def operator(h: numpy.ndarray): def operator(h: numpy.ndarray) -> numpy.ndarray:
""" """
Approximate inverse Maxwell operator for Bloch eigenmode simulation. Approximate inverse Maxwell operator for Bloch eigenmode simulation.
@ -429,7 +429,7 @@ def find_k(frequency: float,
direction = numpy.array(direction) / norm(direction) direction = numpy.array(direction) / norm(direction)
def get_f(k0_mag: float, band: int = 0): def get_f(k0_mag: float, band: int = 0) -> numpy.ndarray:
k0 = direction * k0_mag k0 = direction * k0_mag
n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu) n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
f = numpy.sqrt(numpy.abs(numpy.real(n[band]))) f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
@ -552,12 +552,12 @@ def eigsolve(num_modes: int,
symZtD = _symmetrize(Z.conj().T @ D) symZtD = _symmetrize(Z.conj().T @ D)
symZtAD = _symmetrize(Z.conj().T @ AD) symZtAD = _symmetrize(Z.conj().T @ AD)
Qi_memo = [None, None] Qi_memo: List[Optional[float]] = [None, None]
def Qi_func(theta): def Qi_func(theta: float) -> float:
nonlocal Qi_memo nonlocal Qi_memo
if Qi_memo[0] == theta: if Qi_memo[0] == theta:
return Qi_memo[1] return cast(float, Qi_memo[1])
c = numpy.cos(theta) c = numpy.cos(theta)
s = numpy.sin(theta) s = numpy.sin(theta)
@ -579,7 +579,7 @@ def eigsolve(num_modes: int,
Qi_memo[1] = Qi Qi_memo[1] = Qi
return Qi return Qi
def trace_func(theta): def trace_func(theta: float) -> float:
c = numpy.cos(theta) c = numpy.cos(theta)
s = numpy.sin(theta) s = numpy.sin(theta)
Qi = Qi_func(theta) Qi = Qi_func(theta)
@ -685,9 +685,9 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to
return x, fx, dfx return x, fx, dfx
''' '''
def _rtrace_AtB(A, B): def _rtrace_AtB(A: numpy.ndarray, B: numpy.ndarray) -> numpy.ndarray:
return real(numpy.sum(A.conj() * B)) return real(numpy.sum(A.conj() * B))
def _symmetrize(A): def _symmetrize(A: numpy.ndarray) -> numpy.ndarray:
return (A + A.conj().T) * 0.5 return (A + A.conj().T) * 0.5

View File

@ -36,13 +36,13 @@ def e_full(omega: complex,
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def op_1(e): def op_1(e: fdfield_t) -> fdfield_t:
curls = ch(ce(e)) curls = ch(ce(e))
return curls - omega ** 2 * epsilon * e return curls - omega ** 2 * epsilon * e # type: ignore # issues with numpy/mypy
def op_mu(e): def op_mu(e: fdfield_t) -> fdfield_t:
curls = ch(mu * ce(e)) curls = ch(mu * ce(e))
return curls - omega ** 2 * epsilon * e return curls - omega ** 2 * epsilon * e # type: ignore # issues with numpy/mypy
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return op_1 return op_1
@ -72,13 +72,13 @@ def eh_full(omega: complex,
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def op_1(e, h): def op_1(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]:
return (ch(h) - 1j * omega * epsilon * e, return (ch(h) - 1j * omega * epsilon * e,
ce(e) + 1j * omega * h) ce(e) + 1j * omega * h) # type: ignore # issues with numpy/mypy
def op_mu(e, h): def op_mu(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]:
return (ch(h) - 1j * omega * epsilon * e, return (ch(h) - 1j * omega * epsilon * e,
ce(e) + 1j * omega * mu * h) ce(e) + 1j * omega * mu * h) # type: ignore # issues with numpy/mypy
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return op_1 return op_1
@ -105,11 +105,11 @@ def e2h(omega: complex,
""" """
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def e2h_1_1(e): def e2h_1_1(e: fdfield_t) -> fdfield_t:
return ce(e) / (-1j * omega) return ce(e) / (-1j * omega) # type: ignore # issues with numpy/mypy
def e2h_mu(e): def e2h_mu(e: fdfield_t) -> fdfield_t:
return ce(e) / (-1j * omega * mu) return ce(e) / (-1j * omega * mu) # type: ignore # issues with numpy/mypy
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return e2h_1_1 return e2h_1_1
@ -137,13 +137,13 @@ def m2j(omega: complex,
""" """
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
def m2j_mu(m): def m2j_mu(m: fdfield_t) -> fdfield_t:
J = ch(m / mu) / (-1j * omega) J = ch(m / mu) / (-1j * omega)
return J return J # type: ignore # issues with numpy/mypy
def m2j_1(m): def m2j_1(m: fdfield_t) -> fdfield_t:
J = ch(m) / (-1j * omega) J = ch(m) / (-1j * omega)
return J return J # type: ignore # issues with numpy/mypy
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return m2j_1 return m2j_1
@ -177,7 +177,7 @@ def e_tfsf_source(TF_region: fdfield_t,
# TODO documentation # TODO documentation
A = e_full(omega, dxes, epsilon, mu) A = e_full(omega, dxes, epsilon, mu)
def op(e): def op(e: fdfield_t) -> fdfield_t:
neg_iwj = A(TF_region * e) - TF_region * A(e) neg_iwj = A(TF_region * e) - TF_region * A(e)
return neg_iwj / (-1j * omega) return neg_iwj / (-1j * omega)
return op return op
@ -205,7 +205,7 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdf
Returns: Returns:
Function `f` that returns E x H as required for the poynting vector. Function `f` that returns E x H as required for the poynting vector.
""" """
def exh(e: fdfield_t, h: fdfield_t): def exh(e: fdfield_t, h: fdfield_t) -> fdfield_t:
s = numpy.empty_like(e) s = numpy.empty_like(e)
ex = e[0] * dxes[0][0][:, None, None] ex = e[0] * dxes[0][0][:, None, None]
ey = e[1] * dxes[0][1][None, :, None] ey = e[1] * dxes[0][1][None, :, None]

View File

@ -416,13 +416,14 @@ def e_boundary_source(mask: vfdfield_t,
shape = [len(dxe) for dxe in dxes[0]] shape = [len(dxe) for dxe in dxes[0]]
jmask = numpy.zeros_like(mask, dtype=bool) jmask = numpy.zeros_like(mask, dtype=bool)
if periodic_mask_edges: def shift_rot(axis: int, polarity: int) -> sparse.spmatrix:
def shift(axis, polarity):
return rotation(axis=axis, shape=shape, shift_distance=polarity) return rotation(axis=axis, shape=shape, shift_distance=polarity)
else:
def shift(axis, polarity): def shift_mir(axis: int, polarity: int) -> sparse.spmatrix:
return shift_with_mirror(axis=axis, shape=shape, shift_distance=polarity) return shift_with_mirror(axis=axis, shape=shape, shift_distance=polarity)
shift = shift_rot if periodic_mask_edges else shift_mir
for axis in (0, 1, 2): for axis in (0, 1, 2):
if shape[axis] == 1: if shape[axis] == 1:
continue continue

View File

@ -2,11 +2,9 @@
Functions for creating stretched coordinate perfectly matched layer (PML) absorbers. Functions for creating stretched coordinate perfectly matched layer (PML) absorbers.
""" """
from typing import Sequence, Union, Callable, Optional from typing import Sequence, Union, Callable, Optional, List
import numpy # type: ignore import numpy # type: ignore
from ..fdmath import dx_lists_t, dx_lists_mut
__author__ = 'Jan Petykiewicz' __author__ = 'Jan Petykiewicz'
@ -42,7 +40,7 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
omega: float, omega: float,
epsilon_effective: float = 1.0, epsilon_effective: float = 1.0,
s_function: Optional[s_function_t] = None, s_function: Optional[s_function_t] = None,
) -> dx_lists_mut: ) -> List[List[numpy.ndarray]]:
""" """
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.
@ -69,7 +67,7 @@ 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, n, t): def ll(u: numpy.ndarray, n: numpy.ndarray, t: numpy.ndarray) -> numpy.ndarray:
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
@ -90,14 +88,14 @@ 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: dx_lists_mut, def stretch_with_scpml(dxes: List[List[numpy.ndarray]],
axis: int, axis: int,
polarity: int, polarity: int,
omega: float, omega: float,
epsilon_effective: float = 1.0, epsilon_effective: float = 1.0,
thickness: int = 10, thickness: int = 10,
s_function: Optional[s_function_t] = None, s_function: Optional[s_function_t] = None,
) -> dx_lists_t: ) -> List[List[numpy.ndarray]]:
""" """
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.
@ -134,7 +132,7 @@ def stretch_with_scpml(dxes: dx_lists_mut,
bound = pos[thickness] bound = pos[thickness]
d = bound - pos[0] d = bound - pos[0]
def l_d(x): def l_d(x: numpy.ndarray) -> numpy.ndarray:
return (bound - x) / (bound - pos[0]) return (bound - x) / (bound - pos[0])
slc = slice(thickness) slc = slice(thickness)
@ -144,7 +142,7 @@ def stretch_with_scpml(dxes: dx_lists_mut,
bound = pos[-thickness - 1] bound = pos[-thickness - 1]
d = pos[-1] - bound d = pos[-1] - bound
def l_d(x): def l_d(x: numpy.ndarray) -> numpy.ndarray:
return (x - bound) / (pos[-1] - bound) return (x - bound) / (pos[-1] - bound)
if thickness == 0: if thickness == 0:

View File

@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)
def _scipy_qmr(A: scipy.sparse.csr_matrix, def _scipy_qmr(A: scipy.sparse.csr_matrix,
b: numpy.ndarray, b: numpy.ndarray,
**kwargs **kwargs: Any,
) -> numpy.ndarray: ) -> numpy.ndarray:
""" """
Wrapper for scipy.sparse.linalg.qmr Wrapper for scipy.sparse.linalg.qmr
@ -37,14 +37,14 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix,
''' '''
ii = 0 ii = 0
def log_residual(xk): def log_residual(xk: numpy.ndarray) -> 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): def augmented_callback(xk: numpy.ndarray) -> None:
log_residual(xk) log_residual(xk)
kwargs['callback'](xk) kwargs['callback'](xk)

View File

@ -146,7 +146,7 @@ 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 from typing import List, Tuple, Optional, Any
import numpy # type: ignore import numpy # type: ignore
from numpy.linalg import norm # type: ignore from numpy.linalg import norm # type: ignore
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
@ -721,8 +721,8 @@ def solve_modes(mode_numbers: List[int],
def solve_mode(mode_number: int, def solve_mode(mode_number: int,
*args, *args: Any,
**kwargs **kwargs: Any,
) -> Tuple[vfdfield_t, complex]: ) -> Tuple[vfdfield_t, complex]:
""" """
Wrapper around `solve_modes()` that solves for a single mode. Wrapper around `solve_modes()` that solves for a single mode.

View File

@ -67,7 +67,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, s): def mirrored_range(n: int, s: int) -> numpy.ndarray:
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)
@ -103,7 +103,7 @@ def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij') dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij')
def deriv(axis): def deriv(axis: int) -> sparse.spmatrix:
return rotation(axis, shape, 1) - sparse.eye(n) return rotation(axis, shape, 1) - sparse.eye(n)
Ds = [sparse.diags(+1 / dx.ravel(order='C')) @ deriv(a) Ds = [sparse.diags(+1 / dx.ravel(order='C')) @ deriv(a)
@ -128,7 +128,7 @@ def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij') dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij')
def deriv(axis): def deriv(axis: int) -> sparse.spmatrix:
return rotation(axis, shape, -1) - sparse.eye(n) return rotation(axis, shape, -1) - sparse.eye(n)
Ds = [sparse.diags(-1 / dx.ravel(order='C')) @ deriv(a) Ds = [sparse.diags(-1 / dx.ravel(order='C')) @ deriv(a)

View File

@ -130,7 +130,7 @@ $$
\\end{aligned} \\end{aligned}
$$ $$
This result is exact an should practically hold to within numerical precision. No time- This result is exact and should practically hold to within numerical precision. No time-
or spatial-averaging is necessary. or spatial-averaging is necessary.
Note that each value of $J$ contributes to the energy twice (i.e. once per field update) Note that each value of $J$ contributes to the energy twice (i.e. once per field update)

View File

@ -24,13 +24,13 @@ def conducting_boundary(direction: int,
boundary_slice[direction] = 0 boundary_slice[direction] = 0
shifted1_slice[direction] = 1 shifted1_slice[direction] = 1
def en(e: fdfield_t): def en(e: fdfield_t) -> fdfield_t:
e[direction][boundary_slice] = 0 e[direction][boundary_slice] = 0
e[u][boundary_slice] = e[u][shifted1_slice] e[u][boundary_slice] = e[u][shifted1_slice]
e[v][boundary_slice] = e[v][shifted1_slice] e[v][boundary_slice] = e[v][shifted1_slice]
return e return e
def hn(h: fdfield_t): def hn(h: fdfield_t) -> fdfield_t:
h[direction][boundary_slice] = h[direction][shifted1_slice] h[direction][boundary_slice] = h[direction][shifted1_slice]
h[u][boundary_slice] = 0 h[u][boundary_slice] = 0
h[v][boundary_slice] = 0 h[v][boundary_slice] = 0
@ -46,14 +46,14 @@ def conducting_boundary(direction: int,
shifted1_slice[direction] = -2 shifted1_slice[direction] = -2
shifted2_slice[direction] = -3 shifted2_slice[direction] = -3
def ep(e: fdfield_t): def ep(e: fdfield_t) -> fdfield_t:
e[direction][boundary_slice] = -e[direction][shifted2_slice] e[direction][boundary_slice] = -e[direction][shifted2_slice]
e[direction][shifted1_slice] = 0 e[direction][shifted1_slice] = 0
e[u][boundary_slice] = e[u][shifted1_slice] e[u][boundary_slice] = e[u][shifted1_slice]
e[v][boundary_slice] = e[v][shifted1_slice] e[v][boundary_slice] = e[v][shifted1_slice]
return e return e
def hp(h: fdfield_t): def hp(h: fdfield_t) -> fdfield_t:
h[direction][boundary_slice] = h[direction][shifted1_slice] h[direction][boundary_slice] = h[direction][shifted1_slice]
h[u][boundary_slice] = -h[u][shifted2_slice] h[u][boundary_slice] = -h[u][shifted2_slice]
h[u][shifted1_slice] = 0 h[u][shifted1_slice] = 0

View File

@ -5,6 +5,9 @@ from ..fdmath import dx_lists_t, fdfield_t
from ..fdmath.functional import deriv_back from ..fdmath.functional import deriv_back
# TODO documentation
def poynting(e: fdfield_t, def poynting(e: fdfield_t,
h: fdfield_t, h: fdfield_t,
dxes: Optional[dx_lists_t] = None, dxes: Optional[dx_lists_t] = None,

View File

@ -63,7 +63,7 @@ def cpml(direction: int,
expand_slice_l[direction] = slice(None) expand_slice_l[direction] = slice(None)
expand_slice = tuple(expand_slice_l) expand_slice = tuple(expand_slice_l)
def par(x): def par(x: numpy.ndarray) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
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)

View File

@ -3,6 +3,7 @@
Test fixtures Test fixtures
""" """
from typing import Tuple, Iterable, List
import numpy # type: ignore import numpy # type: ignore
import pytest # type: ignore import pytest # type: ignore
@ -14,22 +15,26 @@ from .utils import PRNG
(5, 5, 5), (5, 5, 5),
# (7, 7, 7), # (7, 7, 7),
]) ])
def shape(request): def shape(request: pytest.FixtureRequest) -> Iterable[Tuple[int, ...]]:
yield (3, *request.param) yield (3, *request.param)
@pytest.fixture(scope='module', params=[1.0, 1.5]) @pytest.fixture(scope='module', params=[1.0, 1.5])
def epsilon_bg(request): def epsilon_bg(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@pytest.fixture(scope='module', params=[1.0, 2.5]) @pytest.fixture(scope='module', params=[1.0, 2.5])
def epsilon_fg(request): def epsilon_fg(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@pytest.fixture(scope='module', params=['center', '000', 'random']) @pytest.fixture(scope='module', params=['center', '000', 'random'])
def epsilon(request, shape, epsilon_bg, epsilon_fg): def epsilon(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
epsilon_bg: float,
epsilon_fg: float,
) -> Iterable[numpy.ndarray]:
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':
@ -53,17 +58,20 @@ def epsilon(request, shape, epsilon_bg, epsilon_fg):
@pytest.fixture(scope='module', params=[1.0]) # 1.5 @pytest.fixture(scope='module', params=[1.0]) # 1.5
def j_mag(request): def j_mag(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@pytest.fixture(scope='module', params=[1.0, 1.5]) @pytest.fixture(scope='module', params=[1.0, 1.5])
def dx(request): def dx(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@pytest.fixture(scope='module', params=['uniform', 'centerbig']) @pytest.fixture(scope='module', params=['uniform', 'centerbig'])
def dxes(request, shape, dx): def dxes(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
dx: float,
) -> Iterable[List[List[numpy.ndarray]]]:
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':

View File

@ -1,4 +1,4 @@
from typing import List, Tuple 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 # type: ignore
@ -9,14 +9,14 @@ from ..fdmath import vec, unvec
from .utils import assert_close # , assert_fields_close from .utils import assert_close # , assert_fields_close
def test_residual(sim): def test_residual(sim: 'FDResult') -> None:
A = fdfd.operators.e_full(sim.omega, sim.dxes, vec(sim.epsilon)).tocsr() A = fdfd.operators.e_full(sim.omega, sim.dxes, vec(sim.epsilon)).tocsr()
b = -1j * sim.omega * vec(sim.j) b = -1j * sim.omega * vec(sim.j)
residual = A @ vec(sim.e) - b residual = A @ vec(sim.e) - b
assert numpy.linalg.norm(residual) < 1e-10 assert numpy.linalg.norm(residual) < 1e-10
def test_poynting_planes(sim): def test_poynting_planes(sim: 'FDResult') -> None:
mask = (sim.j != 0).any(axis=0) mask = (sim.j != 0).any(axis=0)
if mask.sum() != 2: if mask.sum() != 2:
pytest.skip(f'test_poynting_planes will only test 2-point sources, got {mask.sum()}') pytest.skip(f'test_poynting_planes will only test 2-point sources, got {mask.sum()}')
@ -53,17 +53,17 @@ def test_poynting_planes(sim):
# Also see conftest.py # Also see conftest.py
@pytest.fixture(params=[1 / 1500]) @pytest.fixture(params=[1 / 1500])
def omega(request): def omega(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@pytest.fixture(params=[None]) @pytest.fixture(params=[None])
def pec(request): def pec(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
yield request.param yield request.param
@pytest.fixture(params=[None]) @pytest.fixture(params=[None])
def pmc(request): def pmc(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
yield request.param yield request.param
@ -74,7 +74,10 @@ def pmc(request):
@pytest.fixture(params=['diag']) # 'center' @pytest.fixture(params=['diag']) # 'center'
def j_distribution(request, shape, j_mag): def j_distribution(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
j_mag: float,
) -> Iterable[numpy.ndarray]:
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
@ -89,7 +92,7 @@ def j_distribution(request, shape, j_mag):
@dataclasses.dataclass() @dataclasses.dataclass()
class FDResult: class FDResult:
shape: Tuple[int] shape: Tuple[int, ...]
dxes: List[List[numpy.ndarray]] dxes: List[List[numpy.ndarray]]
epsilon: numpy.ndarray epsilon: numpy.ndarray
omega: complex omega: complex
@ -100,7 +103,15 @@ class FDResult:
@pytest.fixture() @pytest.fixture()
def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc): def sim(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
dxes: List[List[numpy.ndarray]],
j_distribution: numpy.ndarray,
omega: float,
pec: Optional[numpy.ndarray],
pmc: Optional[numpy.ndarray],
) -> FDResult:
""" """
Build simulation from parts Build simulation from parts
""" """

View File

@ -1,15 +1,15 @@
##################################### from typing import Optional, Tuple, Iterable, List
import pytest # type: ignore import pytest # type: ignore
import numpy # type: ignore import numpy # type: ignore
from numpy.testing import assert_allclose # type: ignore from numpy.testing import assert_allclose # type: ignore
from .. import fdfd from .. import fdfd
from ..fdmath import vec, unvec from ..fdmath import vec, unvec, dx_lists_mut
#from .utils import assert_close, assert_fields_close #from .utils import assert_close, assert_fields_close
from .test_fdfd import FDResult from .test_fdfd import FDResult
def test_pml(sim, src_polarity): def test_pml(sim: FDResult, src_polarity: int) -> None:
e_sqr = numpy.squeeze((sim.e.conj() * sim.e).sum(axis=0)) e_sqr = numpy.squeeze((sim.e.conj() * sim.e).sum(axis=0))
# from matplotlib import pyplot # from matplotlib import pyplot
@ -42,34 +42,40 @@ def test_pml(sim, src_polarity):
# Also see conftest.py # Also see conftest.py
@pytest.fixture(params=[1 / 1500]) @pytest.fixture(params=[1 / 1500])
def omega(request): def omega(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@pytest.fixture(params=[None]) @pytest.fixture(params=[None])
def pec(request): def pec(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
yield request.param yield request.param
@pytest.fixture(params=[None]) @pytest.fixture(params=[None])
def pmc(request): def pmc(request: pytest.FixtureRequest) -> Iterable[Optional[numpy.ndarray]]:
yield request.param yield request.param
@pytest.fixture(params=[(30, 1, 1), @pytest.fixture(params=[(30, 1, 1),
(1, 30, 1), (1, 30, 1),
(1, 1, 30)]) (1, 1, 30)])
def shape(request): def shape(request: pytest.FixtureRequest) -> Iterable[Tuple[int, ...]]:
yield (3, *request.param) yield (3, *request.param)
@pytest.fixture(params=[+1, -1]) @pytest.fixture(params=[+1, -1])
def src_polarity(request): def src_polarity(request: pytest.FixtureRequest) -> Iterable[int]:
yield request.param yield request.param
@pytest.fixture() @pytest.fixture()
def j_distribution(request, shape, epsilon, dxes, omega, src_polarity): def j_distribution(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
dxes: dx_lists_mut,
omega: float,
src_polarity: int,
) -> Iterable[numpy.ndarray]:
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
@ -101,13 +107,22 @@ def j_distribution(request, shape, epsilon, dxes, omega, src_polarity):
@pytest.fixture() @pytest.fixture()
def epsilon(request, shape, epsilon_bg, epsilon_fg): def epsilon(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
epsilon_bg: float,
epsilon_fg: float,
) -> Iterable[numpy.ndarray]:
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, shape, dx, omega, epsilon_fg): def dxes(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
dx: float,
omega: float,
epsilon_fg: float,
) -> Iterable[List[List[numpy.ndarray]]]:
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
@ -120,7 +135,15 @@ def dxes(request, shape, dx, omega, epsilon_fg):
@pytest.fixture() @pytest.fixture()
def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc): def sim(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
dxes: dx_lists_mut,
j_distribution: numpy.ndarray,
omega: float,
pec: Optional[numpy.ndarray],
pmc: Optional[numpy.ndarray],
) -> 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(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec,
@ -129,7 +152,7 @@ def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
sim = FDResult( sim = FDResult(
shape=shape, shape=shape,
dxes=dxes, dxes=[list(d) for d in dxes],
epsilon=epsilon, epsilon=epsilon,
j=j_distribution, j=j_distribution,
e=e, e=e,

View File

@ -1,4 +1,4 @@
from typing import List, Tuple from typing import List, Tuple, Iterable
import dataclasses import dataclasses
import pytest # type: ignore import pytest # type: ignore
import numpy # type: ignore import numpy # type: ignore
@ -8,7 +8,7 @@ from .. import fdtd
from .utils import assert_close, assert_fields_close, PRNG from .utils import assert_close, assert_fields_close, PRNG
def test_initial_fields(sim): def test_initial_fields(sim: 'TDResult') -> None:
# Make sure initial fields didn't change # Make sure initial fields didn't change
e0 = sim.es[0] e0 = sim.es[0]
h0 = sim.hs[0] h0 = sim.hs[0]
@ -20,7 +20,7 @@ def test_initial_fields(sim):
assert not h0.any() assert not h0.any()
def test_initial_energy(sim): def test_initial_energy(sim: 'TDResult') -> None:
""" """
Assumes fields start at 0 before J0 is added Assumes fields start at 0 before J0 is added
""" """
@ -41,7 +41,7 @@ def test_initial_energy(sim):
assert_fields_close(e0_dot_j0, u0) assert_fields_close(e0_dot_j0, u0)
def test_energy_conservation(sim): def test_energy_conservation(sim: 'TDResult') -> None:
""" """
Assumes fields start at 0 before J0 is added Assumes fields start at 0 before J0 is added
""" """
@ -63,7 +63,7 @@ def test_energy_conservation(sim):
assert_close(u_estep.sum(), u) assert_close(u_estep.sum(), u)
def test_poynting_divergence(sim): def test_poynting_divergence(sim: 'TDResult') -> None:
args = {'dxes': sim.dxes, args = {'dxes': sim.dxes,
'epsilon': sim.epsilon} 'epsilon': sim.epsilon}
@ -90,7 +90,7 @@ def test_poynting_divergence(sim):
u_eprev = u_estep u_eprev = u_estep
def test_poynting_planes(sim): def test_poynting_planes(sim: 'TDResult') -> None:
mask = (sim.js[0] != 0).any(axis=0) mask = (sim.js[0] != 0).any(axis=0)
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()))
@ -140,30 +140,33 @@ def test_poynting_planes(sim):
@pytest.fixture(params=[0.3]) @pytest.fixture(params=[0.3])
def dt(request): def dt(request: pytest.FixtureRequest) -> Iterable[float]:
yield request.param yield request.param
@dataclasses.dataclass() @dataclasses.dataclass()
class TDResult: class TDResult:
shape: Tuple[int] shape: Tuple[int, ...]
dt: float dt: float
dxes: List[List[numpy.ndarray]] dxes: List[List[numpy.ndarray]]
epsilon: numpy.ndarray epsilon: numpy.ndarray
j_distribution: numpy.ndarray j_distribution: numpy.ndarray
j_steps: Tuple[int] j_steps: Tuple[int, ...]
es: List[numpy.ndarray] = dataclasses.field(default_factory=list) es: List[numpy.ndarray] = dataclasses.field(default_factory=list)
hs: 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) js: List[numpy.ndarray] = dataclasses.field(default_factory=list)
@pytest.fixture(params=[(0, 4, 8)]) # (0,) @pytest.fixture(params=[(0, 4, 8)]) # (0,)
def j_steps(request): def j_steps(request: pytest.fixtureRequest) -> Iterable[Tuple[int, ...]]:
yield request.param yield request.param
@pytest.fixture(params=['center', 'random']) @pytest.fixture(params=['center', 'random'])
def j_distribution(request, shape, j_mag): def j_distribution(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
j_mag: float,
) -> Iterable[numpy.ndarray]:
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
@ -175,7 +178,14 @@ def j_distribution(request, shape, j_mag):
@pytest.fixture() @pytest.fixture()
def sim(request, shape, epsilon, dxes, dt, j_distribution, j_steps): def sim(request: pytest.FixtureRequest,
shape: Tuple[int, ...],
epsilon: numpy.ndarray,
dxes: List[List[numpy.ndarray]],
dt: float,
j_distribution: numpy.ndarray,
j_steps: Tuple[int, ...],
) -> TDResult:
is3d = (numpy.array(shape) == 1).sum() == 0 is3d = (numpy.array(shape) == 1).sum() == 0
if is3d: if is3d:
if dt != 0.3: if dt != 0.3:

View File

@ -1,13 +1,22 @@
from typing import Any
import numpy # type: ignore import numpy # type: ignore
PRNG = numpy.random.RandomState(12345) PRNG = numpy.random.RandomState(12345)
def assert_fields_close(x, y, *args, **kwargs): def assert_fields_close(x: numpy.ndarray,
y: numpy.ndarray,
*args: Any,
**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, y, *args, **kwargs): def assert_close(x: numpy.ndarray,
y: numpy.ndarray,
*args: Any,
**kwargs: Any,
) -> None:
numpy.testing.assert_allclose(x, y, *args, **kwargs) numpy.testing.assert_allclose(x, y, *args, **kwargs)