From faecc79179c63651fcddea41d0b15f1718b4d15f Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Tue, 4 Oct 2022 14:32:40 -0700 Subject: [PATCH] typing and formatting updates --- meanas/eigensolvers.py | 46 +++++---- meanas/fdfd/bloch.py | 20 ++-- meanas/fdfd/farfield.py | 32 +++--- meanas/fdfd/functional.py | 49 +++++---- meanas/fdfd/operators.py | 88 ++++++++-------- meanas/fdfd/scpml.py | 53 ++++++---- meanas/fdfd/solvers.py | 47 +++++---- meanas/fdfd/waveguide_2d.py | 182 ++++++++++++++++++--------------- meanas/fdfd/waveguide_3d.py | 83 ++++++++------- meanas/fdfd/waveguide_cyl.py | 34 +++--- meanas/fdmath/functional.py | 29 ++++-- meanas/fdmath/operators.py | 37 +++++-- meanas/fdmath/types.py | 39 +++---- meanas/fdmath/vectorization.py | 13 +-- meanas/fdtd/boundaries.py | 7 +- meanas/fdtd/energy.py | 114 +++++++++++---------- meanas/fdtd/pml.py | 5 +- meanas/test/conftest.py | 23 +++-- meanas/test/test_fdfd.py | 41 ++++---- meanas/test/test_fdfd_pml.py | 79 ++++++++------ meanas/test/test_fdtd.py | 61 ++++++----- meanas/test/utils.py | 25 +++-- 22 files changed, 621 insertions(+), 486 deletions(-) diff --git a/meanas/eigensolvers.py b/meanas/eigensolvers.py index 8c1739e..aa8b9ba 100644 --- a/meanas/eigensolvers.py +++ b/meanas/eigensolvers.py @@ -2,16 +2,18 @@ Solvers for eigenvalue / eigenvector problems """ from typing import Tuple, Callable, Optional, Union -import numpy # type: ignore -from numpy.linalg import norm # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike +from numpy.linalg import norm from scipy import sparse # type: ignore import scipy.sparse.linalg as spalg # type: ignore -def power_iteration(operator: sparse.spmatrix, - guess_vector: Optional[numpy.ndarray] = None, - iterations: int = 20, - ) -> Tuple[complex, numpy.ndarray]: +def power_iteration( + operator: sparse.spmatrix, + guess_vector: Optional[NDArray[numpy.float64]] = None, + iterations: int = 20, + ) -> Tuple[complex, NDArray[numpy.float64]]: """ Use power iteration to estimate the dominant eigenvector of a matrix. @@ -37,12 +39,13 @@ def power_iteration(operator: sparse.spmatrix, return lm_eigval, v -def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOperator], - guess_vector: numpy.ndarray, - iterations: int = 40, - tolerance: float = 1e-13, - solver: Optional[Callable[..., numpy.ndarray]] = None, - ) -> Tuple[complex, numpy.ndarray]: +def rayleigh_quotient_iteration( + operator: Union[sparse.spmatrix, spalg.LinearOperator], + guess_vector: NDArray[numpy.float64], + iterations: int = 40, + tolerance: float = 1e-13, + solver: Optional[Callable[..., NDArray[numpy.float64]]] = None, + ) -> Tuple[complex, NDArray[numpy.float64]]: """ Use Rayleigh quotient iteration to refine an eigenvector guess. @@ -69,11 +72,13 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe solver = spalg.spsolve except TypeError: def shift(eigval: float) -> spalg.LinearOperator: - return spalg.LinearOperator(shape=operator.shape, - dtype=operator.dtype, - matvec=lambda v: eigval * v) + return spalg.LinearOperator( + shape=operator.shape, + dtype=operator.dtype, + matvec=lambda v: eigval * v, + ) if solver is None: - def solver(A: spalg.LinearOperator, b: numpy.ndarray) -> numpy.ndarray: + def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.float64]: return spalg.bicgstab(A, b)[0] assert(solver is not None) @@ -90,10 +95,11 @@ def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOpe return eigval, v -def signed_eigensolve(operator: Union[sparse.spmatrix, spalg.LinearOperator], - how_many: int, - negative: bool = False, - ) -> Tuple[numpy.ndarray, numpy.ndarray]: +def signed_eigensolve( + operator: Union[sparse.spmatrix, spalg.LinearOperator], + how_many: int, + negative: bool = False, + ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]: """ Find the largest-magnitude positive-only (or negative-only) eigenvalues and eigenvectors of the provided matrix. diff --git a/meanas/fdfd/bloch.py b/meanas/fdfd/bloch.py index 040c867..872781c 100644 --- a/meanas/fdfd/bloch.py +++ b/meanas/fdfd/bloch.py @@ -80,7 +80,7 @@ This module contains functions for generating and solving the ''' -from typing import Tuple, Callable, Any, List, Optional, cast +from typing import Tuple, Callable, Any, List, Optional, cast, Union import logging import numpy from numpy import pi, real, trace @@ -433,11 +433,10 @@ def find_k( `(k, actual_frequency)` The found k-vector and its frequency. """ - direction = numpy.array(direction) / norm(direction) def get_f(k0_mag: float, band: int = 0) -> float: - k0 = direction * k0_mag + k0 = direction * k0_mag # type: ignore n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu) f = numpy.sqrt(numpy.abs(numpy.real(n[band]))) if solve_callback: @@ -482,6 +481,8 @@ def eigsolve( `(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the vector `eigenvectors[i, :]` """ + k0 = numpy.array(k0, copy=False) + h_size = 2 * epsilon[0].size kmag = norm(G_matrix @ k0) @@ -497,9 +498,9 @@ def eigsolve( y_shape = (h_size, num_modes) - prev_E = 0 - d_scale = 1 - prev_traceGtKG = 0 + prev_E = 0.0 + d_scale = 1.0 + prev_traceGtKG = 0.0 #prev_theta = 0.5 D = numpy.zeros(shape=y_shape, dtype=complex) @@ -545,7 +546,7 @@ def eigsolve( if prev_traceGtKG == 0 or i % reset_iters == 0: logger.info('CG reset') - gamma = 0 + gamma = 0.0 else: gamma = traceGtKG / prev_traceGtKG @@ -695,7 +696,10 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to return x, fx, dfx ''' -def _rtrace_AtB(A: NDArray[numpy.float64], B: NDArray[numpy.float64]) -> NDArray[numpy.float64]: +def _rtrace_AtB( + A: NDArray[numpy.float64], + B: Union[NDArray[numpy.float64], float], + ) -> float: return real(numpy.sum(A.conj() * B)) def _symmetrize(A: NDArray[numpy.float64]) -> NDArray[numpy.float64]: diff --git a/meanas/fdfd/farfield.py b/meanas/fdfd/farfield.py index eb53b24..e1a725d 100644 --- a/meanas/fdfd/farfield.py +++ b/meanas/fdfd/farfield.py @@ -2,19 +2,20 @@ Functions for performing near-to-farfield transformation (and the reverse). """ from typing import Dict, List, Any -import numpy # type: ignore -from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift # type: ignore -from numpy import pi # type: ignore +import numpy +from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift +from numpy import pi from ..fdmath import fdfield_t -def near_to_farfield(E_near: fdfield_t, - H_near: fdfield_t, - dx: float, - dy: float, - padded_size: List[int] = None - ) -> Dict[str, Any]: +def near_to_farfield( + E_near: fdfield_t, + H_near: fdfield_t, + dx: float, + dy: float, + padded_size: List[int] = None + ) -> Dict[str, Any]: """ Compute the farfield, i.e. the distribution of the fields after propagation through several wavelengths of uniform medium. @@ -120,12 +121,13 @@ def near_to_farfield(E_near: fdfield_t, return outputs -def far_to_nearfield(E_far: fdfield_t, - H_far: fdfield_t, - dkx: float, - dky: float, - padded_size: List[int] = None - ) -> Dict[str, Any]: +def far_to_nearfield( + E_far: fdfield_t, + H_far: fdfield_t, + dkx: float, + dky: float, + padded_size: List[int] = None + ) -> Dict[str, Any]: """ Compute the farfield, i.e. the distribution of the fields after propagation through several wavelengths of uniform medium. diff --git a/meanas/fdfd/functional.py b/meanas/fdfd/functional.py index f16151f..9a83910 100644 --- a/meanas/fdfd/functional.py +++ b/meanas/fdfd/functional.py @@ -5,8 +5,8 @@ Functional versions of many FDFD operators. These can be useful for performing The functions generated here expect `fdfield_t` inputs with shape (3, X, Y, Z), e.g. E = [E_x, E_y, E_z] where each component has shape (X, Y, Z) """ -from typing import Callable, Tuple -import numpy # type: ignore +from typing import Callable, Tuple, Optional +import numpy from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t from ..fdmath.functional import curl_forward, curl_back @@ -15,11 +15,12 @@ from ..fdmath.functional import curl_forward, curl_back __author__ = 'Jan Petykiewicz' -def e_full(omega: complex, - dxes: dx_lists_t, - epsilon: fdfield_t, - mu: fdfield_t = None - ) -> fdfield_updater_t: +def e_full( + omega: complex, + dxes: dx_lists_t, + epsilon: fdfield_t, + mu: fdfield_t = None + ) -> fdfield_updater_t: """ Wave operator for use with E-field. See `operators.e_full` for details. @@ -50,11 +51,12 @@ def e_full(omega: complex, return op_mu -def eh_full(omega: complex, - dxes: dx_lists_t, - epsilon: fdfield_t, - mu: fdfield_t = None - ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]: +def eh_full( + omega: complex, + dxes: dx_lists_t, + epsilon: fdfield_t, + mu: fdfield_t = None + ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]: """ Wave operator for full (both E and H) field representation. See `operators.eh_full`. @@ -86,9 +88,10 @@ def eh_full(omega: complex, return op_mu -def e2h(omega: complex, +def e2h( + omega: complex, dxes: dx_lists_t, - mu: fdfield_t = None, + mu: Optional[fdfield_t] = None, ) -> fdfield_updater_t: """ Utility operator for converting the `E` field into the `H` field. @@ -117,9 +120,10 @@ def e2h(omega: complex, return e2h_mu -def m2j(omega: complex, +def m2j( + omega: complex, dxes: dx_lists_t, - mu: fdfield_t = None, + mu: Optional[fdfield_t] = None, ) -> fdfield_updater_t: """ Utility operator for converting magnetic current `M` distribution @@ -151,12 +155,13 @@ def m2j(omega: complex, return m2j_mu -def e_tfsf_source(TF_region: fdfield_t, - omega: complex, - dxes: dx_lists_t, - epsilon: fdfield_t, - mu: fdfield_t = None, - ) -> fdfield_updater_t: +def e_tfsf_source( + TF_region: fdfield_t, + omega: complex, + dxes: dx_lists_t, + epsilon: fdfield_t, + mu: Optional[fdfield_t] = None, + ) -> fdfield_updater_t: """ Operator that turns an E-field distribution into a total-field/scattered-field (TFSF) source. diff --git a/meanas/fdfd/operators.py b/meanas/fdfd/operators.py index 7f0a7ba..f36080b 100644 --- a/meanas/fdfd/operators.py +++ b/meanas/fdfd/operators.py @@ -28,7 +28,7 @@ The following operators are included: """ from typing import Tuple, Optional -import numpy # type: ignore +import numpy import scipy.sparse as sparse # type: ignore from ..fdmath import vec, dx_lists_t, vfdfield_t @@ -38,13 +38,14 @@ from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl __author__ = 'Jan Petykiewicz' -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, - ) -> sparse.spmatrix: +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, + ) -> sparse.spmatrix: """ Wave operator $$ \\nabla \\times (\\frac{1}{\\mu} \\nabla \\times) - \\Omega^2 \\epsilon $$ @@ -96,8 +97,9 @@ def e_full(omega: complex, return op -def e_full_preconditioners(dxes: dx_lists_t - ) -> Tuple[sparse.spmatrix, sparse.spmatrix]: +def e_full_preconditioners( + dxes: dx_lists_t, + ) -> Tuple[sparse.spmatrix, sparse.spmatrix]: """ Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator. @@ -122,13 +124,14 @@ def e_full_preconditioners(dxes: dx_lists_t return P_left, P_right -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, - ) -> sparse.spmatrix: +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, + ) -> sparse.spmatrix: """ Wave operator $$ \\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu $$ @@ -178,13 +181,14 @@ def h_full(omega: complex, return A -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 - ) -> sparse.spmatrix: +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, + ) -> sparse.spmatrix: """ Wave operator for `[E, H]` field representation. This operator implements Maxwell's equations without cancelling out either E or H. The operator is @@ -247,7 +251,8 @@ def eh_full(omega: complex, return A -def e2h(omega: complex, +def e2h( + omega: complex, dxes: dx_lists_t, mu: Optional[vfdfield_t] = None, pmc: Optional[vfdfield_t] = None, @@ -278,9 +283,10 @@ def e2h(omega: complex, return op -def m2j(omega: complex, +def m2j( + omega: complex, dxes: dx_lists_t, - mu: Optional[vfdfield_t] = None + mu: Optional[vfdfield_t] = None, ) -> sparse.spmatrix: """ Operator for converting a magnetic current M into an electric current J. @@ -357,12 +363,13 @@ def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: return P -def e_tfsf_source(TF_region: vfdfield_t, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - ) -> sparse.spmatrix: +def e_tfsf_source( + TF_region: vfdfield_t, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + ) -> sparse.spmatrix: """ Operator that turns a desired E-field distribution into a total-field/scattered-field (TFSF) source. @@ -387,13 +394,14 @@ def e_tfsf_source(TF_region: vfdfield_t, return (A @ Q - Q @ A) / (-1j * omega) -def e_boundary_source(mask: vfdfield_t, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - periodic_mask_edges: bool = False, - ) -> sparse.spmatrix: +def e_boundary_source( + mask: vfdfield_t, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + periodic_mask_edges: bool = False, + ) -> sparse.spmatrix: """ Operator that turns an E-field distrubtion into a current (J) distribution along the edges (external and internal) of the provided mask. This is just an diff --git a/meanas/fdfd/scpml.py b/meanas/fdfd/scpml.py index 67d58ca..0f9c92c 100644 --- a/meanas/fdfd/scpml.py +++ b/meanas/fdfd/scpml.py @@ -3,7 +3,9 @@ Functions for creating stretched coordinate perfectly matched layer (PML) absorb """ from typing import Sequence, Union, Callable, Optional, List -import numpy # type: ignore + +import numpy +from numpy.typing import ArrayLike, NDArray __author__ = 'Jan Petykiewicz' @@ -13,9 +15,10 @@ s_function_t = Callable[[float], float] """Typedef for s-functions, see `prepare_s_function()`""" -def prepare_s_function(ln_R: float = -16, - m: float = 4 - ) -> s_function_t: +def prepare_s_function( + ln_R: float = -16, + m: float = 4 + ) -> s_function_t: """ Create an s_function to pass to the SCPML functions. This is used when you would like to customize the PML parameters. @@ -29,18 +32,19 @@ def prepare_s_function(ln_R: float = -16, of the cell width; needs to be divided by `sqrt(epilon_effective) * real(omega))` before use. """ - def s_factor(distance: numpy.ndarray) -> numpy.ndarray: + def s_factor(distance: NDArray[numpy.float64]) -> NDArray[numpy.float64]: s_max = (m + 1) * ln_R / 2 # / 2 because we assume periodic boundaries return s_max * (distance ** m) return s_factor -def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]], - thicknesses: Union[numpy.ndarray, Sequence[int]], - omega: float, - epsilon_effective: float = 1.0, - s_function: Optional[s_function_t] = None, - ) -> List[List[numpy.ndarray]]: +def uniform_grid_scpml( + shape: ArrayLike, # ints + thicknesses: ArrayLike, # ints + omega: float, + epsilon_effective: float = 1.0, + s_function: Optional[s_function_t] = None, + ) -> List[List[NDArray[numpy.float64]]]: """ Create dx arrays for a uniform grid with a cell width of 1 and a pml. @@ -67,7 +71,11 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]], s_function = prepare_s_function() # Normalized distance to nearest boundary - def ll(u: numpy.ndarray, n: numpy.ndarray, t: numpy.ndarray) -> numpy.ndarray: + def ll( + u: NDArray[numpy.float64], + n: NDArray[numpy.float64], + t: NDArray[numpy.float64], + ) -> NDArray[numpy.float64]: return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t dx_a = [numpy.array(numpy.inf)] * 3 @@ -88,14 +96,15 @@ def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]], return [dx_a, dx_b] -def stretch_with_scpml(dxes: List[List[numpy.ndarray]], - axis: int, - polarity: int, - omega: float, - epsilon_effective: float = 1.0, - thickness: int = 10, - s_function: Optional[s_function_t] = None, - ) -> List[List[numpy.ndarray]]: +def stretch_with_scpml( + dxes: List[List[NDArray[numpy.float64]]], + axis: int, + polarity: int, + omega: float, + epsilon_effective: float = 1.0, + thickness: int = 10, + s_function: Optional[s_function_t] = None, + ) -> List[List[NDArray[numpy.float64]]]: """ Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis. @@ -132,7 +141,7 @@ def stretch_with_scpml(dxes: List[List[numpy.ndarray]], bound = pos[thickness] d = bound - pos[0] - def l_d(x: numpy.ndarray) -> numpy.ndarray: + def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]: return (bound - x) / (bound - pos[0]) slc = slice(thickness) @@ -142,7 +151,7 @@ def stretch_with_scpml(dxes: List[List[numpy.ndarray]], bound = pos[-thickness - 1] d = pos[-1] - bound - def l_d(x: numpy.ndarray) -> numpy.ndarray: + def l_d(x: NDArray[numpy.float64]) -> NDArray[numpy.float64]: return (x - bound) / (pos[-1] - bound) if thickness == 0: diff --git a/meanas/fdfd/solvers.py b/meanas/fdfd/solvers.py index c9f1ac5..0688966 100644 --- a/meanas/fdfd/solvers.py +++ b/meanas/fdfd/solvers.py @@ -2,11 +2,12 @@ Solvers and solver interface for FDFD problems. """ -from typing import Callable, Dict, Any +from typing import Callable, Dict, Any, Optional import logging -import numpy # type: ignore -from numpy.linalg import norm # type: ignore +import numpy +from numpy.typing import ArrayLike, NDArray +from numpy.linalg import norm import scipy.sparse.linalg # type: ignore from ..fdmath import dx_lists_t, vfdfield_t @@ -16,10 +17,11 @@ from . import operators logger = logging.getLogger(__name__) -def _scipy_qmr(A: scipy.sparse.csr_matrix, - b: numpy.ndarray, - **kwargs: Any, - ) -> numpy.ndarray: +def _scipy_qmr( + A: scipy.sparse.csr_matrix, + b: ArrayLike, + **kwargs: Any, + ) -> NDArray[numpy.float64]: """ Wrapper for scipy.sparse.linalg.qmr @@ -37,14 +39,14 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix, ''' ii = 0 - def log_residual(xk: numpy.ndarray) -> None: + def log_residual(xk: ArrayLike) -> None: nonlocal ii ii += 1 if ii % 100 == 0: logger.info('Solver residual at iteration {} : {}'.format(ii, norm(A @ xk - b))) if 'callback' in kwargs: - def augmented_callback(xk: numpy.ndarray) -> None: + def augmented_callback(xk: ArrayLike) -> None: log_residual(xk) kwargs['callback'](xk) @@ -60,17 +62,18 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix, return x -def generic(omega: complex, - dxes: dx_lists_t, - J: vfdfield_t, - epsilon: vfdfield_t, - mu: vfdfield_t = None, - pec: vfdfield_t = None, - pmc: vfdfield_t = None, - adjoint: bool = False, - matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr, - matrix_solver_opts: Dict[str, Any] = None, - ) -> vfdfield_t: +def generic( + omega: complex, + dxes: dx_lists_t, + J: vfdfield_t, + epsilon: vfdfield_t, + mu: vfdfield_t = None, + pec: vfdfield_t = None, + pmc: vfdfield_t = None, + adjoint: bool = False, + matrix_solver: Callable[..., ArrayLike] = _scipy_qmr, + matrix_solver_opts: Optional[Dict[str, Any]] = None, + ) -> vfdfield_t: """ Conjugate gradient FDFD solver using CSR sparse matrices. @@ -90,8 +93,8 @@ def generic(omega: complex, adjoint: If true, solves the adjoint problem. matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`, where `A`: `scipy.sparse.csr_matrix`; - `b`: `numpy.ndarray`; - `x`: `numpy.ndarray`; + `b`: `ArrayLike`; + `x`: `ArrayLike`; Default is a wrapped version of `scipy.sparse.linalg.qmr()` which doesn't return convergence info and logs the residual every 100 iterations. diff --git a/meanas/fdfd/waveguide_2d.py b/meanas/fdfd/waveguide_2d.py index 39463bf..177d8d3 100644 --- a/meanas/fdfd/waveguide_2d.py +++ b/meanas/fdfd/waveguide_2d.py @@ -179,8 +179,9 @@ to account for numerical dispersion if the result is introduced into a space wit # TODO update module docs from typing import List, Tuple, Optional, Any -import numpy # type: ignore -from numpy.linalg import norm # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike +from numpy.linalg import norm import scipy.sparse as sparse # type: ignore from ..fdmath.operators import deriv_forward, deriv_back, cross @@ -191,11 +192,12 @@ from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration __author__ = 'Jan Petykiewicz' -def operator_e(omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - ) -> sparse.spmatrix: +def operator_e( + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + ) -> sparse.spmatrix: """ Waveguide operator of the form @@ -257,11 +259,12 @@ def operator_e(omega: complex, return op -def operator_h(omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - ) -> sparse.spmatrix: +def operator_h( + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + ) -> sparse.spmatrix: """ Waveguide operator of the form @@ -324,14 +327,15 @@ def operator_h(omega: complex, return op -def normalized_fields_e(e_xy: numpy.ndarray, - wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - prop_phase: float = 0, - ) -> Tuple[vfdfield_t, vfdfield_t]: +def normalized_fields_e( + e_xy: ArrayLike, + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + prop_phase: float = 0, + ) -> Tuple[vfdfield_t, vfdfield_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. @@ -358,14 +362,15 @@ def normalized_fields_e(e_xy: numpy.ndarray, return e_norm, h_norm -def normalized_fields_h(h_xy: numpy.ndarray, - wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - prop_phase: float = 0, - ) -> Tuple[vfdfield_t, vfdfield_t]: +def normalized_fields_h( + h_xy: ArrayLike, + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + prop_phase: float = 0, + ) -> Tuple[vfdfield_t, vfdfield_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. @@ -392,14 +397,15 @@ def normalized_fields_h(h_xy: numpy.ndarray, return e_norm, h_norm -def _normalized_fields(e: numpy.ndarray, - h: numpy.ndarray, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None, - prop_phase: float = 0, - ) -> Tuple[vfdfield_t, vfdfield_t]: +def _normalized_fields( + e: ArrayLike, + h: ArrayLike, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None, + prop_phase: float = 0, + ) -> Tuple[vfdfield_t, vfdfield_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)] @@ -434,12 +440,13 @@ def _normalized_fields(e: numpy.ndarray, return e, h -def exy2h(wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None - ) -> sparse.spmatrix: +def exy2h( + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, into a vectorized H containing all three H components @@ -459,12 +466,13 @@ def exy2h(wavenumber: complex, return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon) -def hxy2e(wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: Optional[vfdfield_t] = None - ) -> sparse.spmatrix: +def hxy2e( + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, into a vectorized E containing all three E components @@ -484,10 +492,11 @@ def hxy2e(wavenumber: complex, return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu) -def hxy2h(wavenumber: complex, - dxes: dx_lists_t, - mu: Optional[vfdfield_t] = None - ) -> sparse.spmatrix: +def hxy2h( + wavenumber: complex, + dxes: dx_lists_t, + mu: Optional[vfdfield_t] = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, into a vectorized H containing all three H components @@ -517,10 +526,11 @@ def hxy2h(wavenumber: complex, return op -def exy2e(wavenumber: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - ) -> sparse.spmatrix: +def exy2e( + wavenumber: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + ) -> sparse.spmatrix: """ Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, into a vectorized E containing all three E components @@ -550,7 +560,8 @@ def exy2e(wavenumber: complex, return op -def e2h(wavenumber: complex, +def e2h( + wavenumber: complex, omega: complex, dxes: dx_lists_t, mu: Optional[vfdfield_t] = None @@ -574,7 +585,8 @@ def e2h(wavenumber: complex, return op -def h2e(wavenumber: complex, +def h2e( + wavenumber: complex, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t @@ -636,13 +648,14 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: return cross([Dbx, Dby, Bz]) -def h_err(h: vfdfield_t, - wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t = None - ) -> float: +def h_err( + h: vfdfield_t, + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: Optional[vfdfield_t] = None + ) -> float: """ Calculates the relative error in the H field @@ -670,13 +683,14 @@ def h_err(h: vfdfield_t, return norm(op) / norm(h) -def e_err(e: vfdfield_t, - wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t = None - ) -> float: +def e_err( + e: vfdfield_t, + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t = Optional[None] + ) -> float: """ Calculates the relative error in the E field @@ -703,13 +717,14 @@ def e_err(e: vfdfield_t, return norm(op) / norm(e) -def solve_modes(mode_numbers: List[int], - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t = None, - mode_margin: int = 2, - ) -> Tuple[numpy.ndarray, List[complex]]: +def solve_modes( + mode_numbers: List[int], + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t = None, + mode_margin: int = 2, + ) -> Tuple[NDArray[numpy.float64], List[complex]]: """ Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers. @@ -752,10 +767,11 @@ def solve_modes(mode_numbers: List[int], return e_xys, wavenumbers -def solve_mode(mode_number: int, - *args: Any, - **kwargs: Any, - ) -> Tuple[vfdfield_t, complex]: +def solve_mode( + mode_number: int, + *args: Any, + **kwargs: Any, + ) -> Tuple[vfdfield_t, complex]: """ Wrapper around `solve_modes()` that solves for a single mode. diff --git a/meanas/fdfd/waveguide_3d.py b/meanas/fdfd/waveguide_3d.py index 4a65453..a6a2cba 100644 --- a/meanas/fdfd/waveguide_3d.py +++ b/meanas/fdfd/waveguide_3d.py @@ -5,21 +5,23 @@ This module relies heavily on `waveguide_2d` and mostly just transforms its parameters into 2D equivalents and expands the results back into 3D. """ from typing import Dict, Optional, Sequence, Union, Any -import numpy # type: ignore +import numpy +from numpy.typing import NDArray from ..fdmath import vec, unvec, dx_lists_t, fdfield_t from . import operators, waveguide_2d -def solve_mode(mode_number: int, - omega: complex, - dxes: dx_lists_t, - axis: int, - polarity: int, - slices: Sequence[slice], - epsilon: fdfield_t, - mu: Optional[fdfield_t] = None, - ) -> Dict[str, Union[complex, numpy.ndarray]]: +def solve_mode( + mode_number: int, + omega: complex, + dxes: dx_lists_t, + axis: int, + polarity: int, + slices: Sequence[slice], + epsilon: fdfield_t, + mu: Optional[fdfield_t] = None, + ) -> Dict[str, Union[complex, NDArray[numpy.float_]]]: """ Given a 3D grid, selects a slice from the grid and attempts to solve for an eigenmode propagating through that slice. @@ -36,7 +38,13 @@ def solve_mode(mode_number: int, mu: Magnetic permeability (default 1 everywhere) Returns: - `{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}` + ``` + { + 'E': List[NDArray[numpy.float_]], + 'H': List[NDArray[numpy.float_]], + 'wavenumber': complex, + } + ``` """ if mu is None: mu = numpy.ones_like(epsilon) @@ -97,16 +105,17 @@ def solve_mode(mode_number: int, return results -def compute_source(E: fdfield_t, - wavenumber: complex, - omega: complex, - dxes: dx_lists_t, - axis: int, - polarity: int, - slices: Sequence[slice], - epsilon: fdfield_t, - mu: Optional[fdfield_t] = None, - ) -> fdfield_t: +def compute_source( + E: fdfield_t, + wavenumber: complex, + omega: complex, + dxes: dx_lists_t, + axis: int, + polarity: int, + slices: Sequence[slice], + epsilon: fdfield_t, + mu: Optional[fdfield_t] = None, + ) -> fdfield_t: """ Given an eigenmode obtained by `solve_mode`, returns the current source distribution necessary to position a unidirectional source at the slice location. @@ -142,18 +151,21 @@ def compute_source(E: fdfield_t, return J -def compute_overlap_e(E: fdfield_t, - wavenumber: complex, - dxes: dx_lists_t, - axis: int, - polarity: int, - slices: Sequence[slice], - ) -> fdfield_t: # TODO DOCS +def compute_overlap_e( + E: fdfield_t, + wavenumber: complex, + dxes: dx_lists_t, + axis: int, + polarity: int, + slices: Sequence[slice], + ) -> fdfield_t: # TODO DOCS """ Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn) [assumes reflection symmetry]. + TODO: add reference + Args: E: E-field of the mode H: H-field of the mode (advanced by half of a Yee cell from E) @@ -187,13 +199,14 @@ def compute_overlap_e(E: fdfield_t, return Etgt -def expand_e(E: fdfield_t, - wavenumber: complex, - dxes: dx_lists_t, - axis: int, - polarity: int, - slices: Sequence[slice], - ) -> fdfield_t: +def expand_e( + E: fdfield_t, + wavenumber: complex, + dxes: dx_lists_t, + axis: int, + polarity: int, + slices: Sequence[slice], + ) -> fdfield_t: """ Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D slice where the mode was calculated to the entire domain (along the propagation diff --git a/meanas/fdfd/waveguide_cyl.py b/meanas/fdfd/waveguide_cyl.py index da75f0c..a1a0633 100644 --- a/meanas/fdfd/waveguide_cyl.py +++ b/meanas/fdfd/waveguide_cyl.py @@ -9,7 +9,7 @@ As the z-dependence is known, all the functions in this file assume a 2D grid # TODO update module docs from typing import Dict, Union -import numpy # type: ignore +import numpy import scipy.sparse as sparse # type: ignore from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t @@ -17,11 +17,12 @@ from ..fdmath.operators import deriv_forward, deriv_back from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration -def cylindrical_operator(omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - r0: float, - ) -> sparse.spmatrix: +def cylindrical_operator( + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + r0: float, + ) -> sparse.spmatrix: """ Cylindrical coordinate waveguide operator of the form @@ -78,12 +79,13 @@ def cylindrical_operator(omega: complex, return op -def solve_mode(mode_number: int, - omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - r0: float, - ) -> Dict[str, Union[complex, fdfield_t]]: +def solve_mode( + mode_number: int, + omega: complex, + dxes: dx_lists_t, + epsilon: vfdfield_t, + r0: float, + ) -> Dict[str, Union[complex, fdfield_t]]: """ TODO: fixup Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode @@ -99,7 +101,13 @@ def solve_mode(mode_number: int, r within the simulation domain. Returns: - `{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}` + ``` + { + 'E': List[NDArray[numpy.float_]], + 'H': List[NDArray[numpy.float_]], + 'wavenumber': complex, + } + ``` """ ''' diff --git a/meanas/fdmath/functional.py b/meanas/fdmath/functional.py index a62655d..aad8a33 100644 --- a/meanas/fdmath/functional.py +++ b/meanas/fdmath/functional.py @@ -5,13 +5,15 @@ Basic discrete calculus etc. """ from typing import Sequence, Tuple, Optional, Callable -import numpy # type: ignore +import numpy +from numpy.typing import NDArray from .types import fdfield_t, fdfield_updater_t -def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None - ) -> Tuple[fdfield_updater_t, fdfield_updater_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]: """ Utility operators for taking discretized derivatives (backward variant). @@ -33,8 +35,9 @@ def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None return derivs -def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None - ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: +def deriv_back( + dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None, + ) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]: """ Utility operators for taking discretized derivatives (forward variant). @@ -56,7 +59,9 @@ def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None return derivs -def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t: +def curl_forward( + dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None, + ) -> fdfield_updater_t: """ Curl operator for use with the E field. @@ -83,7 +88,9 @@ def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_upda return ce_fun -def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t: +def curl_back( + dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None, + ) -> fdfield_updater_t: """ Create a function which takes the backward curl of a field. @@ -110,7 +117,9 @@ def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater return ch_fun -def curl_forward_parts(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> Callable: +def curl_forward_parts( + dx_e: Optional[Sequence[NDArray[numpy.float_]]] = None, + ) -> Callable: Dx, Dy, Dz = deriv_forward(dx_e) def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]: @@ -121,7 +130,9 @@ def curl_forward_parts(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> Callab return mkparts_fwd -def curl_back_parts(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> Callable: +def curl_back_parts( + dx_h: Optional[Sequence[NDArray[numpy.float_]]] = None, + ) -> Callable: Dx, Dy, Dz = deriv_back(dx_e) def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]: diff --git a/meanas/fdmath/operators.py b/meanas/fdmath/operators.py index 45e0cea..d90261a 100644 --- a/meanas/fdmath/operators.py +++ b/meanas/fdmath/operators.py @@ -4,13 +4,18 @@ Matrix operators for finite difference simulations Basic discrete calculus etc. """ from typing import Sequence, List -import numpy # type: ignore +import numpy +from numpy.typing import NDArray import scipy.sparse as sparse # type: ignore from .types import vfdfield_t -def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> sparse.spmatrix: +def shift_circ( + axis: int, + shape: Sequence[int], + shift_distance: int = 1, + ) -> sparse.spmatrix: """ Utility operator for performing a circular shift along a specified axis by a specified number of elements. @@ -46,7 +51,11 @@ def shift_circ(axis: int, shape: Sequence[int], shift_distance: int = 1) -> spar return d -def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1) -> sparse.spmatrix: +def shift_with_mirror( + axis: int, + shape: Sequence[int], + shift_distance: int = 1, + ) -> sparse.spmatrix: """ Utility operator for performing an n-element shift along a specified axis, with mirror boundary conditions applied to the cells beyond the receding edge. @@ -67,7 +76,7 @@ def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1) raise Exception('Shift ({}) is too large for axis {} of size {}'.format( shift_distance, axis, shape[axis])) - def mirrored_range(n: int, s: int) -> numpy.ndarray: + def mirrored_range(n: int, s: int) -> NDArray[numpy.int_]: v = numpy.arange(n) + s v = numpy.where(v >= n, 2 * n - v - 1, v) v = numpy.where(v < 0, - 1 - v, v) @@ -87,7 +96,9 @@ def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int = 1) return d -def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]: +def deriv_forward( + dx_e: Sequence[NDArray[numpy.float_]], + ) -> List[sparse.spmatrix]: """ Utility operators for taking discretized derivatives (forward variant). @@ -112,7 +123,9 @@ def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]: return Ds -def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]: +def deriv_back( + dx_h: Sequence[NDArray[numpy.float_]], + ) -> List[sparse.spmatrix]: """ Utility operators for taking discretized derivatives (backward variant). @@ -137,7 +150,9 @@ def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]: return Ds -def cross(B: Sequence[sparse.spmatrix]) -> sparse.spmatrix: +def cross( + B: Sequence[sparse.spmatrix], + ) -> sparse.spmatrix: """ Cross product operator @@ -203,7 +218,9 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix: return avg_forward(axis, shape).T -def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix: +def curl_forward( + dx_e: Sequence[NDArray[numpy.float_]], + ) -> sparse.spmatrix: """ Curl operator for use with the E field. @@ -217,7 +234,9 @@ def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix: return cross(deriv_forward(dx_e)) -def curl_back(dx_h: Sequence[numpy.ndarray]) -> sparse.spmatrix: +def curl_back( + dx_h: Sequence[NDArray[numpy.float_]], + ) -> sparse.spmatrix: """ Curl operator for use with the H field. diff --git a/meanas/fdmath/types.py b/meanas/fdmath/types.py index 2676bc5..2dd0040 100644 --- a/meanas/fdmath/types.py +++ b/meanas/fdmath/types.py @@ -2,31 +2,20 @@ Types shared across multiple submodules """ from typing import Sequence, Callable, MutableSequence -import numpy # type: ignore +import numpy +from numpy.typing import NDArray # Field types -# TODO: figure out a better way to set the docstrings without creating actual subclasses? -# Probably not a big issue since they're only used for type hinting -class fdfield_t(numpy.ndarray): - """ - Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`) +fdfield_t = NDArray[numpy.float_] +"""Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)""" - This is actually is just an unaltered `numpy.ndarray` - """ - pass - -class vfdfield_t(numpy.ndarray): - """ - Linearized vector field (single vector of length 3*X*Y*Z) - - This is actually just an unaltered `numpy.ndarray` - """ - pass +vfdfield_t = NDArray[numpy.float_] +"""Linearized vector field (single vector of length 3*X*Y*Z)""" -dx_lists_t = Sequence[Sequence[numpy.ndarray]] -''' +dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]] +""" 'dxes' datastructure which contains grid cell width information in the following format: [[[dx_e[0], dx_e[1], ...], [dy_e[0], ...], [dz_e[0], ...]], @@ -34,15 +23,11 @@ dx_lists_t = Sequence[Sequence[numpy.ndarray]] where `dx_e[0]` is the x-width of the `x=0` cells, as used when calculating dE/dx, and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc. -''' +""" -dx_lists_mut = MutableSequence[MutableSequence[numpy.ndarray]] -''' - Mutable version of `dx_lists_t` -''' +dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]] +"""Mutable version of `dx_lists_t`""" fdfield_updater_t = Callable[..., fdfield_t] -''' - Convenience type for functions which take and return an fdfield_t -''' +"""Convenience type for functions which take and return an fdfield_t""" diff --git a/meanas/fdmath/vectorization.py b/meanas/fdmath/vectorization.py index 23e3d9c..6b6c49e 100644 --- a/meanas/fdmath/vectorization.py +++ b/meanas/fdmath/vectorization.py @@ -5,7 +5,8 @@ Vectorized versions of the field use row-major (ie., C-style) ordering. """ from typing import Optional, overload, Union, List -import numpy # type: ignore +import numpy +from numpy.typing import ArrayLike from .types import fdfield_t, vfdfield_t @@ -15,10 +16,10 @@ def vec(f: None) -> None: pass @overload -def vec(f: Union[fdfield_t, List[numpy.ndarray]]) -> vfdfield_t: +def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t: pass -def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield_t]: +def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]: """ Create a 1D ndarray from a 3D vector field which spans a 1-3D region. @@ -37,14 +38,14 @@ def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield @overload -def unvec(v: None, shape: numpy.ndarray) -> None: +def unvec(v: None, shape: ArrayLike) -> None: pass @overload -def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t: +def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t: pass -def unvec(v: Optional[vfdfield_t], shape: numpy.ndarray) -> Optional[fdfield_t]: +def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> Optional[fdfield_t]: """ Perform the inverse of vec(): take a 1D ndarray and output a 3D field of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional diff --git a/meanas/fdtd/boundaries.py b/meanas/fdtd/boundaries.py index d03a976..4171936 100644 --- a/meanas/fdtd/boundaries.py +++ b/meanas/fdtd/boundaries.py @@ -9,9 +9,10 @@ from typing import Tuple, Any, List from ..fdmath import fdfield_t, fdfield_updater_t -def conducting_boundary(direction: int, - polarity: int - ) -> Tuple[fdfield_updater_t, fdfield_updater_t]: +def conducting_boundary( + direction: int, + polarity: int + ) -> Tuple[fdfield_updater_t, fdfield_updater_t]: dirs = [0, 1, 2] if direction not in dirs: raise Exception('Invalid direction: {}'.format(direction)) diff --git a/meanas/fdtd/energy.py b/meanas/fdtd/energy.py index 93eedf0..ca7d308 100644 --- a/meanas/fdtd/energy.py +++ b/meanas/fdtd/energy.py @@ -1,5 +1,5 @@ from typing import Optional, Union -import numpy # type: ignore +import numpy from ..fdmath import dx_lists_t, fdfield_t from ..fdmath.functional import deriv_back @@ -8,10 +8,11 @@ from ..fdmath.functional import deriv_back # TODO documentation -def poynting(e: fdfield_t, - h: fdfield_t, - dxes: Optional[dx_lists_t] = None, - ) -> fdfield_t: +def poynting( + e: fdfield_t, + h: fdfield_t, + dxes: Optional[dx_lists_t] = None, + ) -> fdfield_t: """ Calculate the poynting vector `S` ($S$). @@ -87,12 +88,13 @@ def poynting(e: fdfield_t, return s -def poynting_divergence(s: Optional[fdfield_t] = None, - *, - e: Optional[fdfield_t] = None, - h: Optional[fdfield_t] = None, - dxes: Optional[dx_lists_t] = None, - ) -> fdfield_t: +def poynting_divergence( + s: Optional[fdfield_t] = None, + *, + e: Optional[fdfield_t] = None, + h: Optional[fdfield_t] = None, + dxes: Optional[dx_lists_t] = None, + ) -> fdfield_t: """ Calculate the divergence of the poynting vector. @@ -124,13 +126,14 @@ def poynting_divergence(s: Optional[fdfield_t] = None, return ds -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, - ) -> fdfield_t: +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, + ) -> fdfield_t: """ Calculate energy `U` at the time of the provided H-field `h1`. @@ -151,13 +154,14 @@ def energy_hstep(e0: fdfield_t, return u -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, - ) -> fdfield_t: +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, + ) -> fdfield_t: """ Calculate energy `U` at the time of the provided E-field `e1`. @@ -178,15 +182,16 @@ def energy_estep(h0: fdfield_t, return u -def delta_energy_h2e(dt: float, - e0: fdfield_t, - 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, - ) -> fdfield_t: +def delta_energy_h2e( + dt: float, + e0: fdfield_t, + 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, + ) -> fdfield_t: """ Change in energy during the half-step from `h1` to `e2`. @@ -210,15 +215,16 @@ def delta_energy_h2e(dt: float, return du -def delta_energy_e2h(dt: float, - h0: fdfield_t, - 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, - ) -> fdfield_t: +def delta_energy_e2h( + dt: float, + h0: fdfield_t, + 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, + ) -> fdfield_t: """ Change in energy during the half-step from `e1` to `h2`. @@ -242,10 +248,11 @@ def delta_energy_e2h(dt: float, return du -def delta_energy_j(j0: fdfield_t, - e1: fdfield_t, - dxes: Optional[dx_lists_t] = None, - ) -> fdfield_t: +def delta_energy_j( + j0: fdfield_t, + e1: fdfield_t, + dxes: Optional[dx_lists_t] = None, + ) -> fdfield_t: """ Calculate @@ -264,12 +271,13 @@ def delta_energy_j(j0: fdfield_t, return du -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 - ) -> fdfield_t: +def dxmul( + ee: fdfield_t, + hh: fdfield_t, + epsilon: Optional[Union[fdfield_t, float]] = None, + mu: Optional[Union[fdfield_t, float]] = None, + dxes: Optional[dx_lists_t] = None, + ) -> fdfield_t: if epsilon is None: epsilon = 1 if mu is None: diff --git a/meanas/fdtd/pml.py b/meanas/fdtd/pml.py index 066cca8..6f7aff7 100644 --- a/meanas/fdtd/pml.py +++ b/meanas/fdtd/pml.py @@ -8,7 +8,8 @@ PML implementations # TODO retest pmls! from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional -import numpy # type: ignore +import numpy +from typing import NDArray from ..fdmath import fdfield_t, dx_lists_t from ..fdmath.functional import deriv_forward, deriv_back @@ -61,7 +62,7 @@ def cpml_params( expand_slice_l[axis] = slice(None) expand_slice = tuple(expand_slice_l) - def par(x: numpy.ndarray) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: + def par(x: NDArray[numpy.float64]) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]: scaling = (x / thickness) ** m sigma = scaling * sigma_max kappa = 1 + scaling * (kappa_max - 1) diff --git a/meanas/test/conftest.py b/meanas/test/conftest.py index 311aae3..ba6a3a8 100644 --- a/meanas/test/conftest.py +++ b/meanas/test/conftest.py @@ -4,7 +4,8 @@ Test fixtures """ from typing import Tuple, Iterable, List, Any -import numpy # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike import pytest # type: ignore from .utils import PRNG @@ -34,11 +35,12 @@ def epsilon_fg(request: FixtureRequest) -> Iterable[float]: @pytest.fixture(scope='module', params=['center', '000', 'random']) -def epsilon(request: FixtureRequest, - shape: Tuple[int, ...], - epsilon_bg: float, - epsilon_fg: float, - ) -> Iterable[numpy.ndarray]: +def epsilon( + request: FixtureRequest, + shape: Tuple[int, ...], + epsilon_bg: float, + epsilon_fg: float, + ) -> Iterable[NDArray[numpy.float64]]: is3d = (numpy.array(shape) == 1).sum() == 0 if is3d: if request.param == '000': @@ -72,10 +74,11 @@ def dx(request: FixtureRequest) -> Iterable[float]: @pytest.fixture(scope='module', params=['uniform', 'centerbig']) -def dxes(request: FixtureRequest, - shape: Tuple[int, ...], - dx: float, - ) -> Iterable[List[List[numpy.ndarray]]]: +def dxes( + request: FixtureRequest, + shape: Tuple[int, ...], + dx: float, + ) -> 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': diff --git a/meanas/test/test_fdfd.py b/meanas/test/test_fdfd.py index 076cb52..fef80fd 100644 --- a/meanas/test/test_fdfd.py +++ b/meanas/test/test_fdfd.py @@ -1,7 +1,8 @@ from typing import List, Tuple, Iterable, Optional import dataclasses import pytest # type: ignore -import numpy # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike #from numpy.testing import assert_allclose, assert_array_equal from .. import fdfd @@ -59,12 +60,12 @@ def omega(request: FixtureRequest) -> Iterable[float]: @pytest.fixture(params=[None]) -def pec(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]: +def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]: yield request.param @pytest.fixture(params=[None]) -def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]: +def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]: yield request.param @@ -75,10 +76,11 @@ def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]: @pytest.fixture(params=['diag']) # 'center' -def j_distribution(request: FixtureRequest, - shape: Tuple[int, ...], - j_mag: float, - ) -> Iterable[numpy.ndarray]: +def j_distribution( + request: FixtureRequest, + shape: Tuple[int, ...], + j_mag: float, + ) -> Iterable[NDArray[numpy.float64]]: j = numpy.zeros(shape, dtype=complex) center_mask = numpy.zeros(shape, dtype=bool) center_mask[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = True @@ -94,24 +96,25 @@ def j_distribution(request: FixtureRequest, @dataclasses.dataclass() class FDResult: shape: Tuple[int, ...] - dxes: List[List[numpy.ndarray]] - epsilon: numpy.ndarray + dxes: List[List[NDArray[numpy.float64]]] + epsilon: NDArray[numpy.float64] omega: complex - j: numpy.ndarray - e: numpy.ndarray - pmc: numpy.ndarray - pec: numpy.ndarray + j: NDArray[numpy.float64] + e: NDArray[numpy.float64] + pmc: Optional[NDArray[numpy.float64]] + pec: Optional[NDArray[numpy.float64]] @pytest.fixture() -def sim(request: FixtureRequest, +def sim( + request: FixtureRequest, shape: Tuple[int, ...], - epsilon: numpy.ndarray, - dxes: List[List[numpy.ndarray]], - j_distribution: numpy.ndarray, + epsilon: NDArray[numpy.float64], + dxes: List[List[NDArray[numpy.float64]]], + j_distribution: NDArray[numpy.float64], omega: float, - pec: Optional[numpy.ndarray], - pmc: Optional[numpy.ndarray], + pec: Optional[NDArray[numpy.float64]], + pmc: Optional[NDArray[numpy.float64]], ) -> FDResult: """ Build simulation from parts diff --git a/meanas/test/test_fdfd_pml.py b/meanas/test/test_fdfd_pml.py index cf9c05a..30eb32d 100644 --- a/meanas/test/test_fdfd_pml.py +++ b/meanas/test/test_fdfd_pml.py @@ -1,7 +1,8 @@ from typing import Optional, Tuple, Iterable, List import pytest # type: ignore -import numpy # type: ignore -from numpy.testing import assert_allclose # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike +from numpy.testing import assert_allclose from .. import fdfd from ..fdmath import vec, unvec, dx_lists_mut @@ -48,12 +49,12 @@ def omega(request: FixtureRequest) -> Iterable[float]: @pytest.fixture(params=[None]) -def pec(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]: +def pec(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]: yield request.param @pytest.fixture(params=[None]) -def pmc(request: FixtureRequest) -> Iterable[Optional[numpy.ndarray]]: +def pmc(request: FixtureRequest) -> Iterable[Optional[NDArray[numpy.float64]]]: yield request.param @@ -70,13 +71,14 @@ def src_polarity(request: FixtureRequest) -> Iterable[int]: @pytest.fixture() -def j_distribution(request: FixtureRequest, - shape: Tuple[int, ...], - epsilon: numpy.ndarray, - dxes: dx_lists_mut, - omega: float, - src_polarity: int, - ) -> Iterable[numpy.ndarray]: +def j_distribution( + request: FixtureRequest, + shape: Tuple[int, ...], + epsilon: NDArray[numpy.float64], + dxes: dx_lists_mut, + omega: float, + src_polarity: int, + ) -> Iterable[NDArray[numpy.float64]]: j = numpy.zeros(shape, dtype=complex) dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis @@ -108,47 +110,60 @@ def j_distribution(request: FixtureRequest, @pytest.fixture() -def epsilon(request: FixtureRequest, - shape: Tuple[int, ...], - epsilon_bg: float, - epsilon_fg: float, - ) -> Iterable[numpy.ndarray]: +def epsilon( + request: FixtureRequest, + shape: Tuple[int, ...], + epsilon_bg: float, + epsilon_fg: float, + ) -> Iterable[NDArray[numpy.float64]]: epsilon = numpy.full(shape, epsilon_fg, dtype=float) yield epsilon @pytest.fixture(params=['uniform']) -def dxes(request: FixtureRequest, - shape: Tuple[int, ...], - dx: float, - omega: float, - epsilon_fg: float, - ) -> Iterable[List[List[numpy.ndarray]]]: +def dxes( + request: FixtureRequest, + shape: Tuple[int, ...], + dx: float, + omega: float, + epsilon_fg: float, + ) -> Iterable[List[List[NDArray[numpy.float64]]]]: if request.param == 'uniform': dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)] dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis for axis in (dim,): for polarity in (-1, 1): - dxes = fdfd.scpml.stretch_with_scpml(dxes, axis=axis, polarity=polarity, - omega=omega, epsilon_effective=epsilon_fg, - thickness=10) + dxes = fdfd.scpml.stretch_with_scpml( + dxes, + axis=axis, + polarity=polarity, + omega=omega, + epsilon_effective=epsilon_fg, + thickness=10, + ) yield dxes @pytest.fixture() -def sim(request: FixtureRequest, +def sim( + request: FixtureRequest, shape: Tuple[int, ...], - epsilon: numpy.ndarray, + epsilon: NDArray[numpy.float64], dxes: dx_lists_mut, - j_distribution: numpy.ndarray, + j_distribution: NDArray[numpy.float64], omega: float, - pec: Optional[numpy.ndarray], - pmc: Optional[numpy.ndarray], + pec: Optional[NDArray[numpy.float64]], + pmc: Optional[NDArray[numpy.float64]], ) -> FDResult: j_vec = vec(j_distribution) eps_vec = vec(epsilon) - e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec, - matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11}) + e_vec = fdfd.solvers.generic( + J=j_vec, + omega=omega, + dxes=dxes, + epsilon=eps_vec, + matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11}, + ) e = unvec(e_vec, shape[1:]) sim = FDResult( diff --git a/meanas/test/test_fdtd.py b/meanas/test/test_fdtd.py index 8f5e013..b46a3ca 100644 --- a/meanas/test/test_fdtd.py +++ b/meanas/test/test_fdtd.py @@ -1,8 +1,9 @@ -from typing import List, Tuple, Iterable +from typing import List, Tuple, Iterable, Any, Dict import dataclasses import pytest # type: ignore -import numpy # type: ignore -#from numpy.testing import assert_allclose, assert_array_equal # type: ignore +import numpy +from numpy.typing import NDArray, ArrayLike +#from numpy.testing import assert_allclose, assert_array_equal from .. import fdtd from .utils import assert_close, assert_fields_close, PRNG @@ -32,8 +33,10 @@ def test_initial_energy(sim: 'TDResult') -> None: dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0) u0 = (j0 * j0.conj() / sim.epsilon * dV).sum(axis=0) - args = {'dxes': sim.dxes, - 'epsilon': sim.epsilon} + args: Dict[str, Any] = { + 'dxes': sim.dxes, + 'epsilon': sim.epsilon, + } # Make sure initial energy and E dot J are correct energy0 = fdtd.energy_estep(h0=h0, e1=e0, h2=h1, **args) @@ -49,8 +52,10 @@ def test_energy_conservation(sim: 'TDResult') -> None: e0 = sim.es[0] j0 = sim.js[0] u = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes).sum() - args = {'dxes': sim.dxes, - 'epsilon': sim.epsilon} + args: Dict[str, Any] = { + 'dxes': sim.dxes, + 'epsilon': sim.epsilon, + } for ii in range(1, 8): u_hstep = fdtd.energy_hstep(e0=sim.es[ii - 1], h1=sim.hs[ii], e2=sim.es[ii], **args) @@ -65,8 +70,10 @@ def test_energy_conservation(sim: 'TDResult') -> None: def test_poynting_divergence(sim: 'TDResult') -> None: - args = {'dxes': sim.dxes, - 'epsilon': sim.epsilon} + args: Dict[str, Any] = { + 'dxes': sim.dxes, + 'epsilon': sim.epsilon, + } u_eprev = None for ii in range(1, 8): @@ -96,8 +103,10 @@ def test_poynting_planes(sim: 'TDResult') -> None: if mask.sum() > 1: pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum())) - args = {'dxes': sim.dxes, - 'epsilon': sim.epsilon} + args: Dict[str, Any] = { + 'dxes': sim.dxes, + 'epsilon': sim.epsilon, + } mx = numpy.roll(mask, -1, axis=0) my = numpy.roll(mask, -1, axis=1) @@ -149,13 +158,13 @@ def dt(request: FixtureRequest) -> Iterable[float]: class TDResult: shape: Tuple[int, ...] dt: float - dxes: List[List[numpy.ndarray]] - epsilon: numpy.ndarray - j_distribution: numpy.ndarray + dxes: List[List[NDArray[numpy.float64]]] + epsilon: NDArray[numpy.float64] + j_distribution: NDArray[numpy.float64] j_steps: Tuple[int, ...] - es: List[numpy.ndarray] = dataclasses.field(default_factory=list) - hs: List[numpy.ndarray] = dataclasses.field(default_factory=list) - js: List[numpy.ndarray] = dataclasses.field(default_factory=list) + es: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list) + hs: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list) + js: List[NDArray[numpy.float64]] = dataclasses.field(default_factory=list) @pytest.fixture(params=[(0, 4, 8)]) # (0,) @@ -164,10 +173,11 @@ def j_steps(request: FixtureRequest) -> Iterable[Tuple[int, ...]]: @pytest.fixture(params=['center', 'random']) -def j_distribution(request: FixtureRequest, - shape: Tuple[int, ...], - j_mag: float, - ) -> Iterable[numpy.ndarray]: +def j_distribution( + request: FixtureRequest, + shape: Tuple[int, ...], + j_mag: float, + ) -> Iterable[NDArray[numpy.float64]]: j = numpy.zeros(shape) if request.param == 'center': j[:, shape[1] // 2, shape[2] // 2, shape[3] // 2] = j_mag @@ -179,12 +189,13 @@ def j_distribution(request: FixtureRequest, @pytest.fixture() -def sim(request: FixtureRequest, +def sim( + request: FixtureRequest, shape: Tuple[int, ...], - epsilon: numpy.ndarray, - dxes: List[List[numpy.ndarray]], + epsilon: NDArray[numpy.float64], + dxes: List[List[NDArray[numpy.float64]]], dt: float, - j_distribution: numpy.ndarray, + j_distribution: NDArray[numpy.float64], j_steps: Tuple[int, ...], ) -> TDResult: is3d = (numpy.array(shape) == 1).sum() == 0 diff --git a/meanas/test/utils.py b/meanas/test/utils.py index f76b910..b4cb3ab 100644 --- a/meanas/test/utils.py +++ b/meanas/test/utils.py @@ -1,24 +1,27 @@ from typing import Any -import numpy # type: ignore +import numpy +from typing import ArrayLike PRNG = numpy.random.RandomState(12345) -def assert_fields_close(x: numpy.ndarray, - y: numpy.ndarray, - *args: Any, - **kwargs: Any, - ) -> None: +def assert_fields_close( + x: ArrayLike, + y: ArrayLike, + *args: Any, + **kwargs: Any, + ) -> None: numpy.testing.assert_allclose( x, y, verbose=False, err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1), numpy.rollaxis(y, -1)), *args, **kwargs) -def assert_close(x: numpy.ndarray, - y: numpy.ndarray, - *args: Any, - **kwargs: Any, - ) -> None: +def assert_close( + x: ArrayLike, + y: ArrayLike, + *args: Any, + **kwargs: Any, + ) -> None: numpy.testing.assert_allclose(x, y, *args, **kwargs)