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

modernize type annotations

Browse Source
This commit is contained in:
Jan Petykiewicz 2023-05-22 10:53:13 -07:00
parent 7009e505e7
commit 44465f1bc9
22 changed files with 245 additions and 248 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
meanas/eigensolvers.py
View File

@ -1,7 +1,7 @@
"""
Solvers for eigenvalue / eigenvector problems
"""
from typing import Tuple, Callable, Optional, Union
from typing import Callable
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm
@ -11,9 +11,9 @@ import scipy.sparse.linalg as spalg # type: ignore
def power_iteration(
operator: sparse.spmatrix,
guess_vector: Optional[NDArray[numpy.complex128]] = None,
guess_vector: NDArray[numpy.complex128] | None = None,
iterations: int = 20,
) -> Tuple[complex, NDArray[numpy.complex128]]:
) -> tuple[complex, NDArray[numpy.complex128]]:
"""
Use power iteration to estimate the dominant eigenvector of a matrix.
@ -40,12 +40,12 @@ def power_iteration(
def rayleigh_quotient_iteration(
operator: Union[sparse.spmatrix, spalg.LinearOperator],
operator: sparse.spmatrix | spalg.LinearOperator,
guess_vector: NDArray[numpy.complex128],
iterations: int = 40,
tolerance: float = 1e-13,
solver: Optional[Callable[..., NDArray[numpy.complex128]]] = None,
) -> Tuple[complex, NDArray[numpy.complex128]]:
solver: Callable[..., NDArray[numpy.complex128]] | None = None,
) -> tuple[complex, NDArray[numpy.complex128]]:
"""
Use Rayleigh quotient iteration to refine an eigenvector guess.
@ -80,7 +80,7 @@ def rayleigh_quotient_iteration(
if solver is None:
def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.complex128]:
return spalg.bicgstab(A, b)[0]
assert(solver is not None)
assert solver is not None
v = numpy.squeeze(guess_vector)
v /= norm(v)
@ -96,10 +96,10 @@ def rayleigh_quotient_iteration(
def signed_eigensolve(
operator: Union[sparse.spmatrix, spalg.LinearOperator],
operator: sparse.spmatrix | spalg.LinearOperator,
how_many: int,
negative: bool = False,
) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
"""
Find the largest-magnitude positive-only (or negative-only) eigenvalues and
eigenvectors of the provided matrix.

55
meanas/fdfd/bloch.py
View File

@ -94,7 +94,7 @@ This module contains functions for generating and solving the
'''
from typing import Tuple, Callable, Any, List, Optional, cast, Union, Sequence
from typing import Callable, Any, cast, Sequence
import logging
import numpy
from numpy import pi, real, trace
@ -140,7 +140,7 @@ def generate_kmn(
k0: ArrayLike,
G_matrix: ArrayLike,
shape: Sequence[int],
) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
"""
Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.
@ -155,8 +155,9 @@ def generate_kmn(
All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
"""
k0 = numpy.array(k0)
G_matrix = numpy.array(G_matrix, copy=False)
Gi_grids = numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij')
Gi_grids = numpy.array(numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij'))
Gi = numpy.moveaxis(Gi_grids, 0, -1)
k_G = k0[None, None, None, :] - Gi
@ -183,7 +184,7 @@ def maxwell_operator(
k0: ArrayLike,
G_matrix: ArrayLike,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None
mu: fdfield_t | None = None
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
"""
Generate the Maxwell operator
@ -237,7 +238,7 @@ def maxwell_operator(
# cross product and transform into xyz basis
d_xyz = (n * hin_m
- m * hin_n) * k_mag
- m * hin_n) * k_mag # noqa: E128
# divide by epsilon
temp = ifftn(d_xyz, axes=range(3)) # reuses d_xyz if using pyfftw
@ -253,7 +254,7 @@ def maxwell_operator(
else:
# transform from mn to xyz
b_xyz = (m * b_m[:, :, :, None]
+ n * b_n[:, :, :, None])
+ n * b_n[:, :, :, None]) # noqa: E128
# divide by mu
temp = ifftn(b_xyz, axes=range(3))
@ -302,7 +303,7 @@ def hmn_2_exyz(
def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
d_xyz = (n * hin_m
- m * hin_n) * k_mag
- m * hin_n) * k_mag # noqa: E128
# divide by epsilon
return numpy.array([ei for ei in numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)]) # TODO avoid copy
@ -340,7 +341,7 @@ def hmn_2_hxyz(
def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
h_xyz = (m * hin_m
+ n * hin_n)
+ n * hin_n) # noqa: E128
return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])
return operator
@ -350,7 +351,7 @@ def inverse_maxwell_operator_approx(
k0: ArrayLike,
G_matrix: ArrayLike,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
"""
Generate an approximate inverse of the Maxwell operator,
@ -400,7 +401,7 @@ def inverse_maxwell_operator_approx(
else:
# transform from mn to xyz
h_xyz = (m * hin_m[:, :, :, None]
+ n * hin_n[:, :, :, None])
+ n * hin_n[:, :, :, None]) # noqa: E128
# multiply by mu
temp = ifftn(h_xyz, axes=range(3))
@ -413,7 +414,7 @@ def inverse_maxwell_operator_approx(
# cross product and transform into xyz basis
e_xyz = (n * b_m
- m * b_n) / k_mag
- m * b_n) / k_mag # noqa: E128
# multiply by epsilon
temp = ifftn(e_xyz, axes=range(3))
@ -436,13 +437,13 @@ def find_k(
direction: ArrayLike,
G_matrix: ArrayLike,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
band: int = 0,
k_bounds: Tuple[float, float] = (0, 0.5),
k_guess: Optional[float] = None,
solve_callback: Optional[Callable[..., None]] = None,
iter_callback: Optional[Callable[..., None]] = None,
) -> Tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
k_bounds: tuple[float, float] = (0, 0.5),
k_guess: float | None = None,
solve_callback: Callable[..., None] | None = None,
iter_callback: Callable[..., None] | None = None,
) -> tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
"""
Search for a bloch vector that has a given frequency.
@ -503,13 +504,13 @@ def eigsolve(
k0: ArrayLike,
G_matrix: ArrayLike,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
tolerance: float = 1e-20,
max_iters: int = 10000,
reset_iters: int = 100,
y0: Optional[ArrayLike] = None,
callback: Optional[Callable[..., None]] = None,
) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
y0: ArrayLike | None = None,
callback: Callable[..., None] | None = None,
) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
"""
Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
k0 of the specified structure.
@ -555,6 +556,7 @@ def eigsolve(
#prev_theta = 0.5
D = numpy.zeros(shape=y_shape, dtype=complex)
Z: NDArray[numpy.complex128]
if y0 is None:
Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
else:
@ -593,7 +595,8 @@ def eigsolve(
# via (1 - ZUZt)
if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7):
logger.info('Optimization succeded: '
logger.info(
'Optimization succeded: '
f'[change in trace] {abs(E - prev_E)} - 5e-8 '
f'< {tolerance} [tolerance] * {(E + prev_E) / 2} [value of trace]'
)
@ -635,7 +638,7 @@ def eigsolve(
symZtD = _symmetrize(Zt @ D)
symZtAD = _symmetrize(Zt @ AD)
Qi_memo: List[Optional[float]] = [None, None]
Qi_memo: list[float | None] = [None, None]
def Qi_func(theta: float) -> float:
nonlocal Qi_memo
@ -659,8 +662,8 @@ def eigsolve(
else:
raise Exception('Inexplicable singularity in trace_func')
Qi_memo[0] = theta
Qi_memo[1] = Qi
return Qi
Qi_memo[1] = cast(float, Qi)
return cast(float, Qi)
def trace_func(theta: float) -> float:
c = numpy.cos(theta)
@ -772,7 +775,7 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to
def _rtrace_AtB(
A: NDArray[numpy.complex128],
B: Union[NDArray[numpy.complex128], float],
B: NDArray[numpy.complex128] | float,
) -> float:
return real(numpy.sum(A.conj() * B))

36
meanas/fdfd/farfield.py
View File

@ -1,7 +1,7 @@
"""
Functions for performing near-to-farfield transformation (and the reverse).
"""
from typing import Dict, List, Any, Union, Sequence
from typing import Any, Sequence, cast
import numpy
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
from numpy import pi
@ -14,8 +14,8 @@ def near_to_farfield(
H_near: cfdfield_t,
dx: float,
dy: float,
padded_size: Union[List[int], int, None] = None
) -> Dict[str, Any]:
padded_size: list[int] | int | None = None
) -> dict[str, Any]:
"""
Compute the farfield, i.e. the distribution of the fields after propagation
through several wavelengths of uniform medium.
@ -62,14 +62,15 @@ def near_to_farfield(
padded_size = (2**numpy.ceil(numpy.log2(s))).astype(int)
if not hasattr(padded_size, '__len__'):
padded_size = (padded_size, padded_size) # type: ignore # checked if sequence
padded_shape = cast(Sequence[int], padded_size)
En_fft = [fftshift(fft2(fftshift(Eni), s=padded_size)) for Eni in E_near]
Hn_fft = [fftshift(fft2(fftshift(Hni), s=padded_size)) for Hni in H_near]
En_fft = [fftshift(fft2(fftshift(Eni), s=padded_shape)) for Eni in E_near]
Hn_fft = [fftshift(fft2(fftshift(Hni), s=padded_shape)) for Hni in H_near]
# Propagation vectors kx, ky
k = 2 * pi
kxx = 2 * pi * fftshift(fftfreq(padded_size[0], dx))
kyy = 2 * pi * fftshift(fftfreq(padded_size[1], dy))
kxx = 2 * pi * fftshift(fftfreq(padded_shape[0], dx))
kyy = 2 * pi * fftshift(fftfreq(padded_shape[1], dy))
kx, ky = numpy.meshgrid(kxx, kyy, indexing='ij')
kxy2 = kx * kx + ky * ky
@ -85,14 +86,14 @@ def near_to_farfield(
# Normalized vector potentials N, L
N = [-Hn_fft[1] * cos_phi * cos_th + Hn_fft[0] * cos_phi * sin_th,
Hn_fft[1] * sin_th + Hn_fft[0] * cos_th]
Hn_fft[1] * sin_th + Hn_fft[0] * cos_th] # noqa: E127
L = [ En_fft[1] * cos_phi * cos_th - En_fft[0] * cos_phi * sin_th,
-En_fft[1] * sin_th - En_fft[0] * cos_th]
-En_fft[1] * sin_th - En_fft[0] * cos_th] # noqa: E128
E_far = [-L[1] - N[0],
L[0] - N[1]]
L[0] - N[1]] # noqa: E127
H_far = [-E_far[1],
E_far[0]]
E_far[0]] # noqa: E127
theta = numpy.arctan2(ky, kx)
phi = numpy.arccos(cos_phi)
@ -126,8 +127,8 @@ def far_to_nearfield(
H_far: cfdfield_t,
dkx: float,
dky: float,
padded_size: Union[List[int], int, None] = None
) -> Dict[str, Any]:
padded_size: list[int] | int | None = None
) -> dict[str, Any]:
"""
Compute the farfield, i.e. the distribution of the fields after propagation
through several wavelengths of uniform medium.
@ -170,6 +171,7 @@ def far_to_nearfield(
padded_size = (2 ** numpy.ceil(numpy.log2(s))).astype(int)
if not hasattr(padded_size, '__len__'):
padded_size = (padded_size, padded_size) # type: ignore # checked if sequence
padded_shape = cast(Sequence[int], padded_size)
k = 2 * pi
kxs = fftshift(fftfreq(s[0], 1 / (s[0] * dkx)))
@ -203,9 +205,9 @@ def far_to_nearfield(
# Normalized vector potentials N, L
L = [0.5 * E_far[1],
-0.5 * E_far[0]]
-0.5 * E_far[0]] # noqa: E128
N = [L[1],
-L[0]]
-L[0]] # noqa: E128
En_fft = [-( L[0] * sin_th + L[1] * cos_phi * cos_th) / cos_phi,
-(-L[0] * cos_th + L[1] * cos_phi * sin_th) / cos_phi]
@ -217,8 +219,8 @@ def far_to_nearfield(
En_fft[i][cos_phi == 0] = 0
Hn_fft[i][cos_phi == 0] = 0
E_near = [ifftshift(ifft2(ifftshift(Ei), s=padded_size)) for Ei in En_fft]
H_near = [ifftshift(ifft2(ifftshift(Hi), s=padded_size)) for Hi in Hn_fft]
E_near = [ifftshift(ifft2(ifftshift(Ei), s=padded_shape)) for Ei in En_fft]
H_near = [ifftshift(ifft2(ifftshift(Hi), s=padded_shape)) for Hi in Hn_fft]
dx = 2 * pi / (s[0] * dkx)
dy = 2 * pi / (s[0] * dky)

18
meanas/fdfd/functional.py
View File

@ -5,7 +5,7 @@ Functional versions of many FDFD operators. These can be useful for performing
The functions generated here expect `cfdfield_t` inputs with shape (3, X, Y, Z),
e.g. E = [E_x, E_y, E_z] where each (complex) component has shape (X, Y, Z)
"""
from typing import Callable, Tuple, Optional
from typing import Callable
import numpy
from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t
@ -19,7 +19,7 @@ def e_full(
omega: complex,
dxes: dx_lists_t,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
) -> cfdfield_updater_t:
"""
Wave operator for use with E-field. See `operators.e_full` for details.
@ -55,8 +55,8 @@ def eh_full(
omega: complex,
dxes: dx_lists_t,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
) -> Callable[[cfdfield_t, cfdfield_t], Tuple[cfdfield_t, cfdfield_t]]:
mu: fdfield_t | None = None,
) -> Callable[[cfdfield_t, cfdfield_t], tuple[cfdfield_t, cfdfield_t]]:
"""
Wave operator for full (both E and H) field representation.
See `operators.eh_full`.
@ -74,11 +74,11 @@ def eh_full(
ch = curl_back(dxes[1])
ce = curl_forward(dxes[0])
def op_1(e: cfdfield_t, h: cfdfield_t) -> Tuple[cfdfield_t, cfdfield_t]:
def op_1(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]:
return (ch(h) - 1j * omega * epsilon * e,
ce(e) + 1j * omega * h)
def op_mu(e: cfdfield_t, h: cfdfield_t) -> Tuple[cfdfield_t, cfdfield_t]:
def op_mu(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]:
return (ch(h) - 1j * omega * epsilon * e,
ce(e) + 1j * omega * mu * h) # type: ignore # mu=None ok
@ -91,7 +91,7 @@ def eh_full(
def e2h(
omega: complex,
dxes: dx_lists_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
) -> cfdfield_updater_t:
"""
Utility operator for converting the `E` field into the `H` field.
@ -123,7 +123,7 @@ def e2h(
def m2j(
omega: complex,
dxes: dx_lists_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
) -> cfdfield_updater_t:
"""
Utility operator for converting magnetic current `M` distribution
@ -160,7 +160,7 @@ def e_tfsf_source(
omega: complex,
dxes: dx_lists_t,
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
) -> cfdfield_updater_t:
"""
Operator that turns an E-field distribution into a total-field/scattered-field

31
meanas/fdfd/operators.py
View File

@ -27,7 +27,6 @@ The following operators are included:
- Cross product matrices
"""
from typing import Tuple, Optional
import numpy
import scipy.sparse as sparse # type: ignore
@ -42,9 +41,9 @@ def e_full(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
pec: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
pec: vfdfield_t | None = None,
pmc: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Wave operator
@ -99,7 +98,7 @@ def e_full(
def e_full_preconditioners(
dxes: dx_lists_t,
) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
) -> tuple[sparse.spmatrix, sparse.spmatrix]:
"""
Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator.
@ -128,9 +127,9 @@ def h_full(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
pec: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
pec: vfdfield_t | None = None,
pmc: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Wave operator
@ -185,9 +184,9 @@ def eh_full(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
pec: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
pec: vfdfield_t | None = None,
pmc: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Wave operator for `[E, H]` field representation. This operator implements Maxwell's
@ -254,8 +253,8 @@ def eh_full(
def e2h(
omega: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None,
pmc: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
pmc: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Utility operator for converting the E field into the H field.
@ -286,7 +285,7 @@ def e2h(
def m2j(
omega: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Operator for converting a magnetic current M into an electric current J.
@ -368,7 +367,7 @@ def e_tfsf_source(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Operator that turns a desired E-field distribution into a
@ -399,7 +398,7 @@ def e_boundary_source(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
periodic_mask_edges: bool = False,
) -> sparse.spmatrix:
"""

16
meanas/fdfd/scpml.py
View File

@ -2,10 +2,10 @@
Functions for creating stretched coordinate perfectly matched layer (PML) absorbers.
"""
from typing import Sequence, Union, Callable, Optional, List
from typing import Sequence, Callable
import numpy
from numpy.typing import ArrayLike, NDArray
from numpy.typing import NDArray
__author__ = 'Jan Petykiewicz'
@ -43,8 +43,8 @@ def uniform_grid_scpml(
thicknesses: Sequence[int],
omega: float,
epsilon_effective: float = 1.0,
s_function: Optional[s_function_t] = None,
) -> List[List[NDArray[numpy.float64]]]:
s_function: s_function_t | None = None,
) -> list[list[NDArray[numpy.float64]]]:
"""
Create dx arrays for a uniform grid with a cell width of 1 and a pml.
@ -86,7 +86,7 @@ def uniform_grid_scpml(
for k, th in enumerate(thicknesses):
s = shape[k]
if th > 0:
sr = numpy.arange(s)
sr = numpy.arange(s, dtype=numpy.float64)
dx_a[k] = 1 + 1j * s_function(ll(sr, s, th)) / s_correction
dx_b[k] = 1 + 1j * s_function(ll(sr + 0.5, s, th)) / s_correction
else:
@ -96,14 +96,14 @@ def uniform_grid_scpml(
def stretch_with_scpml(
dxes: List[List[NDArray[numpy.float64]]],
dxes: list[list[NDArray[numpy.float64]]],
axis: int,
polarity: int,
omega: float,
epsilon_effective: float = 1.0,
thickness: int = 10,
s_function: Optional[s_function_t] = None,
) -> List[List[NDArray[numpy.float64]]]:
s_function: s_function_t | None = None,
) -> list[list[NDArray[numpy.float64]]]:
"""
Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis.

42
meanas/fdfd/waveguide_2d.py
View File

@ -178,7 +178,7 @@ to account for numerical dispersion if the result is introduced into a space wit
"""
# TODO update module docs
from typing import List, Tuple, Optional, Any
from typing import Any
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm
@ -196,7 +196,7 @@ def operator_e(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Waveguide operator of the form
@ -263,7 +263,7 @@ def operator_h(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
) -> sparse.spmatrix:
"""
Waveguide operator of the form
@ -333,9 +333,9 @@ def normalized_fields_e(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
prop_phase: float = 0,
) -> Tuple[vcfdfield_t, vcfdfield_t]:
) -> tuple[vcfdfield_t, vcfdfield_t]:
"""
Given a vector `e_xy` containing the vectorized E_x and E_y fields,
returns normalized, vectorized E and H fields for the system.
@ -368,9 +368,9 @@ def normalized_fields_h(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
prop_phase: float = 0,
) -> Tuple[vcfdfield_t, vcfdfield_t]:
) -> tuple[vcfdfield_t, vcfdfield_t]:
"""
Given a vector `h_xy` containing the vectorized H_x and H_y fields,
returns normalized, vectorized E and H fields for the system.
@ -403,9 +403,9 @@ def _normalized_fields(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
prop_phase: float = 0,
) -> Tuple[vcfdfield_t, vcfdfield_t]:
) -> tuple[vcfdfield_t, vcfdfield_t]:
# TODO documentation
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)]
@ -445,7 +445,7 @@ def exy2h(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None
mu: vfdfield_t | None = None
) -> sparse.spmatrix:
"""
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
@ -471,7 +471,7 @@ def hxy2e(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None
mu: vfdfield_t | None = None
) -> sparse.spmatrix:
"""
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
@ -495,7 +495,7 @@ def hxy2e(
def hxy2h(
wavenumber: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None
mu: vfdfield_t | None = None
) -> sparse.spmatrix:
"""
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
@ -564,7 +564,7 @@ def e2h(
wavenumber: complex,
omega: complex,
dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None
mu: vfdfield_t | None = None
) -> sparse.spmatrix:
"""
Returns an operator which, when applied to a vectorized E eigenfield, produces
@ -654,7 +654,7 @@ def h_err(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None
mu: vfdfield_t | None = None
) -> float:
"""
Calculates the relative error in the H field
@ -680,7 +680,7 @@ def h_err(
else:
op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h)
return norm(op) / norm(h)
return float(norm(op) / norm(h))
def e_err(
@ -689,7 +689,7 @@ def e_err(
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
) -> float:
"""
Calculates the relative error in the E field
@ -714,17 +714,17 @@ def e_err(
mu_inv = sparse.diags(1 / mu)
op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
return norm(op) / norm(e)
return float(norm(op) / norm(e))
def solve_modes(
mode_numbers: List[int],
mode_numbers: list[int],
omega: complex,
dxes: dx_lists_t,
epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None,
mu: vfdfield_t | None = None,
mode_margin: int = 2,
) -> Tuple[NDArray[numpy.float64], NDArray[numpy.complex128]]:
) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
"""
Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
@ -772,7 +772,7 @@ def solve_mode(
mode_number: int,
*args: Any,
**kwargs: Any,
) -> Tuple[vcfdfield_t, complex]:
) -> tuple[vcfdfield_t, complex]:
"""
Wrapper around `solve_modes()` that solves for a single mode.

14
meanas/fdfd/waveguide_3d.py
View File

@ -4,7 +4,7 @@ Tools for working with waveguide modes in 3D domains.
This module relies heavily on `waveguide_2d` and mostly just transforms
its parameters into 2D equivalents and expands the results back into 3D.
"""
from typing import Dict, Optional, Sequence, Union, Any
from typing import Sequence, Any
import numpy
from numpy.typing import NDArray
@ -20,8 +20,8 @@ def solve_mode(
polarity: int,
slices: Sequence[slice],
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
) -> Dict[str, Union[complex, NDArray[numpy.float_]]]:
mu: fdfield_t | None = None,
) -> dict[str, complex | NDArray[numpy.float_]]:
"""
Given a 3D grid, selects a slice from the grid and attempts to
solve for an eigenmode propagating through that slice.
@ -40,8 +40,8 @@ def solve_mode(
Returns:
```
{
'E': List[NDArray[numpy.float_]],
'H': List[NDArray[numpy.float_]],
'E': list[NDArray[numpy.float_]],
'H': list[NDArray[numpy.float_]],
'wavenumber': complex,
}
```
@ -63,7 +63,7 @@ def solve_mode(
dx_prop = 0.5 * dxab_forward.sum()
# Reduce to 2D and solve the 2D problem
args_2d: Dict[str, Any] = {
args_2d: dict[str, Any] = {
'omega': omega,
'dxes': [[dx[i][slices[i]] for i in order[:2]] for dx in dxes],
'epsilon': vec([epsilon[i][slices].transpose(order) for i in order]),
@ -114,7 +114,7 @@ def compute_source(
polarity: int,
slices: Sequence[slice],
epsilon: fdfield_t,
mu: Optional[fdfield_t] = None,
mu: fdfield_t | None = None,
) -> cfdfield_t:
"""
Given an eigenmode obtained by `solve_mode`, returns the current source distribution

7
meanas/fdfd/waveguide_cyl.py
View File

@ -8,7 +8,6 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
"""
# TODO update module docs
from typing import Dict, Union
import numpy
import scipy.sparse as sparse # type: ignore
@ -85,7 +84,7 @@ def solve_mode(
dxes: dx_lists_t,
epsilon: vfdfield_t,
r0: float,
) -> Dict[str, Union[complex, cfdfield_t]]:
) -> dict[str, complex | cfdfield_t]:
"""
TODO: fixup
Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@ -103,8 +102,8 @@ def solve_mode(
Returns:
```
{
'E': List[NDArray[numpy.complex_]],
'H': List[NDArray[numpy.complex_]],
'E': list[NDArray[numpy.complex_]],
'H': list[NDArray[numpy.complex_]],
'wavenumber': complex,
}
```

2
meanas/fdmath/__init__.py
View File

@ -8,7 +8,7 @@ Fields, Functions, and Operators
Discrete fields are stored in one of two forms:
- The `fdfield_t` form is a multidimensional `numpy.ndarray`
- The `fdfield_t` form is a multidimensional `numpy.NDArray`
+ For a scalar field, this is just `U[m, n, p]`, where `m`, `n`, and `p` are
discrete indices referring to positions on the x, y, and z axes respectively.
+ For a vector field, the first index specifies which vector component is accessed:

22
meanas/fdmath/functional.py
View File

@ -3,7 +3,7 @@ Math functions for finite difference simulations
Basic discrete calculus etc.
"""
from typing import Sequence, Tuple, Optional, Callable
from typing import Sequence, Callable
import numpy
from numpy.typing import NDArray
@ -12,8 +12,8 @@ from .types import fdfield_t, fdfield_updater_t
def deriv_forward(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
dx_e: Sequence[NDArray[numpy.float_]] | None = None,
) -> tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
"""
Utility operators for taking discretized derivatives (backward variant).
@ -36,8 +36,8 @@ def deriv_forward(
def deriv_back(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
dx_h: Sequence[NDArray[numpy.float_]] | None = None,
) -> tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
"""
Utility operators for taking discretized derivatives (forward variant).
@ -60,7 +60,7 @@ def deriv_back(
def curl_forward(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
dx_e: Sequence[NDArray[numpy.float_]] | None = None,
) -> fdfield_updater_t:
"""
Curl operator for use with the E field.
@ -89,7 +89,7 @@ def curl_forward(
def curl_back(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
dx_h: Sequence[NDArray[numpy.float_]] | None = None,
) -> fdfield_updater_t:
"""
Create a function which takes the backward curl of a field.
@ -118,11 +118,11 @@ def curl_back(
def curl_forward_parts(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None,
dx_e: Sequence[NDArray[numpy.float_]] | None = None,
) -> Callable:
Dx, Dy, Dz = deriv_forward(dx_e)
def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, fdfield_t], ...]:
def mkparts_fwd(e: fdfield_t) -> tuple[tuple[fdfield_t, fdfield_t], ...]:
return ((-Dz(e[1]), Dy(e[2])),
( Dz(e[0]), -Dx(e[2])),
(-Dy(e[0]), Dx(e[1])))
@ -131,11 +131,11 @@ def curl_forward_parts(
def curl_back_parts(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None,
dx_h: Sequence[NDArray[numpy.float_]] | None = None,
) -> Callable:
Dx, Dy, Dz = deriv_back(dx_h)
def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, fdfield_t], ...]:
def mkparts_back(h: fdfield_t) -> tuple[tuple[fdfield_t, fdfield_t], ...]:
return ((-Dz(h[1]), Dy(h[2])),
( Dz(h[0]), -Dx(h[2])),
(-Dy(h[0]), Dx(h[1])))

6
meanas/fdmath/operators.py
View File

@ -3,7 +3,7 @@ Matrix operators for finite difference simulations
Basic discrete calculus etc.
"""
from typing import Sequence, List
from typing import Sequence
import numpy
from numpy.typing import NDArray
import scipy.sparse as sparse # type: ignore
@ -98,7 +98,7 @@ def shift_with_mirror(
def deriv_forward(
dx_e: Sequence[NDArray[numpy.float_]],
) -> List[sparse.spmatrix]:
) -> list[sparse.spmatrix]:
"""
Utility operators for taking discretized derivatives (forward variant).
@ -125,7 +125,7 @@ def deriv_forward(
def deriv_back(
dx_h: Sequence[NDArray[numpy.float_]],
) -> List[sparse.spmatrix]:
) -> list[sparse.spmatrix]:
"""
Utility operators for taking discretized derivatives (backward variant).

8
meanas/fdmath/vectorization.py
View File

@ -4,7 +4,7 @@ and a 1D array representation of that field `[f_x0, f_x1, f_x2,... f_y0,... f_z0
Vectorized versions of the field use row-major (ie., C-style) ordering.
"""
from typing import Optional, overload, Union, Sequence
from typing import overload, Sequence
import numpy
from numpy.typing import ArrayLike
@ -24,10 +24,10 @@ def vec(f: cfdfield_t) -> vcfdfield_t:
pass
@overload
def vec(f: ArrayLike) -> Union[vfdfield_t, vcfdfield_t]:
def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t:
pass
def vec(f: Union[fdfield_t, cfdfield_t, ArrayLike, None]) -> Union[vfdfield_t, vcfdfield_t, None]:
def vec(f: fdfield_t | cfdfield_t | ArrayLike | None) -> vfdfield_t | vcfdfield_t | None:
"""
Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
@ -57,7 +57,7 @@ def unvec(v: vfdfield_t, shape: Sequence[int]) -> fdfield_t:
def unvec(v: vcfdfield_t, shape: Sequence[int]) -> cfdfield_t:
pass
def unvec(v: Union[vfdfield_t, vcfdfield_t, None], shape: Sequence[int]) -> Union[fdfield_t, cfdfield_t, None]:
def unvec(v: vfdfield_t | vcfdfield_t | None, shape: Sequence[int]) -> fdfield_t | cfdfield_t | None:
"""
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

10
meanas/fdtd/base.py
View File

@ -3,8 +3,6 @@ Basic FDTD field updates
"""
from typing import Union, Optional
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
from ..fdmath.functional import curl_forward, curl_back
@ -14,7 +12,7 @@ __author__ = 'Jan Petykiewicz'
def maxwell_e(
dt: float,
dxes: Optional[dx_lists_t] = None,
dxes: dx_lists_t | None = None,
) -> fdfield_updater_t:
"""
Build a function which performs a portion the time-domain E-field update,
@ -49,7 +47,7 @@ def maxwell_e(
else:
curl_h_fun = curl_back()
def me_fun(e: fdfield_t, h: fdfield_t, epsilon: Union[fdfield_t, float]) -> fdfield_t:
def me_fun(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t | float) -> fdfield_t:
"""
Update the E-field.
@ -69,7 +67,7 @@ def maxwell_e(
def maxwell_h(
dt: float,
dxes: Optional[dx_lists_t] = None,
dxes: dx_lists_t | None = None,
) -> fdfield_updater_t:
"""
Build a function which performs part of the time-domain H-field update,
@ -105,7 +103,7 @@ def maxwell_h(
else:
curl_e_fun = curl_forward()
def mh_fun(e: fdfield_t, h: fdfield_t, mu: Union[fdfield_t, float, None] = None) -> fdfield_t:
def mh_fun(e: fdfield_t, h: fdfield_t, mu: fdfield_t | float | None = None) -> fdfield_t:
"""
Update the H-field.

10
meanas/fdtd/boundaries.py
View File

@ -4,7 +4,7 @@ Boundary conditions
#TODO conducting boundary documentation
"""
from typing import Tuple, Any, List
from typing import Any
from ..fdmath import fdfield_t, fdfield_updater_t
@ -12,7 +12,7 @@ from ..fdmath import fdfield_t, fdfield_updater_t
def conducting_boundary(
direction: int,
polarity: int
) -> Tuple[fdfield_updater_t, fdfield_updater_t]:
) -> tuple[fdfield_updater_t, fdfield_updater_t]:
dirs = [0, 1, 2]
if direction not in dirs:
raise Exception('Invalid direction: {}'.format(direction))
@ -20,8 +20,8 @@ def conducting_boundary(
u, v = dirs
if polarity < 0:
boundary_slice = [slice(None)] * 3 # type: List[Any]
shifted1_slice = [slice(None)] * 3 # type: List[Any]
boundary_slice = [slice(None)] * 3 # type: list[Any]
shifted1_slice = [slice(None)] * 3 # type: list[Any]
boundary_slice[direction] = 0
shifted1_slice[direction] = 1
@ -42,7 +42,7 @@ def conducting_boundary(
if polarity > 0:
boundary_slice = [slice(None)] * 3
shifted1_slice = [slice(None)] * 3
shifted2_slice = [slice(None)] * 3 # type: List[Any]
shifted2_slice = [slice(None)] * 3 # type: list[Any]
boundary_slice[direction] = -1
shifted1_slice[direction] = -2
shifted2_slice[direction] = -3

49
meanas/fdtd/energy.py
View File

@ -1,4 +1,3 @@
from typing import Optional, Union
import numpy
from ..fdmath import dx_lists_t, fdfield_t
@ -11,7 +10,7 @@ from ..fdmath.functional import deriv_back
def poynting(
e: fdfield_t,
h: fdfield_t,
dxes: Optional[dx_lists_t] = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Calculate the poynting vector `S` ($S$).
@ -89,11 +88,11 @@ def poynting(
def poynting_divergence(
s: Optional[fdfield_t] = None,
s: fdfield_t | None = None,
*,
e: Optional[fdfield_t] = None,
h: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None,
e: fdfield_t | None = None,
h: fdfield_t | None = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Calculate the divergence of the poynting vector.
@ -116,9 +115,9 @@ def poynting_divergence(
energy cell.
"""
if s is None:
assert(e is not None)
assert(h is not None)
assert(dxes is not None)
assert e is not None
assert h is not None
assert dxes is not None
s = poynting(e, h, dxes=dxes)
Dx, Dy, Dz = deriv_back()
@ -130,9 +129,9 @@ def energy_hstep(
e0: fdfield_t,
h1: fdfield_t,
e2: fdfield_t,
epsilon: Optional[fdfield_t] = None,
mu: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None,
epsilon: fdfield_t | None = None,
mu: fdfield_t | None = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Calculate energy `U` at the time of the provided H-field `h1`.
@ -158,9 +157,9 @@ def energy_estep(
h0: fdfield_t,
e1: fdfield_t,
h2: fdfield_t,
epsilon: Optional[fdfield_t] = None,
mu: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None,
epsilon: fdfield_t | None = None,
mu: fdfield_t | None = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Calculate energy `U` at the time of the provided E-field `e1`.
@ -188,9 +187,9 @@ def delta_energy_h2e(
h1: fdfield_t,
e2: fdfield_t,
h3: fdfield_t,
epsilon: Optional[fdfield_t] = None,
mu: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None,
epsilon: fdfield_t | None = None,
mu: fdfield_t | None = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Change in energy during the half-step from `h1` to `e2`.
@ -221,9 +220,9 @@ def delta_energy_e2h(
e1: fdfield_t,
h2: fdfield_t,
e3: fdfield_t,
epsilon: Optional[fdfield_t] = None,
mu: Optional[fdfield_t] = None,
dxes: Optional[dx_lists_t] = None,
epsilon: fdfield_t | None = None,
mu: fdfield_t | None = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Change in energy during the half-step from `e1` to `h2`.
@ -251,7 +250,7 @@ def delta_energy_e2h(
def delta_energy_j(
j0: fdfield_t,
e1: fdfield_t,
dxes: Optional[dx_lists_t] = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
"""
Calculate
@ -274,9 +273,9 @@ def delta_energy_j(
def dxmul(
ee: fdfield_t,
hh: fdfield_t,
epsilon: Optional[Union[fdfield_t, float]] = None,
mu: Optional[Union[fdfield_t, float]] = None,
dxes: Optional[dx_lists_t] = None,
epsilon: fdfield_t | float | None = None,
mu: fdfield_t | float | None = None,
dxes: dx_lists_t | None = None,
) -> fdfield_t:
if epsilon is None:
epsilon = 1

23
meanas/fdtd/pml.py
View File

@ -7,7 +7,7 @@ PML implementations
"""
# TODO retest pmls!
from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional
from typing import Callable, Sequence, Any
from copy import deepcopy
import numpy
from numpy.typing import NDArray, DTypeLike
@ -30,7 +30,7 @@ def cpml_params(
m: float = 3.5,
ma: float = 1,
cfs_alpha: float = 0,
) -> Dict[str, Any]:
) -> dict[str, Any]:
if axis not in range(3):
raise Exception('Invalid axis: {}'.format(axis))
@ -59,11 +59,11 @@ def cpml_params(
else:
raise Exception('Bad polarity!')
expand_slice_l: List[Any] = [None, None, None]
expand_slice_l: list[Any] = [None, None, None]
expand_slice_l[axis] = slice(None)
expand_slice = tuple(expand_slice_l)
def par(x: NDArray[numpy.float64]) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
def par(x: NDArray[numpy.float64]) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
scaling = (x / thickness) ** m
sigma = scaling * sigma_max
kappa = 1 + scaling * (kappa_max - 1)
@ -93,23 +93,22 @@ def cpml_params(
def updates_with_cpml(
cpml_params: Sequence[Sequence[Optional[Dict[str, Any]]]],
cpml_params: Sequence[Sequence[dict[str, Any] | None]],
dt: float,
dxes: dx_lists_t,
epsilon: fdfield_t,
*,
dtype: DTypeLike = numpy.float32,
) -> Tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None],
) -> tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None],
Callable[[fdfield_t, fdfield_t, fdfield_t], None]]:
Dfx, Dfy, Dfz = deriv_forward(dxes[1])
Dbx, Dby, Dbz = deriv_back(dxes[1])
psi_E: List[List[Tuple[Any, Any]]] = [[(None, None) for _ in range(2)] for _ in range(3)]
psi_H: List[List[Tuple[Any, Any]]] = deepcopy(psi_E)
params_E: List[List[Tuple[Any, Any, Any, Any]]] = [[(None, None, None, None) for _ in range(2)] for _ in range(3)]
params_H: List[List[Tuple[Any, Any, Any, Any]]] = deepcopy(params_E)
psi_E: list[list[tuple[Any, Any]]] = [[(None, None) for _ in range(2)] for _ in range(3)]
psi_H: list[list[tuple[Any, Any]]] = deepcopy(psi_E)
params_E: list[list[tuple[Any, Any, Any, Any]]] = [[(None, None, None, None) for _ in range(2)] for _ in range(3)]
params_H: list[list[tuple[Any, Any, Any, Any]]] = deepcopy(params_E)
for axis in range(3):
for pp, polarity in enumerate((-1, 1)):
@ -133,7 +132,6 @@ def updates_with_cpml(
params_E[axis][pp] = cpml_param['param_e'] + (region,)
params_H[axis][pp] = cpml_param['param_h'] + (region,)
pE = numpy.empty_like(epsilon, dtype=dtype)
pH = numpy.empty_like(epsilon, dtype=dtype)
@ -183,7 +181,6 @@ def updates_with_cpml(
e[1] += dt / epsilon[1] * (dzHx - dxHz + pE[1])
e[2] += dt / epsilon[2] * (dxHy - dyHx + pE[2])
def update_H(
e: fdfield_t,
h: fdfield_t,

12
meanas/test/conftest.py
View File

@ -3,9 +3,9 @@
Test fixtures
"""
from typing import Tuple, Iterable, List, Any
from typing import Iterable, Any
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.typing import NDArray
import pytest # type: ignore
from .utils import PRNG
@ -20,7 +20,7 @@ FixtureRequest = Any
(5, 5, 5),
# (7, 7, 7),
])
def shape(request: FixtureRequest) -> Iterable[Tuple[int, ...]]:
def shape(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
yield (3, *request.param)
@ -37,7 +37,7 @@ def epsilon_fg(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(scope='module', params=['center', '000', 'random'])
def epsilon(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
epsilon_bg: float,
epsilon_fg: float,
) -> Iterable[NDArray[numpy.float64]]:
@ -76,9 +76,9 @@ def dx(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(scope='module', params=['uniform', 'centerbig'])
def dxes(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
dx: float,
) -> Iterable[List[List[NDArray[numpy.float64]]]]:
) -> Iterable[list[list[NDArray[numpy.float64]]]]:
if request.param == 'uniform':
dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
elif request.param == 'centerbig':

26
meanas/test/test_fdfd.py
View File

@ -1,8 +1,8 @@
from typing import List, Tuple, Iterable, Optional
from typing import Iterable
import dataclasses
import pytest # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.typing import NDArray
#from numpy.testing import assert_allclose, assert_array_equal
from .. import fdfd
@ -60,12 +60,12 @@ def omega(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(params=[None])
def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
def pec(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
yield request.param
@pytest.fixture(params=[None])
def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
def pmc(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
yield request.param
@ -78,7 +78,7 @@ def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
@pytest.fixture(params=['diag']) # 'center'
def j_distribution(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
j_mag: float,
) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape, dtype=complex)
@ -95,26 +95,26 @@ def j_distribution(
@dataclasses.dataclass()
class FDResult:
shape: Tuple[int, ...]
dxes: List[List[NDArray[numpy.float64]]]
shape: tuple[int, ...]
dxes: list[list[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64]
omega: complex
j: NDArray[numpy.complex128]
e: NDArray[numpy.complex128]
pmc: Optional[NDArray[numpy.float64]]
pec: Optional[NDArray[numpy.float64]]
pmc: NDArray[numpy.float64] | None
pec: NDArray[numpy.float64] | None
@pytest.fixture()
def sim(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]],
dxes: list[list[NDArray[numpy.float64]]],
j_distribution: NDArray[numpy.complex128],
omega: float,
pec: Optional[NDArray[numpy.float64]],
pmc: Optional[NDArray[numpy.float64]],
pec: NDArray[numpy.float64] | None,
pmc: NDArray[numpy.float64] | None,
) -> FDResult:
"""
Build simulation from parts

26
meanas/test/test_fdfd_pml.py
View File

@ -1,7 +1,7 @@
from typing import Optional, Tuple, Iterable, List
from typing import Iterable
import pytest # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.typing import NDArray
from numpy.testing import assert_allclose
from .. import fdfd
@ -49,19 +49,19 @@ def omega(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(params=[None])
def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
def pec(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
yield request.param
@pytest.fixture(params=[None])
def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
def pmc(request: FixtureRequest) -> Iterable[NDArray[numpy.float64] | None]:
yield request.param
@pytest.fixture(params=[(30, 1, 1),
(1, 30, 1),
(1, 1, 30)])
def shape(request: FixtureRequest) -> Iterable[Tuple[int, ...]]:
def shape(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
yield (3, *request.param)
@ -73,7 +73,7 @@ def src_polarity(request: FixtureRequest) -> Iterable[int]:
@pytest.fixture()
def j_distribution(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut,
omega: float,
@ -86,7 +86,7 @@ def j_distribution(
other_dims.remove(dim)
dx_prop = (dxes[0][dim][shape[dim + 1] // 2]
+ dxes[1][dim][shape[dim + 1] // 2]) / 2 # TODO is this right for nonuniform dxes?
+ dxes[1][dim][shape[dim + 1] // 2]) / 2 # noqa: E128 # TODO is this right for nonuniform dxes?
# Mask only contains components orthogonal to propagation direction
center_mask = numpy.zeros(shape, dtype=bool)
@ -112,7 +112,7 @@ def j_distribution(
@pytest.fixture()
def epsilon(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
epsilon_bg: float,
epsilon_fg: float,
) -> Iterable[NDArray[numpy.float64]]:
@ -123,11 +123,11 @@ def epsilon(
@pytest.fixture(params=['uniform'])
def dxes(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
dx: float,
omega: float,
epsilon_fg: float,
) -> Iterable[List[List[NDArray[numpy.float64]]]]:
) -> Iterable[list[list[NDArray[numpy.float64]]]]:
if request.param == 'uniform':
dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
@ -147,13 +147,13 @@ def dxes(
@pytest.fixture()
def sim(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut,
j_distribution: NDArray[numpy.complex128],
omega: float,
pec: Optional[NDArray[numpy.float64]],
pmc: Optional[NDArray[numpy.float64]],
pec: NDArray[numpy.float64] | None,
pmc: NDArray[numpy.float64] | None,
) -> FDResult:
j_vec = vec(j_distribution)
eps_vec = vec(epsilon)

34
meanas/test/test_fdtd.py
View File

@ -1,8 +1,8 @@
from typing import List, Tuple, Iterable, Any, Dict
from typing import Iterable, Any
import dataclasses
import pytest # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from numpy.typing import NDArray
#from numpy.testing import assert_allclose, assert_array_equal
from .. import fdtd
@ -33,7 +33,7 @@ def test_initial_energy(sim: 'TDResult') -> None:
dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0)
u0 = (j0 * j0.conj() / sim.epsilon * dV).sum(axis=0)
args: Dict[str, Any] = {
args: dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
@ -52,7 +52,7 @@ def test_energy_conservation(sim: 'TDResult') -> None:
e0 = sim.es[0]
j0 = sim.js[0]
u = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes).sum()
args: Dict[str, Any] = {
args: dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
@ -70,7 +70,7 @@ def test_energy_conservation(sim: 'TDResult') -> None:
def test_poynting_divergence(sim: 'TDResult') -> None:
args: Dict[str, Any] = {
args: dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
@ -103,7 +103,7 @@ def test_poynting_planes(sim: 'TDResult') -> None:
if mask.sum() > 1:
pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum()))
args: Dict[str, Any] = {
args: dict[str, Any] = {
'dxes': sim.dxes,
'epsilon': sim.epsilon,
}
@ -156,26 +156,26 @@ def dt(request: FixtureRequest) -> Iterable[float]:
@dataclasses.dataclass()
class TDResult:
shape: Tuple[int, ...]
shape: tuple[int, ...]
dt: float
dxes: List[List[NDArray[numpy.float64]]]
dxes: list[list[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64]
j_distribution: NDArray[numpy.float64]
j_steps: Tuple[int, ...]
es: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
hs: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
js: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
j_steps: tuple[int, ...]
es: list[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
hs: list[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
js: list[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
@pytest.fixture(params=[(0, 4, 8)]) # (0,)
def j_steps(request: FixtureRequest) -> Iterable[Tuple[int, ...]]:
def j_steps(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
yield request.param
@pytest.fixture(params=['center', 'random'])
def j_distribution(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
j_mag: float,
) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape)
@ -191,12 +191,12 @@ def j_distribution(
@pytest.fixture()
def sim(
request: FixtureRequest,
shape: Tuple[int, ...],
shape: tuple[int, ...],
epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]],
dxes: list[list[NDArray[numpy.float64]]],
dt: float,
j_distribution: NDArray[numpy.float64],
j_steps: Tuple[int, ...],
j_steps: tuple[int, ...],
) -> TDResult:
is3d = (numpy.array(shape) == 1).sum() == 0
if is3d:

4
meanas/test/utils.py
View File

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