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 Solvers for eigenvalue / eigenvector problems
""" """
from typing import Tuple, Callable, Optional, Union from typing import Callable
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm from numpy.linalg import norm
@ -11,9 +11,9 @@ import scipy.sparse.linalg as spalg # type: ignore
def power_iteration( def power_iteration(
operator: sparse.spmatrix, operator: sparse.spmatrix,
guess_vector: Optional[NDArray[numpy.complex128]] = None, guess_vector: NDArray[numpy.complex128] | None = None,
iterations: int = 20, iterations: int = 20,
) -> Tuple[complex, NDArray[numpy.complex128]]: ) -> tuple[complex, NDArray[numpy.complex128]]:
""" """
Use power iteration to estimate the dominant eigenvector of a matrix. Use power iteration to estimate the dominant eigenvector of a matrix.
@ -40,12 +40,12 @@ def power_iteration(
def rayleigh_quotient_iteration( def rayleigh_quotient_iteration(
operator: Union[sparse.spmatrix, spalg.LinearOperator], operator: sparse.spmatrix | spalg.LinearOperator,
guess_vector: NDArray[numpy.complex128], guess_vector: NDArray[numpy.complex128],
iterations: int = 40, iterations: int = 40,
tolerance: float = 1e-13, tolerance: float = 1e-13,
solver: Optional[Callable[..., NDArray[numpy.complex128]]] = None, solver: Callable[..., NDArray[numpy.complex128]] | None = None,
) -> Tuple[complex, NDArray[numpy.complex128]]: ) -> tuple[complex, NDArray[numpy.complex128]]:
""" """
Use Rayleigh quotient iteration to refine an eigenvector guess. Use Rayleigh quotient iteration to refine an eigenvector guess.
@ -80,7 +80,7 @@ def rayleigh_quotient_iteration(
if solver is None: if solver is None:
def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.complex128]: def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.complex128]:
return spalg.bicgstab(A, b)[0] return spalg.bicgstab(A, b)[0]
assert(solver is not None) assert solver is not None
v = numpy.squeeze(guess_vector) v = numpy.squeeze(guess_vector)
v /= norm(v) v /= norm(v)
@ -96,10 +96,10 @@ def rayleigh_quotient_iteration(
def signed_eigensolve( def signed_eigensolve(
operator: Union[sparse.spmatrix, spalg.LinearOperator], operator: sparse.spmatrix | spalg.LinearOperator,
how_many: int, how_many: int,
negative: bool = False, 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 Find the largest-magnitude positive-only (or negative-only) eigenvalues and
eigenvectors of the provided matrix. 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 logging
import numpy import numpy
from numpy import pi, real, trace from numpy import pi, real, trace
@ -140,7 +140,7 @@ def generate_kmn(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
shape: Sequence[int], 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. 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)`). All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
""" """
k0 = numpy.array(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) Gi = numpy.moveaxis(Gi_grids, 0, -1)
k_G = k0[None, None, None, :] - Gi k_G = k0[None, None, None, :] - Gi
@ -183,7 +184,7 @@ def maxwell_operator(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None mu: fdfield_t | None = None
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
""" """
Generate the Maxwell operator Generate the Maxwell operator
@ -237,7 +238,7 @@ def maxwell_operator(
# cross product and transform into xyz basis # cross product and transform into xyz basis
d_xyz = (n * hin_m d_xyz = (n * hin_m
- m * hin_n) * k_mag - m * hin_n) * k_mag # noqa: E128
# divide by epsilon # divide by epsilon
temp = ifftn(d_xyz, axes=range(3)) # reuses d_xyz if using pyfftw temp = ifftn(d_xyz, axes=range(3)) # reuses d_xyz if using pyfftw
@ -253,7 +254,7 @@ def maxwell_operator(
else: else:
# transform from mn to xyz # transform from mn to xyz
b_xyz = (m * b_m[:, :, :, None] b_xyz = (m * b_m[:, :, :, None]
+ n * b_n[:, :, :, None]) + n * b_n[:, :, :, None]) # noqa: E128
# divide by mu # divide by mu
temp = ifftn(b_xyz, axes=range(3)) temp = ifftn(b_xyz, axes=range(3))
@ -302,7 +303,7 @@ def hmn_2_exyz(
def operator(h: NDArray[numpy.complex128]) -> cfdfield_t: def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)] hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
d_xyz = (n * hin_m d_xyz = (n * hin_m
- m * hin_n) * k_mag - m * hin_n) * k_mag # noqa: E128
# divide by epsilon # divide by epsilon
return numpy.array([ei for ei in numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)]) # TODO avoid copy 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: def operator(h: NDArray[numpy.complex128]) -> cfdfield_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) # noqa: E128
return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)]) return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])
return operator return operator
@ -350,7 +351,7 @@ def inverse_maxwell_operator_approx(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
""" """
Generate an approximate inverse of the Maxwell operator, Generate an approximate inverse of the Maxwell operator,
@ -400,7 +401,7 @@ def inverse_maxwell_operator_approx(
else: else:
# transform from mn to xyz # transform from mn to xyz
h_xyz = (m * hin_m[:, :, :, None] h_xyz = (m * hin_m[:, :, :, None]
+ n * hin_n[:, :, :, None]) + n * hin_n[:, :, :, None]) # noqa: E128
# multiply by mu # multiply by mu
temp = ifftn(h_xyz, axes=range(3)) temp = ifftn(h_xyz, axes=range(3))
@ -413,7 +414,7 @@ def inverse_maxwell_operator_approx(
# cross product and transform into xyz basis # cross product and transform into xyz basis
e_xyz = (n * b_m e_xyz = (n * b_m
- m * b_n) / k_mag - m * b_n) / k_mag # noqa: E128
# multiply by epsilon # multiply by epsilon
temp = ifftn(e_xyz, axes=range(3)) temp = ifftn(e_xyz, axes=range(3))
@ -436,13 +437,13 @@ def find_k(
direction: ArrayLike, direction: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
band: int = 0, band: int = 0,
k_bounds: Tuple[float, float] = (0, 0.5), k_bounds: tuple[float, float] = (0, 0.5),
k_guess: Optional[float] = None, k_guess: float | None = None,
solve_callback: Optional[Callable[..., None]] = None, solve_callback: Callable[..., None] | None = None,
iter_callback: Optional[Callable[..., None]] = None, iter_callback: Callable[..., None] | None = None,
) -> Tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]: ) -> tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
""" """
Search for a bloch vector that has a given frequency. Search for a bloch vector that has a given frequency.
@ -503,13 +504,13 @@ def eigsolve(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
tolerance: float = 1e-20, tolerance: float = 1e-20,
max_iters: int = 10000, max_iters: int = 10000,
reset_iters: int = 100, reset_iters: int = 100,
y0: Optional[ArrayLike] = None, y0: ArrayLike | None = None,
callback: Optional[Callable[..., None]] = None, callback: Callable[..., None] | None = None,
) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
""" """
Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
k0 of the specified structure. k0 of the specified structure.
@ -555,6 +556,7 @@ def eigsolve(
#prev_theta = 0.5 #prev_theta = 0.5
D = numpy.zeros(shape=y_shape, dtype=complex) D = numpy.zeros(shape=y_shape, dtype=complex)
Z: NDArray[numpy.complex128]
if y0 is None: if y0 is None:
Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape) Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
else: else:
@ -593,7 +595,8 @@ def eigsolve(
# via (1 - ZUZt) # via (1 - ZUZt)
if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7): 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'[change in trace] {abs(E - prev_E)} - 5e-8 '
f'< {tolerance} [tolerance] * {(E + prev_E) / 2} [value of trace]' f'< {tolerance} [tolerance] * {(E + prev_E) / 2} [value of trace]'
) )
@ -635,7 +638,7 @@ def eigsolve(
symZtD = _symmetrize(Zt @ D) symZtD = _symmetrize(Zt @ D)
symZtAD = _symmetrize(Zt @ AD) symZtAD = _symmetrize(Zt @ AD)
Qi_memo: List[Optional[float]] = [None, None] Qi_memo: list[float | None] = [None, None]
def Qi_func(theta: float) -> float: def Qi_func(theta: float) -> float:
nonlocal Qi_memo nonlocal Qi_memo
@ -659,8 +662,8 @@ def eigsolve(
else: else:
raise Exception('Inexplicable singularity in trace_func') raise Exception('Inexplicable singularity in trace_func')
Qi_memo[0] = theta Qi_memo[0] = theta
Qi_memo[1] = Qi Qi_memo[1] = cast(float, Qi)
return Qi return cast(float, Qi)
def trace_func(theta: float) -> float: def trace_func(theta: float) -> float:
c = numpy.cos(theta) 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( def _rtrace_AtB(
A: NDArray[numpy.complex128], A: NDArray[numpy.complex128],
B: Union[NDArray[numpy.complex128], float], B: NDArray[numpy.complex128] | float,
) -> float: ) -> float:
return real(numpy.sum(A.conj() * B)) 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). 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 import numpy
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
from numpy import pi from numpy import pi
@ -14,8 +14,8 @@ def near_to_farfield(
H_near: cfdfield_t, H_near: cfdfield_t,
dx: float, dx: float,
dy: float, dy: float,
padded_size: Union[List[int], int, None] = None padded_size: list[int] | int | None = None
) -> Dict[str, Any]: ) -> 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.
@ -62,14 +62,15 @@ def near_to_farfield(
padded_size = (2**numpy.ceil(numpy.log2(s))).astype(int) padded_size = (2**numpy.ceil(numpy.log2(s))).astype(int)
if not hasattr(padded_size, '__len__'): if not hasattr(padded_size, '__len__'):
padded_size = (padded_size, padded_size) # type: ignore # checked if sequence 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] En_fft = [fftshift(fft2(fftshift(Eni), s=padded_shape)) for Eni in E_near]
Hn_fft = [fftshift(fft2(fftshift(Hni), s=padded_size)) for Hni in H_near] Hn_fft = [fftshift(fft2(fftshift(Hni), s=padded_shape)) for Hni in H_near]
# Propagation vectors kx, ky # Propagation vectors kx, ky
k = 2 * pi k = 2 * pi
kxx = 2 * pi * fftshift(fftfreq(padded_size[0], dx)) kxx = 2 * pi * fftshift(fftfreq(padded_shape[0], dx))
kyy = 2 * pi * fftshift(fftfreq(padded_size[1], dy)) kyy = 2 * pi * fftshift(fftfreq(padded_shape[1], dy))
kx, ky = numpy.meshgrid(kxx, kyy, indexing='ij') kx, ky = numpy.meshgrid(kxx, kyy, indexing='ij')
kxy2 = kx * kx + ky * ky kxy2 = kx * kx + ky * ky
@ -85,14 +86,14 @@ def near_to_farfield(
# Normalized vector potentials N, L # Normalized vector potentials N, L
N = [-Hn_fft[1] * cos_phi * cos_th + Hn_fft[0] * cos_phi * sin_th, 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, 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], E_far = [-L[1] - N[0],
L[0] - N[1]] L[0] - N[1]] # noqa: E127
H_far = [-E_far[1], H_far = [-E_far[1],
E_far[0]] E_far[0]] # noqa: E127
theta = numpy.arctan2(ky, kx) theta = numpy.arctan2(ky, kx)
phi = numpy.arccos(cos_phi) phi = numpy.arccos(cos_phi)
@ -126,8 +127,8 @@ def far_to_nearfield(
H_far: cfdfield_t, H_far: cfdfield_t,
dkx: float, dkx: float,
dky: float, dky: float,
padded_size: Union[List[int], int, None] = None padded_size: list[int] | int | None = None
) -> Dict[str, Any]: ) -> 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.
@ -170,6 +171,7 @@ def far_to_nearfield(
padded_size = (2 ** numpy.ceil(numpy.log2(s))).astype(int) padded_size = (2 ** numpy.ceil(numpy.log2(s))).astype(int)
if not hasattr(padded_size, '__len__'): if not hasattr(padded_size, '__len__'):
padded_size = (padded_size, padded_size) # type: ignore # checked if sequence padded_size = (padded_size, padded_size) # type: ignore # checked if sequence
padded_shape = cast(Sequence[int], padded_size)
k = 2 * pi k = 2 * pi
kxs = fftshift(fftfreq(s[0], 1 / (s[0] * dkx))) kxs = fftshift(fftfreq(s[0], 1 / (s[0] * dkx)))
@ -203,9 +205,9 @@ def far_to_nearfield(
# Normalized vector potentials N, L # Normalized vector potentials N, L
L = [0.5 * E_far[1], L = [0.5 * E_far[1],
-0.5 * E_far[0]] -0.5 * E_far[0]] # noqa: E128
N = [L[1], N = [L[1],
-L[0]] -L[0]] # noqa: E128
En_fft = [-( L[0] * sin_th + L[1] * cos_phi * cos_th) / cos_phi, 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] -(-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 En_fft[i][cos_phi == 0] = 0
Hn_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] E_near = [ifftshift(ifft2(ifftshift(Ei), s=padded_shape)) for Ei in En_fft]
H_near = [ifftshift(ifft2(ifftshift(Hi), s=padded_size)) for Hi in Hn_fft] H_near = [ifftshift(ifft2(ifftshift(Hi), s=padded_shape)) for Hi in Hn_fft]
dx = 2 * pi / (s[0] * dkx) dx = 2 * pi / (s[0] * dkx)
dy = 2 * pi / (s[0] * dky) 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), 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) 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 import numpy
from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t
@ -19,7 +19,7 @@ def e_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_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.
@ -55,8 +55,8 @@ def eh_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> Callable[[cfdfield_t, cfdfield_t], Tuple[cfdfield_t, cfdfield_t]]: ) -> Callable[[cfdfield_t, cfdfield_t], tuple[cfdfield_t, cfdfield_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`.
@ -74,11 +74,11 @@ def eh_full(
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) 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, return (ch(h) - 1j * omega * epsilon * e,
ce(e) + 1j * omega * h) 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, return (ch(h) - 1j * omega * epsilon * e,
ce(e) + 1j * omega * mu * h) # type: ignore # mu=None ok ce(e) + 1j * omega * mu * h) # type: ignore # mu=None ok
@ -91,7 +91,7 @@ def eh_full(
def e2h( def e2h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Utility operator for converting the `E` field into the `H` field. Utility operator for converting the `E` field into the `H` field.
@ -123,7 +123,7 @@ def e2h(
def m2j( def m2j(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Utility operator for converting magnetic current `M` distribution Utility operator for converting magnetic current `M` distribution
@ -160,7 +160,7 @@ def e_tfsf_source(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_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

31
meanas/fdfd/operators.py
View File

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

16
meanas/fdfd/scpml.py
View File

@ -2,10 +2,10 @@
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, List from typing import Sequence, Callable
import numpy import numpy
from numpy.typing import ArrayLike, NDArray from numpy.typing import NDArray
__author__ = 'Jan Petykiewicz' __author__ = 'Jan Petykiewicz'
@ -43,8 +43,8 @@ def uniform_grid_scpml(
thicknesses: Sequence[int], thicknesses: 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: s_function_t | None = None,
) -> List[List[NDArray[numpy.float64]]]: ) -> 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.
@ -86,7 +86,7 @@ def uniform_grid_scpml(
for k, th in enumerate(thicknesses): for k, th in enumerate(thicknesses):
s = shape[k] s = shape[k]
if th > 0: 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_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 dx_b[k] = 1 + 1j * s_function(ll(sr + 0.5, s, th)) / s_correction
else: else:
@ -96,14 +96,14 @@ def uniform_grid_scpml(
def stretch_with_scpml( def stretch_with_scpml(
dxes: List[List[NDArray[numpy.float64]]], dxes: list[list[NDArray[numpy.float64]]],
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: s_function_t | None = None,
) -> List[List[NDArray[numpy.float64]]]: ) -> 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.

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 # TODO update module docs
from typing import List, Tuple, Optional, Any from typing import Any
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray, ArrayLike
from numpy.linalg import norm from numpy.linalg import norm
@ -196,7 +196,7 @@ def operator_e(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Waveguide operator of the form Waveguide operator of the form
@ -263,7 +263,7 @@ def operator_h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Waveguide operator of the form Waveguide operator of the form
@ -333,9 +333,9 @@ def normalized_fields_e(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
prop_phase: float = 0, 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, 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.
@ -368,9 +368,9 @@ def normalized_fields_h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
prop_phase: float = 0, 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, 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.
@ -403,9 +403,9 @@ def _normalized_fields(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> Tuple[vcfdfield_t, vcfdfield_t]: ) -> tuple[vcfdfield_t, vcfdfield_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)]
@ -445,7 +445,7 @@ def exy2h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None mu: vfdfield_t | None = None
) -> sparse.spmatrix: ) -> 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,
@ -471,7 +471,7 @@ def hxy2e(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None mu: vfdfield_t | None = None
) -> sparse.spmatrix: ) -> 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,
@ -495,7 +495,7 @@ def hxy2e(
def hxy2h( def hxy2h(
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None mu: vfdfield_t | None = None
) -> sparse.spmatrix: ) -> 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,
@ -564,7 +564,7 @@ def e2h(
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[vfdfield_t] = None mu: vfdfield_t | None = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Returns an operator which, when applied to a vectorized E eigenfield, produces Returns an operator which, when applied to a vectorized E eigenfield, produces
@ -654,7 +654,7 @@ def h_err(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None mu: vfdfield_t | None = None
) -> float: ) -> float:
""" """
Calculates the relative error in the H field Calculates the relative error in the H field
@ -680,7 +680,7 @@ def h_err(
else: else:
op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h) op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h)
return norm(op) / norm(h) return float(norm(op) / norm(h))
def e_err( def e_err(
@ -689,7 +689,7 @@ def e_err(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
) -> float: ) -> float:
""" """
Calculates the relative error in the E field Calculates the relative error in the E field
@ -714,17 +714,17 @@ def e_err(
mu_inv = sparse.diags(1 / mu) mu_inv = sparse.diags(1 / mu)
op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e) op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
return norm(op) / norm(e) return float(norm(op) / norm(e))
def solve_modes( def solve_modes(
mode_numbers: List[int], mode_numbers: list[int],
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: vfdfield_t | None = None,
mode_margin: int = 2, 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. 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, mode_number: int,
*args: Any, *args: Any,
**kwargs: Any, **kwargs: Any,
) -> Tuple[vcfdfield_t, complex]: ) -> tuple[vcfdfield_t, complex]:
""" """
Wrapper around `solve_modes()` that solves for a single mode. 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 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 Sequence, Any
import numpy import numpy
from numpy.typing import NDArray from numpy.typing import NDArray
@ -20,8 +20,8 @@ def solve_mode(
polarity: int, polarity: int,
slices: Sequence[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> Dict[str, Union[complex, NDArray[numpy.float_]]]: ) -> dict[str, 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.
@ -40,8 +40,8 @@ def solve_mode(
Returns: Returns:
``` ```
{ {
'E': List[NDArray[numpy.float_]], 'E': list[NDArray[numpy.float_]],
'H': List[NDArray[numpy.float_]], 'H': list[NDArray[numpy.float_]],
'wavenumber': complex, 'wavenumber': complex,
} }
``` ```
@ -63,7 +63,7 @@ def solve_mode(
dx_prop = 0.5 * dxab_forward.sum() dx_prop = 0.5 * dxab_forward.sum()
# Reduce to 2D and solve the 2D problem # Reduce to 2D and solve the 2D problem
args_2d: Dict[str, Any] = { args_2d: dict[str, Any] = {
'omega': omega, 'omega': omega,
'dxes': [[dx[i][slices[i]] for i in order[:2]] for dx in dxes], '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]), 'epsilon': vec([epsilon[i][slices].transpose(order) for i in order]),
@ -114,7 +114,7 @@ def compute_source(
polarity: int, polarity: int,
slices: Sequence[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: fdfield_t | None = None,
) -> cfdfield_t: ) -> cfdfield_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

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 # TODO update module docs
from typing import Dict, Union
import numpy import numpy
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
@ -85,7 +84,7 @@ def solve_mode(
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
r0: float, r0: float,
) -> Dict[str, Union[complex, cfdfield_t]]: ) -> dict[str, complex | cfdfield_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
@ -103,8 +102,8 @@ def solve_mode(
Returns: Returns:
``` ```
{ {
'E': List[NDArray[numpy.complex_]], 'E': list[NDArray[numpy.complex_]],
'H': List[NDArray[numpy.complex_]], 'H': list[NDArray[numpy.complex_]],
'wavenumber': 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: 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 + 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. 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: + 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. Basic discrete calculus etc.
""" """
from typing import Sequence, Tuple, Optional, Callable from typing import Sequence, Callable
import numpy import numpy
from numpy.typing import NDArray from numpy.typing import NDArray
@ -12,8 +12,8 @@ from .types import fdfield_t, fdfield_updater_t
def deriv_forward( def deriv_forward(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None, dx_e: Sequence[NDArray[numpy.float_]] | None = None,
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: ) -> 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).
@ -36,8 +36,8 @@ def deriv_forward(
def deriv_back( def deriv_back(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None, dx_h: Sequence[NDArray[numpy.float_]] | None = None,
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: ) -> 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).
@ -60,7 +60,7 @@ def deriv_back(
def curl_forward( def curl_forward(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None, dx_e: Sequence[NDArray[numpy.float_]] | None = None,
) -> fdfield_updater_t: ) -> fdfield_updater_t:
""" """
Curl operator for use with the E field. Curl operator for use with the E field.
@ -89,7 +89,7 @@ def curl_forward(
def curl_back( def curl_back(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None, dx_h: Sequence[NDArray[numpy.float_]] | None = None,
) -> fdfield_updater_t: ) -> 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.
@ -118,11 +118,11 @@ def curl_back(
def curl_forward_parts( def curl_forward_parts(
dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None, dx_e: Sequence[NDArray[numpy.float_]] | None = None,
) -> Callable: ) -> 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, fdfield_t], ...]: def mkparts_fwd(e: fdfield_t) -> tuple[tuple[fdfield_t, fdfield_t], ...]:
return ((-Dz(e[1]), Dy(e[2])), return ((-Dz(e[1]), Dy(e[2])),
( Dz(e[0]), -Dx(e[2])), ( Dz(e[0]), -Dx(e[2])),
(-Dy(e[0]), Dx(e[1]))) (-Dy(e[0]), Dx(e[1])))
@ -131,11 +131,11 @@ def curl_forward_parts(
def curl_back_parts( def curl_back_parts(
dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None, dx_h: Sequence[NDArray[numpy.float_]] | None = None,
) -> Callable: ) -> Callable:
Dx, Dy, Dz = deriv_back(dx_h) 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])), return ((-Dz(h[1]), Dy(h[2])),
( Dz(h[0]), -Dx(h[2])), ( Dz(h[0]), -Dx(h[2])),
(-Dy(h[0]), Dx(h[1]))) (-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. Basic discrete calculus etc.
""" """
from typing import Sequence, List from typing import Sequence
import numpy import numpy
from numpy.typing import NDArray from numpy.typing import NDArray
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
@ -98,7 +98,7 @@ def shift_with_mirror(
def deriv_forward( def deriv_forward(
dx_e: Sequence[NDArray[numpy.float_]], dx_e: Sequence[NDArray[numpy.float_]],
) -> List[sparse.spmatrix]: ) -> list[sparse.spmatrix]:
""" """
Utility operators for taking discretized derivatives (forward variant). Utility operators for taking discretized derivatives (forward variant).
@ -125,7 +125,7 @@ def deriv_forward(
def deriv_back( def deriv_back(
dx_h: Sequence[NDArray[numpy.float_]], dx_h: Sequence[NDArray[numpy.float_]],
) -> List[sparse.spmatrix]: ) -> list[sparse.spmatrix]:
""" """
Utility operators for taking discretized derivatives (backward variant). 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. 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 import numpy
from numpy.typing import ArrayLike from numpy.typing import ArrayLike
@ -24,10 +24,10 @@ def vec(f: cfdfield_t) -> vcfdfield_t:
pass pass
@overload @overload
def vec(f: ArrayLike) -> Union[vfdfield_t, vcfdfield_t]: def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t:
pass 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. 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: def unvec(v: vcfdfield_t, shape: Sequence[int]) -> cfdfield_t:
pass 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 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

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 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
@ -14,7 +12,7 @@ __author__ = 'Jan Petykiewicz'
def maxwell_e( def maxwell_e(
dt: float, dt: float,
dxes: Optional[dx_lists_t] = None, dxes: dx_lists_t | None = None,
) -> fdfield_updater_t: ) -> fdfield_updater_t:
""" """
Build a function which performs a portion the time-domain E-field update, Build a function which performs a portion the time-domain E-field update,
@ -49,7 +47,7 @@ def maxwell_e(
else: else:
curl_h_fun = curl_back() 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. Update the E-field.
@ -69,7 +67,7 @@ def maxwell_e(
def maxwell_h( def maxwell_h(
dt: float, dt: float,
dxes: Optional[dx_lists_t] = None, dxes: dx_lists_t | None = None,
) -> fdfield_updater_t: ) -> fdfield_updater_t:
""" """
Build a function which performs part of the time-domain H-field update, Build a function which performs part of the time-domain H-field update,
@ -105,7 +103,7 @@ def maxwell_h(
else: else:
curl_e_fun = curl_forward() 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. Update the H-field.

10
meanas/fdtd/boundaries.py
View File

@ -4,7 +4,7 @@ Boundary conditions
#TODO conducting boundary documentation #TODO conducting boundary documentation
""" """
from typing import Tuple, Any, List from typing import Any
from ..fdmath import fdfield_t, fdfield_updater_t from ..fdmath import fdfield_t, fdfield_updater_t
@ -12,7 +12,7 @@ from ..fdmath import fdfield_t, fdfield_updater_t
def conducting_boundary( def conducting_boundary(
direction: int, direction: int,
polarity: int polarity: int
) -> Tuple[fdfield_updater_t, fdfield_updater_t]: ) -> 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))
@ -20,8 +20,8 @@ def conducting_boundary(
u, v = dirs u, v = dirs
if polarity < 0: if polarity < 0:
boundary_slice = [slice(None)] * 3 # type: List[Any] boundary_slice = [slice(None)] * 3 # type: list[Any]
shifted1_slice = [slice(None)] * 3 # type: List[Any] shifted1_slice = [slice(None)] * 3 # type: list[Any]
boundary_slice[direction] = 0 boundary_slice[direction] = 0
shifted1_slice[direction] = 1 shifted1_slice[direction] = 1
@ -42,7 +42,7 @@ def conducting_boundary(
if polarity > 0: if polarity > 0:
boundary_slice = [slice(None)] * 3 boundary_slice = [slice(None)] * 3
shifted1_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 boundary_slice[direction] = -1
shifted1_slice[direction] = -2 shifted1_slice[direction] = -2
shifted2_slice[direction] = -3 shifted2_slice[direction] = -3

49
meanas/fdtd/energy.py
View File

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

23
meanas/fdtd/pml.py
View File

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

12
meanas/test/conftest.py
View File

@ -3,9 +3,9 @@
Test fixtures Test fixtures
""" """
from typing import Tuple, Iterable, List, Any from typing import Iterable, Any
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray
import pytest # type: ignore import pytest # type: ignore
from .utils import PRNG from .utils import PRNG
@ -20,7 +20,7 @@ FixtureRequest = Any
(5, 5, 5), (5, 5, 5),
# (7, 7, 7), # (7, 7, 7),
]) ])
def shape(request: FixtureRequest) -> Iterable[Tuple[int, ...]]: def shape(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
yield (3, *request.param) yield (3, *request.param)
@ -37,7 +37,7 @@ 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( def epsilon(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
epsilon_bg: float, epsilon_bg: float,
epsilon_fg: float, epsilon_fg: float,
) -> Iterable[NDArray[numpy.float64]]: ) -> Iterable[NDArray[numpy.float64]]:
@ -76,9 +76,9 @@ def dx(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(scope='module', params=['uniform', 'centerbig']) @pytest.fixture(scope='module', params=['uniform', 'centerbig'])
def dxes( def dxes(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
dx: float, dx: float,
) -> Iterable[List[List[NDArray[numpy.float64]]]]: ) -> 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':

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 dataclasses
import pytest # type: ignore import pytest # type: ignore
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray
#from numpy.testing import assert_allclose, assert_array_equal #from numpy.testing import assert_allclose, assert_array_equal
from .. import fdfd from .. import fdfd
@ -60,12 +60,12 @@ def omega(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(params=[None]) @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 yield request.param
@pytest.fixture(params=[None]) @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 yield request.param
@ -78,7 +78,7 @@ def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]:
@pytest.fixture(params=['diag']) # 'center' @pytest.fixture(params=['diag']) # 'center'
def j_distribution( def j_distribution(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
j_mag: float, j_mag: float,
) -> Iterable[NDArray[numpy.float64]]: ) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape, dtype=complex) j = numpy.zeros(shape, dtype=complex)
@ -95,26 +95,26 @@ def j_distribution(
@dataclasses.dataclass() @dataclasses.dataclass()
class FDResult: class FDResult:
shape: Tuple[int, ...] shape: tuple[int, ...]
dxes: List[List[NDArray[numpy.float64]]] dxes: list[list[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64] epsilon: NDArray[numpy.float64]
omega: complex omega: complex
j: NDArray[numpy.complex128] j: NDArray[numpy.complex128]
e: NDArray[numpy.complex128] e: NDArray[numpy.complex128]
pmc: Optional[NDArray[numpy.float64]] pmc: NDArray[numpy.float64] | None
pec: Optional[NDArray[numpy.float64]] pec: NDArray[numpy.float64] | None
@pytest.fixture() @pytest.fixture()
def sim( def sim(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]], dxes: list[list[NDArray[numpy.float64]]],
j_distribution: NDArray[numpy.complex128], j_distribution: NDArray[numpy.complex128],
omega: float, omega: float,
pec: Optional[NDArray[numpy.float64]], pec: NDArray[numpy.float64] | None,
pmc: Optional[NDArray[numpy.float64]], pmc: NDArray[numpy.float64] | None,
) -> FDResult: ) -> FDResult:
""" """
Build simulation from parts 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 pytest # type: ignore
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray
from numpy.testing import assert_allclose from numpy.testing import assert_allclose
from .. import fdfd from .. import fdfd
@ -49,19 +49,19 @@ def omega(request: FixtureRequest) -> Iterable[float]:
@pytest.fixture(params=[None]) @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 yield request.param
@pytest.fixture(params=[None]) @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 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: FixtureRequest) -> Iterable[Tuple[int, ...]]: def shape(request: FixtureRequest) -> Iterable[tuple[int, ...]]:
yield (3, *request.param) yield (3, *request.param)
@ -73,7 +73,7 @@ def src_polarity(request: FixtureRequest) -> Iterable[int]:
@pytest.fixture() @pytest.fixture()
def j_distribution( def j_distribution(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut, dxes: dx_lists_mut,
omega: float, omega: float,
@ -86,7 +86,7 @@ def j_distribution(
other_dims.remove(dim) other_dims.remove(dim)
dx_prop = (dxes[0][dim][shape[dim + 1] // 2] 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 # Mask only contains components orthogonal to propagation direction
center_mask = numpy.zeros(shape, dtype=bool) center_mask = numpy.zeros(shape, dtype=bool)
@ -112,7 +112,7 @@ def j_distribution(
@pytest.fixture() @pytest.fixture()
def epsilon( def epsilon(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
epsilon_bg: float, epsilon_bg: float,
epsilon_fg: float, epsilon_fg: float,
) -> Iterable[NDArray[numpy.float64]]: ) -> Iterable[NDArray[numpy.float64]]:
@ -123,11 +123,11 @@ def epsilon(
@pytest.fixture(params=['uniform']) @pytest.fixture(params=['uniform'])
def dxes( def dxes(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
dx: float, dx: float,
omega: float, omega: float,
epsilon_fg: float, epsilon_fg: float,
) -> Iterable[List[List[NDArray[numpy.float64]]]]: ) -> 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
@ -147,13 +147,13 @@ def dxes(
@pytest.fixture() @pytest.fixture()
def sim( def sim(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut, dxes: dx_lists_mut,
j_distribution: NDArray[numpy.complex128], j_distribution: NDArray[numpy.complex128],
omega: float, omega: float,
pec: Optional[NDArray[numpy.float64]], pec: NDArray[numpy.float64] | None,
pmc: Optional[NDArray[numpy.float64]], pmc: NDArray[numpy.float64] | None,
) -> FDResult: ) -> FDResult:
j_vec = vec(j_distribution) j_vec = vec(j_distribution)
eps_vec = vec(epsilon) 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 dataclasses
import pytest # type: ignore import pytest # type: ignore
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray
#from numpy.testing import assert_allclose, assert_array_equal #from numpy.testing import assert_allclose, assert_array_equal
from .. import fdtd 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) 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: Dict[str, Any] = { args: dict[str, Any] = {
'dxes': sim.dxes, 'dxes': sim.dxes,
'epsilon': sim.epsilon, 'epsilon': sim.epsilon,
} }
@ -52,7 +52,7 @@ 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: Dict[str, Any] = { args: dict[str, Any] = {
'dxes': sim.dxes, 'dxes': sim.dxes,
'epsilon': sim.epsilon, 'epsilon': sim.epsilon,
} }
@ -70,7 +70,7 @@ def test_energy_conservation(sim: 'TDResult') -> None:
def test_poynting_divergence(sim: 'TDResult') -> None: def test_poynting_divergence(sim: 'TDResult') -> None:
args: Dict[str, Any] = { args: dict[str, Any] = {
'dxes': sim.dxes, 'dxes': sim.dxes,
'epsilon': sim.epsilon, 'epsilon': sim.epsilon,
} }
@ -103,7 +103,7 @@ 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: Dict[str, Any] = { args: dict[str, Any] = {
'dxes': sim.dxes, 'dxes': sim.dxes,
'epsilon': sim.epsilon, 'epsilon': sim.epsilon,
} }
@ -156,26 +156,26 @@ def dt(request: FixtureRequest) -> Iterable[float]:
@dataclasses.dataclass() @dataclasses.dataclass()
class TDResult: class TDResult:
shape: Tuple[int, ...] shape: tuple[int, ...]
dt: float dt: float
dxes: List[List[NDArray[numpy.float64]]] dxes: list[list[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64] epsilon: NDArray[numpy.float64]
j_distribution: NDArray[numpy.float64] j_distribution: NDArray[numpy.float64]
j_steps: Tuple[int, ...] j_steps: tuple[int, ...]
es: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list) es: list[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
hs: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list) hs: list[NDArray[numpy.float64]] = dataclasses.field(default_factory=list)
js: 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,) @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 yield request.param
@pytest.fixture(params=['center', 'random']) @pytest.fixture(params=['center', 'random'])
def j_distribution( def j_distribution(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
j_mag: float, j_mag: float,
) -> Iterable[NDArray[numpy.float64]]: ) -> Iterable[NDArray[numpy.float64]]:
j = numpy.zeros(shape) j = numpy.zeros(shape)
@ -191,12 +191,12 @@ def j_distribution(
@pytest.fixture() @pytest.fixture()
def sim( def sim(
request: FixtureRequest, request: FixtureRequest,
shape: Tuple[int, ...], shape: tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]], dxes: list[list[NDArray[numpy.float64]]],
dt: float, dt: float,
j_distribution: NDArray[numpy.float64], 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
if is3d: if is3d:

4
meanas/test/utils.py
View File

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