diff --git a/meanas/eigensolvers.py b/meanas/eigensolvers.py index aa8b9ba..4b96f60 100644 --- a/meanas/eigensolvers.py +++ b/meanas/eigensolvers.py @@ -11,9 +11,9 @@ import scipy.sparse.linalg as spalg # type: ignore def power_iteration( operator: sparse.spmatrix, - guess_vector: Optional[NDArray[numpy.float64]] = None, + guess_vector: Optional[NDArray[numpy.complex128]] = None, iterations: int = 20, - ) -> Tuple[complex, NDArray[numpy.float64]]: + ) -> Tuple[complex, NDArray[numpy.complex128]]: """ Use power iteration to estimate the dominant eigenvector of a matrix. @@ -26,7 +26,7 @@ def power_iteration( (Largest-magnitude eigenvalue, Corresponding eigenvector estimate) """ if numpy.any(numpy.equal(guess_vector, None)): - v = numpy.random.rand(operator.shape[0]) + v = numpy.random.rand(operator.shape[0]) + 1j * numpy.random.rand(operator.shape[0]) else: v = guess_vector @@ -41,11 +41,11 @@ def power_iteration( def rayleigh_quotient_iteration( operator: Union[sparse.spmatrix, spalg.LinearOperator], - guess_vector: NDArray[numpy.float64], + guess_vector: NDArray[numpy.complex128], iterations: int = 40, tolerance: float = 1e-13, - solver: Optional[Callable[..., NDArray[numpy.float64]]] = None, - ) -> Tuple[complex, NDArray[numpy.float64]]: + solver: Optional[Callable[..., NDArray[numpy.complex128]]] = None, + ) -> Tuple[complex, NDArray[numpy.complex128]]: """ Use Rayleigh quotient iteration to refine an eigenvector guess. @@ -78,7 +78,7 @@ def rayleigh_quotient_iteration( matvec=lambda v: eigval * v, ) if solver is None: - def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.float64]: + def solver(A: spalg.LinearOperator, b: ArrayLike) -> NDArray[numpy.complex128]: return spalg.bicgstab(A, b)[0] assert(solver is not None) @@ -99,7 +99,7 @@ def signed_eigensolve( operator: Union[sparse.spmatrix, spalg.LinearOperator], how_many: int, negative: bool = False, - ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]: + ) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: """ 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 a300ab7..0e3ca51 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, Union +from typing import Tuple, Callable, Any, List, Optional, cast, Union, Sequence import logging import numpy from numpy import pi, real, trace @@ -91,7 +91,8 @@ import scipy.optimize # type: ignore from scipy.linalg import norm # type: ignore import scipy.sparse.linalg as spalg # type: ignore -from ..fdmath import fdfield_t +from ..fdmath import fdfield_t, cfdfield_t + logger = logging.getLogger(__name__) @@ -110,10 +111,10 @@ try: 'planner_effort': 'FFTW_PATIENT', } - def fftn(*args: Any, **kwargs: Any) -> NDArray[numpy.float64]: + def fftn(*args: Any, **kwargs: Any) -> NDArray[numpy.complex128]: return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args) - def ifftn(*args: Any, **kwargs: Any) -> NDArray[numpy.float64]: + def ifftn(*args: Any, **kwargs: Any) -> NDArray[numpy.complex128]: return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args) except ImportError: @@ -124,7 +125,7 @@ except ImportError: def generate_kmn( k0: ArrayLike, G_matrix: ArrayLike, - shape: ArrayLike, + shape: Sequence[int], ) -> 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. @@ -142,7 +143,7 @@ def generate_kmn( k0 = numpy.array(k0) Gi_grids = numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij') - Gi = numpy.stack(Gi_grids, axis=3) + Gi = numpy.moveaxis(Gi_grids, 0, -1) k_G = k0[None, None, None, :] - Gi k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2) @@ -169,7 +170,7 @@ def maxwell_operator( G_matrix: ArrayLike, epsilon: fdfield_t, mu: Optional[fdfield_t] = None - ) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]: + ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: """ Generate the Maxwell operator @@ -198,11 +199,11 @@ def maxwell_operator( shape = epsilon[0].shape + (1,) k_mag, m, n = generate_kmn(k0, G_matrix, shape) - epsilon = numpy.stack(epsilon, axis=3) + epsilon = numpy.moveaxis(epsilon, 0, -1) if mu is not None: - mu = numpy.stack(mu, axis=3) + mu = numpy.moveaxis(mu, 0, -1) - def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]: + def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]: """ Maxwell operator for Bloch eigenmode simulation. @@ -251,7 +252,7 @@ def hmn_2_exyz( k0: ArrayLike, G_matrix: ArrayLike, epsilon: fdfield_t, - ) -> Callable[[NDArray[numpy.float64]], fdfield_t]: + ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: """ Generate an operator which converts a vectorized spatial-frequency-space `h_mn` into an E-field distribution, i.e. @@ -272,11 +273,11 @@ def hmn_2_exyz( Function for converting `h_mn` into `E_xyz` """ shape = epsilon[0].shape + (1,) - epsilon = numpy.stack(epsilon, axis=3) + epsilon = numpy.moveaxis(epsilon, 0, -1) k_mag, m, n = generate_kmn(k0, G_matrix, shape) - def operator(h: NDArray[numpy.float64]) -> fdfield_t: + def operator(h: NDArray[numpy.complex128]) -> cfdfield_t: hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)] d_xyz = (n * hin_m - m * hin_n) * k_mag @@ -291,7 +292,7 @@ def hmn_2_hxyz( k0: ArrayLike, G_matrix: ArrayLike, epsilon: fdfield_t - ) -> Callable[[NDArray[numpy.float64]], fdfield_t]: + ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: """ Generate an operator which converts a vectorized spatial-frequency-space `h_mn` into an H-field distribution, i.e. @@ -314,7 +315,7 @@ def hmn_2_hxyz( shape = epsilon[0].shape + (1,) _k_mag, m, n = generate_kmn(k0, G_matrix, shape) - def operator(h: NDArray[numpy.float64]) -> fdfield_t: + def operator(h: NDArray[numpy.complex128]) -> cfdfield_t: hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)] h_xyz = (m * hin_m + n * hin_n) @@ -328,7 +329,7 @@ def inverse_maxwell_operator_approx( G_matrix: ArrayLike, epsilon: fdfield_t, mu: Optional[fdfield_t] = None, - ) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]: + ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: """ Generate an approximate inverse of the Maxwell operator, @@ -350,14 +351,13 @@ def inverse_maxwell_operator_approx( Function which applies the approximate inverse of the maxwell operator to `h_mn`. """ shape = epsilon[0].shape + (1,) - epsilon = numpy.stack(epsilon, axis=3) + epsilon = numpy.moveaxis(epsilon, 0, -1) + if mu is not None: + mu = numpy.moveaxis(mu, 0, -1) k_mag, m, n = generate_kmn(k0, G_matrix, shape) - if mu is not None: - mu = numpy.stack(mu, axis=3) - - def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]: + def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]: """ Approximate inverse Maxwell operator for Bloch eigenmode simulation. @@ -462,7 +462,7 @@ def eigsolve( tolerance: float = 1e-20, max_iters: int = 10000, reset_iters: int = 100, - ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64]]: + ) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: """ Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector k0 of the specified structure. @@ -697,11 +697,11 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to ''' def _rtrace_AtB( - A: NDArray[numpy.float64], - B: Union[NDArray[numpy.float64], float], + A: NDArray[numpy.complex128], + B: Union[NDArray[numpy.complex128], float], ) -> float: return real(numpy.sum(A.conj() * B)) -def _symmetrize(A: NDArray[numpy.float64]) -> NDArray[numpy.float64]: +def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]: return (A + A.conj().T) * 0.5 diff --git a/meanas/fdfd/farfield.py b/meanas/fdfd/farfield.py index e1a725d..bd02eb9 100644 --- a/meanas/fdfd/farfield.py +++ b/meanas/fdfd/farfield.py @@ -6,12 +6,12 @@ import numpy from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift from numpy import pi -from ..fdmath import fdfield_t +from ..fdmath import cfdfield_t def near_to_farfield( - E_near: fdfield_t, - H_near: fdfield_t, + E_near: cfdfield_t, + H_near: cfdfield_t, dx: float, dy: float, padded_size: List[int] = None @@ -122,8 +122,8 @@ def near_to_farfield( def far_to_nearfield( - E_far: fdfield_t, - H_far: fdfield_t, + E_far: cfdfield_t, + H_far: cfdfield_t, dkx: float, dky: float, padded_size: List[int] = None diff --git a/meanas/fdfd/functional.py b/meanas/fdfd/functional.py index 9a83910..1a1506b 100644 --- a/meanas/fdfd/functional.py +++ b/meanas/fdfd/functional.py @@ -2,13 +2,13 @@ Functional versions of many FDFD operators. These can be useful for performing FDFD calculations without needing to construct large matrices in memory. -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) +The functions generated here expect `cfdfield_t` inputs with shape (3, X, Y, Z), +e.g. E = [E_x, E_y, E_z] where each (complex) component has shape (X, Y, Z) """ from typing import Callable, Tuple, Optional import numpy -from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t +from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t from ..fdmath.functional import curl_forward, curl_back @@ -19,8 +19,8 @@ def e_full( omega: complex, dxes: dx_lists_t, epsilon: fdfield_t, - mu: fdfield_t = None - ) -> fdfield_updater_t: + mu: Optional[fdfield_t] = None + ) -> cfdfield_updater_t: """ Wave operator for use with E-field. See `operators.e_full` for details. @@ -37,13 +37,13 @@ def e_full( ch = curl_back(dxes[1]) ce = curl_forward(dxes[0]) - def op_1(e: fdfield_t) -> fdfield_t: + def op_1(e: cfdfield_t) -> cfdfield_t: curls = ch(ce(e)) - return curls - omega ** 2 * epsilon * e # type: ignore # issues with numpy/mypy + return curls - omega ** 2 * epsilon * e - def op_mu(e: fdfield_t) -> fdfield_t: - curls = ch(mu * ce(e)) - return curls - omega ** 2 * epsilon * e # type: ignore # issues with numpy/mypy + def op_mu(e: cfdfield_t) -> cfdfield_t: + curls = ch(mu * ce(e)) # type: ignore # mu = None ok because we don't return the function + return curls - omega ** 2 * epsilon * e if numpy.any(numpy.equal(mu, None)): return op_1 @@ -56,7 +56,7 @@ def eh_full( dxes: dx_lists_t, epsilon: fdfield_t, mu: fdfield_t = None - ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]: + ) -> Callable[[cfdfield_t, cfdfield_t], Tuple[cfdfield_t, cfdfield_t]]: """ Wave operator for full (both E and H) field representation. See `operators.eh_full`. @@ -74,13 +74,13 @@ def eh_full( ch = curl_back(dxes[1]) ce = curl_forward(dxes[0]) - def op_1(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]: + def op_1(e: cfdfield_t, h: cfdfield_t) -> Tuple[cfdfield_t, cfdfield_t]: return (ch(h) - 1j * omega * epsilon * e, - ce(e) + 1j * omega * h) # type: ignore # issues with numpy/mypy + ce(e) + 1j * omega * h) - def op_mu(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]: + def op_mu(e: cfdfield_t, h: cfdfield_t) -> Tuple[cfdfield_t, cfdfield_t]: return (ch(h) - 1j * omega * epsilon * e, - ce(e) + 1j * omega * mu * h) # type: ignore # issues with numpy/mypy + ce(e) + 1j * omega * mu * h) # type: ignore # mu=None ok if numpy.any(numpy.equal(mu, None)): return op_1 @@ -92,7 +92,7 @@ def e2h( omega: complex, dxes: dx_lists_t, mu: Optional[fdfield_t] = None, - ) -> fdfield_updater_t: + ) -> cfdfield_updater_t: """ Utility operator for converting the `E` field into the `H` field. For use with `e_full` -- assumes that there is no magnetic current `M`. @@ -108,11 +108,11 @@ def e2h( """ ce = curl_forward(dxes[0]) - def e2h_1_1(e: fdfield_t) -> fdfield_t: - return ce(e) / (-1j * omega) # type: ignore # issues with numpy/mypy + def e2h_1_1(e: cfdfield_t) -> cfdfield_t: + return ce(e) / (-1j * omega) - def e2h_mu(e: fdfield_t) -> fdfield_t: - return ce(e) / (-1j * omega * mu) # type: ignore # issues with numpy/mypy + def e2h_mu(e: cfdfield_t) -> cfdfield_t: + return ce(e) / (-1j * omega * mu) # type: ignore # mu=None ok if numpy.any(numpy.equal(mu, None)): return e2h_1_1 @@ -124,7 +124,7 @@ def m2j( omega: complex, dxes: dx_lists_t, mu: Optional[fdfield_t] = None, - ) -> fdfield_updater_t: + ) -> cfdfield_updater_t: """ Utility operator for converting magnetic current `M` distribution into equivalent electric current distribution `J`. @@ -141,13 +141,13 @@ def m2j( """ ch = curl_back(dxes[1]) - def m2j_mu(m: fdfield_t) -> fdfield_t: - J = ch(m / mu) / (-1j * omega) - return J # type: ignore # issues with numpy/mypy + def m2j_mu(m: cfdfield_t) -> cfdfield_t: + J = ch(m / mu) / (-1j * omega) # type: ignore # mu=None ok + return J - def m2j_1(m: fdfield_t) -> fdfield_t: + def m2j_1(m: cfdfield_t) -> cfdfield_t: J = ch(m) / (-1j * omega) - return J # type: ignore # issues with numpy/mypy + return J if numpy.any(numpy.equal(mu, None)): return m2j_1 @@ -161,7 +161,7 @@ def e_tfsf_source( dxes: dx_lists_t, epsilon: fdfield_t, mu: Optional[fdfield_t] = None, - ) -> fdfield_updater_t: + ) -> cfdfield_updater_t: """ Operator that turns an E-field distribution into a total-field/scattered-field (TFSF) source. @@ -182,13 +182,13 @@ def e_tfsf_source( # TODO documentation A = e_full(omega, dxes, epsilon, mu) - def op(e: fdfield_t) -> fdfield_t: + def op(e: cfdfield_t) -> cfdfield_t: neg_iwj = A(TF_region * e) - TF_region * A(e) return neg_iwj / (-1j * omega) return op -def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdfield_t]: +def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield_t, cfdfield_t], cfdfield_t]: """ Generates a function that takes the single-frequency `E` and `H` fields and calculates the cross product `E` x `H` = $E \\times H$ as required @@ -210,7 +210,7 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdf Returns: Function `f` that returns E x H as required for the poynting vector. """ - def exh(e: fdfield_t, h: fdfield_t) -> fdfield_t: + def exh(e: cfdfield_t, h: cfdfield_t) -> cfdfield_t: s = numpy.empty_like(e) ex = e[0] * dxes[0][0][:, None, None] ey = e[1] * dxes[0][1][None, :, None] diff --git a/meanas/fdfd/operators.py b/meanas/fdfd/operators.py index f36080b..ea45cba 100644 --- a/meanas/fdfd/operators.py +++ b/meanas/fdfd/operators.py @@ -31,7 +31,7 @@ from typing import Tuple, Optional import numpy import scipy.sparse as sparse # type: ignore -from ..fdmath import vec, dx_lists_t, vfdfield_t +from ..fdmath import vec, dx_lists_t, vfdfield_t, vcfdfield_t from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl_back @@ -91,7 +91,7 @@ def e_full( if numpy.any(numpy.equal(mu, None)): m_div = sparse.eye(epsilon.size) else: - m_div = sparse.diags(1 / mu) # type: ignore # checked mu is not None + m_div = sparse.diags(1 / mu) op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe return op @@ -275,7 +275,7 @@ def e2h( op = curl_forward(dxes[0]) / (-1j * omega) if not numpy.any(numpy.equal(mu, None)): - op = sparse.diags(1 / mu) @ op # type: ignore # checked mu is not None + op = sparse.diags(1 / mu) @ op if not numpy.any(numpy.equal(pmc, None)): op = sparse.diags(numpy.where(pmc, 0, 1)) @ op @@ -303,12 +303,12 @@ def m2j( op = curl_back(dxes[1]) / (1j * omega) if not numpy.any(numpy.equal(mu, None)): - op = op @ sparse.diags(1 / mu) # type: ignore # checked mu is not None + op = op @ sparse.diags(1 / mu) return op -def poynting_e_cross(e: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: +def poynting_e_cross(e: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: """ Operator for computing the Poynting vector, containing the (E x) portion of the Poynting vector. @@ -336,7 +336,7 @@ def poynting_e_cross(e: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: return P -def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: +def poynting_h_cross(h: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: """ Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector. diff --git a/meanas/fdfd/scpml.py b/meanas/fdfd/scpml.py index 0f9c92c..de38854 100644 --- a/meanas/fdfd/scpml.py +++ b/meanas/fdfd/scpml.py @@ -11,7 +11,7 @@ from numpy.typing import ArrayLike, NDArray __author__ = 'Jan Petykiewicz' -s_function_t = Callable[[float], float] +s_function_t = Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]] """Typedef for s-functions, see `prepare_s_function()`""" @@ -39,8 +39,8 @@ def prepare_s_function( def uniform_grid_scpml( - shape: ArrayLike, # ints - thicknesses: ArrayLike, # ints + shape: Sequence[int], + thicknesses: Sequence[int], omega: float, epsilon_effective: float = 1.0, s_function: Optional[s_function_t] = None, @@ -70,12 +70,11 @@ def uniform_grid_scpml( if s_function is None: s_function = prepare_s_function() + shape = tuple(shape) + thicknesses = tuple(thicknesses) + # Normalized distance to nearest boundary - def ll( - u: NDArray[numpy.float64], - n: NDArray[numpy.float64], - t: NDArray[numpy.float64], - ) -> NDArray[numpy.float64]: + def ll(u: NDArray[numpy.float64], n: int, t: int) -> NDArray[numpy.float64]: return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t dx_a = [numpy.array(numpy.inf)] * 3 diff --git a/meanas/fdfd/solvers.py b/meanas/fdfd/solvers.py index 0688966..6ced908 100644 --- a/meanas/fdfd/solvers.py +++ b/meanas/fdfd/solvers.py @@ -10,7 +10,7 @@ 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 +from ..fdmath import dx_lists_t, vfdfield_t, vcfdfield_t from . import operators @@ -65,7 +65,7 @@ def _scipy_qmr( def generic( omega: complex, dxes: dx_lists_t, - J: vfdfield_t, + J: vcfdfield_t, epsilon: vfdfield_t, mu: vfdfield_t = None, pec: vfdfield_t = None, @@ -73,7 +73,7 @@ def generic( adjoint: bool = False, matrix_solver: Callable[..., ArrayLike] = _scipy_qmr, matrix_solver_opts: Optional[Dict[str, Any]] = None, - ) -> vfdfield_t: + ) -> vcfdfield_t: """ Conjugate gradient FDFD solver using CSR sparse matrices. diff --git a/meanas/fdfd/waveguide_2d.py b/meanas/fdfd/waveguide_2d.py index 177d8d3..69d0b19 100644 --- a/meanas/fdfd/waveguide_2d.py +++ b/meanas/fdfd/waveguide_2d.py @@ -185,7 +185,7 @@ from numpy.linalg import norm import scipy.sparse as sparse # type: ignore from ..fdmath.operators import deriv_forward, deriv_back, cross -from ..fdmath import unvec, dx_lists_t, vfdfield_t +from ..fdmath import unvec, dx_lists_t, vfdfield_t, vcfdfield_t from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration @@ -335,7 +335,7 @@ def normalized_fields_e( epsilon: vfdfield_t, mu: Optional[vfdfield_t] = None, prop_phase: float = 0, - ) -> Tuple[vfdfield_t, vfdfield_t]: + ) -> Tuple[vcfdfield_t, vcfdfield_t]: """ Given a vector `e_xy` containing the vectorized E_x and E_y fields, returns normalized, vectorized E and H fields for the system. @@ -370,7 +370,7 @@ def normalized_fields_h( epsilon: vfdfield_t, mu: Optional[vfdfield_t] = None, prop_phase: float = 0, - ) -> Tuple[vfdfield_t, vfdfield_t]: + ) -> Tuple[vcfdfield_t, vcfdfield_t]: """ Given a vector `h_xy` containing the vectorized H_x and H_y fields, returns normalized, vectorized E and H fields for the system. @@ -398,14 +398,14 @@ def normalized_fields_h( def _normalized_fields( - e: ArrayLike, - h: ArrayLike, + e: vcfdfield_t, + h: vcfdfield_t, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, mu: Optional[vfdfield_t] = None, prop_phase: float = 0, - ) -> Tuple[vfdfield_t, vfdfield_t]: + ) -> Tuple[vcfdfield_t, vcfdfield_t]: # TODO documentation shape = [s.size for s in dxes[0]] dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)] @@ -581,7 +581,7 @@ def e2h( """ op = curl_e(wavenumber, dxes) / (-1j * omega) if not numpy.any(numpy.equal(mu, None)): - op = sparse.diags(1 / mu) @ op # type: ignore # checked that mu is not None + op = sparse.diags(1 / mu) @ op return op @@ -649,7 +649,7 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: def h_err( - h: vfdfield_t, + h: vcfdfield_t, wavenumber: complex, omega: complex, dxes: dx_lists_t, @@ -684,12 +684,12 @@ def h_err( def e_err( - e: vfdfield_t, + e: vcfdfield_t, wavenumber: complex, omega: complex, dxes: dx_lists_t, epsilon: vfdfield_t, - mu: vfdfield_t = Optional[None] + mu: Optional[vfdfield_t] = None, ) -> float: """ Calculates the relative error in the E field @@ -711,7 +711,7 @@ def e_err( if numpy.any(numpy.equal(mu, None)): op = ch @ ce @ e - omega ** 2 * (epsilon * e) else: - mu_inv = sparse.diags(1 / mu) # type: ignore # checked that mu is not None + mu_inv = sparse.diags(1 / mu) op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e) return norm(op) / norm(e) @@ -724,7 +724,7 @@ def solve_modes( epsilon: vfdfield_t, mu: vfdfield_t = None, mode_margin: int = 2, - ) -> Tuple[NDArray[numpy.float64], List[complex]]: + ) -> Tuple[NDArray[numpy.float64], NDArray[numpy.complex128]]: """ Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers. @@ -771,7 +771,7 @@ def solve_mode( mode_number: int, *args: Any, **kwargs: Any, - ) -> Tuple[vfdfield_t, complex]: + ) -> Tuple[vcfdfield_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 a6a2cba..9051123 100644 --- a/meanas/fdfd/waveguide_3d.py +++ b/meanas/fdfd/waveguide_3d.py @@ -8,7 +8,7 @@ from typing import Dict, Optional, Sequence, Union, Any import numpy from numpy.typing import NDArray -from ..fdmath import vec, unvec, dx_lists_t, fdfield_t +from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, cfdfield_t from . import operators, waveguide_2d @@ -106,7 +106,7 @@ def solve_mode( def compute_source( - E: fdfield_t, + E: cfdfield_t, wavenumber: complex, omega: complex, dxes: dx_lists_t, @@ -115,7 +115,7 @@ def compute_source( slices: Sequence[slice], epsilon: fdfield_t, mu: Optional[fdfield_t] = None, - ) -> fdfield_t: + ) -> cfdfield_t: """ Given an eigenmode obtained by `solve_mode`, returns the current source distribution necessary to position a unidirectional source at the slice location. @@ -152,13 +152,13 @@ def compute_source( def compute_overlap_e( - E: fdfield_t, + E: cfdfield_t, wavenumber: complex, dxes: dx_lists_t, axis: int, polarity: int, slices: Sequence[slice], - ) -> fdfield_t: # TODO DOCS + ) -> cfdfield_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) @@ -200,13 +200,13 @@ def compute_overlap_e( def expand_e( - E: fdfield_t, + E: cfdfield_t, wavenumber: complex, dxes: dx_lists_t, axis: int, polarity: int, slices: Sequence[slice], - ) -> fdfield_t: + ) -> cfdfield_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 a1a0633..2286d4c 100644 --- a/meanas/fdfd/waveguide_cyl.py +++ b/meanas/fdfd/waveguide_cyl.py @@ -12,7 +12,7 @@ from typing import Dict, Union import numpy import scipy.sparse as sparse # type: ignore -from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t +from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_t, cfdfield_t from ..fdmath.operators import deriv_forward, deriv_back from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration @@ -85,7 +85,7 @@ def solve_mode( dxes: dx_lists_t, epsilon: vfdfield_t, r0: float, - ) -> Dict[str, Union[complex, fdfield_t]]: + ) -> Dict[str, Union[complex, cfdfield_t]]: """ TODO: fixup Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode @@ -103,8 +103,8 @@ def solve_mode( Returns: ``` { - 'E': List[NDArray[numpy.float_]], - 'H': List[NDArray[numpy.float_]], + 'E': List[NDArray[numpy.complex_]], + 'H': List[NDArray[numpy.complex_]], 'wavenumber': complex, } ``` diff --git a/meanas/fdmath/__init__.py b/meanas/fdmath/__init__.py index 4b629f9..40c68e7 100644 --- a/meanas/fdmath/__init__.py +++ b/meanas/fdmath/__init__.py @@ -741,7 +741,8 @@ the true values can be multiplied back in after the simulation is complete if no normalized results are needed. """ -from .types import fdfield_t, vfdfield_t, dx_lists_t, dx_lists_mut, fdfield_updater_t +from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t, dx_lists_t, dx_lists_mut +from .types import fdfield_updater_t, cfdfield_updater_t from .vectorization import vec, unvec from . import operators, functional, types, vectorization diff --git a/meanas/fdmath/functional.py b/meanas/fdmath/functional.py index 2cc5172..27cd44e 100644 --- a/meanas/fdmath/functional.py +++ b/meanas/fdmath/functional.py @@ -24,7 +24,7 @@ def deriv_forward( Returns: List of functions for taking forward derivatives along each axis. """ - if dx_e: + if dx_e is not None: derivs = (lambda f: (numpy.roll(f, -1, axis=0) - f) / dx_e[0][:, None, None], lambda f: (numpy.roll(f, -1, axis=1) - f) / dx_e[1][None, :, None], lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :]) @@ -48,7 +48,7 @@ def deriv_back( Returns: List of functions for taking forward derivatives along each axis. """ - if dx_h: + if dx_h is not None: derivs = (lambda f: (f - numpy.roll(f, 1, axis=0)) / dx_h[0][:, None, None], lambda f: (f - numpy.roll(f, 1, axis=1)) / dx_h[1][None, :, None], lambda f: (f - numpy.roll(f, 1, axis=2)) / dx_h[2][None, None, :]) @@ -122,7 +122,7 @@ def curl_forward_parts( ) -> Callable: Dx, Dy, Dz = deriv_forward(dx_e) - def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]: + def mkparts_fwd(e: fdfield_t) -> Tuple[Tuple[fdfield_t, fdfield_t], ...]: return ((-Dz(e[1]), Dy(e[2])), ( Dz(e[0]), -Dx(e[2])), (-Dy(e[0]), Dx(e[1]))) @@ -135,7 +135,7 @@ def curl_back_parts( ) -> Callable: Dx, Dy, Dz = deriv_back(dx_h) - def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, ...]]: + def mkparts_back(h: fdfield_t) -> Tuple[Tuple[fdfield_t, fdfield_t], ...]: return ((-Dz(h[1]), Dy(h[2])), ( Dz(h[0]), -Dx(h[2])), (-Dy(h[0]), Dx(h[1]))) diff --git a/meanas/fdmath/types.py b/meanas/fdmath/types.py index 2dd0040..aae9594 100644 --- a/meanas/fdmath/types.py +++ b/meanas/fdmath/types.py @@ -13,6 +13,12 @@ fdfield_t = NDArray[numpy.float_] vfdfield_t = NDArray[numpy.float_] """Linearized vector field (single vector of length 3*X*Y*Z)""" +cfdfield_t = NDArray[numpy.complex_] +"""Complex vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)""" + +vcfdfield_t = NDArray[numpy.complex_] +"""Linearized complex vector field (single vector of length 3*X*Y*Z)""" + dx_lists_t = Sequence[Sequence[NDArray[numpy.float_]]] """ @@ -31,3 +37,6 @@ dx_lists_mut = MutableSequence[MutableSequence[NDArray[numpy.float_]]] fdfield_updater_t = Callable[..., fdfield_t] """Convenience type for functions which take and return an fdfield_t""" + +cfdfield_updater_t = Callable[..., cfdfield_t] +"""Convenience type for functions which take and return an cfdfield_t""" diff --git a/meanas/fdmath/vectorization.py b/meanas/fdmath/vectorization.py index 6b6c49e..ef97f7c 100644 --- a/meanas/fdmath/vectorization.py +++ b/meanas/fdmath/vectorization.py @@ -4,11 +4,11 @@ and a 1D array representation of that field `[f_x0, f_x1, f_x2,... f_y0,... f_z0 Vectorized versions of the field use row-major (ie., C-style) ordering. """ -from typing import Optional, overload, Union, List +from typing import Optional, overload, Union, Sequence import numpy from numpy.typing import ArrayLike -from .types import fdfield_t, vfdfield_t +from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t @overload @@ -16,10 +16,18 @@ def vec(f: None) -> None: pass @overload -def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t: +def vec(f: fdfield_t) -> vfdfield_t: pass -def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]: +@overload +def vec(f: cfdfield_t) -> vcfdfield_t: + pass + +@overload +def vec(f: ArrayLike) -> Union[vfdfield_t, vcfdfield_t]: + pass + +def vec(f: Union[fdfield_t, cfdfield_t, ArrayLike, None]) -> Union[vfdfield_t, vcfdfield_t, None]: """ Create a 1D ndarray from a 3D vector field which spans a 1-3D region. @@ -38,14 +46,18 @@ def vec(f: Optional[Union[fdfield_t, List[ArrayLike]]]) -> Optional[vfdfield_t]: @overload -def unvec(v: None, shape: ArrayLike) -> None: +def unvec(v: None, shape: Sequence[int]) -> None: pass @overload -def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t: +def unvec(v: vfdfield_t, shape: Sequence[int]) -> fdfield_t: pass -def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> Optional[fdfield_t]: +@overload +def unvec(v: vcfdfield_t, shape: Sequence[int]) -> cfdfield_t: + pass + +def unvec(v: Union[vfdfield_t, vcfdfield_t, None], shape: Sequence[int]) -> Union[fdfield_t, cfdfield_t, None]: """ 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 @@ -62,5 +74,5 @@ def unvec(v: Optional[vfdfield_t], shape: ArrayLike) -> Optional[fdfield_t]: """ if numpy.any(numpy.equal(v, None)): return None - return v.reshape((3, *shape), order='C') # type: ignore # already check v is not None + return v.reshape((3, *shape), order='C') diff --git a/meanas/fdtd/pml.py b/meanas/fdtd/pml.py index 6f7aff7..1781485 100644 --- a/meanas/fdtd/pml.py +++ b/meanas/fdtd/pml.py @@ -8,8 +8,9 @@ PML implementations # TODO retest pmls! from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional +from copy import deepcopy import numpy -from typing import NDArray +from numpy.typing import NDArray, DTypeLike from ..fdmath import fdfield_t, dx_lists_t from ..fdmath.functional import deriv_forward, deriv_back @@ -97,34 +98,38 @@ def updates_with_cpml( dxes: dx_lists_t, epsilon: fdfield_t, *, - dtype: numpy.dtype = numpy.float32, - ) -> Tuple[Callable[[fdfield_t, fdfield_t], None], - Callable[[fdfield_t, fdfield_t], None]]: + dtype: DTypeLike = numpy.float32, + ) -> Tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None], + Callable[[fdfield_t, fdfield_t, fdfield_t], None]]: Dfx, Dfy, Dfz = deriv_forward(dxes[1]) Dbx, Dby, Dbz = deriv_back(dxes[1]) - psi_E = [[None, None], [None, None], [None, None]] - psi_H = [[None, None], [None, None], [None, None]] - params_E = [[None, None], [None, None], [None, None]] - params_H = [[None, None], [None, None], [None, None]] + + psi_E: List[List[Tuple[Any, Any]]] = [[(None, None) for _ in range(2)] for _ in range(3)] + psi_H: List[List[Tuple[Any, Any]]] = deepcopy(psi_E) + params_E: List[List[Tuple[Any, Any, Any, Any]]] = [[(None, None, None, None) for _ in range(2)] for _ in range(3)] + params_H: List[List[Tuple[Any, Any, Any, Any]]] = deepcopy(params_E) for axis in range(3): for pp, polarity in enumerate((-1, 1)): - if cpml_params[axis][pp] is None: + cpml_param = cpml_params[axis][pp] + if cpml_param is None: psi_E[axis][pp] = (None, None) psi_H[axis][pp] = (None, None) continue - cpml_param = cpml_params[axis][pp] - region = cpml_param['region'] region_shape = epsilon[0][region].shape - psi_E[axis][pp] = (numpy.zeros(region_shape, dtype=dtype), - numpy.zeros(region_shape, dtype=dtype)) - psi_H[axis][pp] = (numpy.zeros(region_shape, dtype=dtype), - numpy.zeros(region_shape, dtype=dtype)) + psi_E[axis][pp] = ( + numpy.zeros(region_shape, dtype=dtype), + numpy.zeros(region_shape, dtype=dtype), + ) + psi_H[axis][pp] = ( + numpy.zeros(region_shape, dtype=dtype), + numpy.zeros(region_shape, dtype=dtype), + ) params_E[axis][pp] = cpml_param['param_e'] + (region,) params_H[axis][pp] = cpml_param['param_h'] + (region,) @@ -132,7 +137,11 @@ def updates_with_cpml( pE = numpy.empty_like(epsilon, dtype=dtype) pH = numpy.empty_like(epsilon, dtype=dtype) - def update_E(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t) -> None: + def update_E( + e: fdfield_t, + h: fdfield_t, + epsilon: fdfield_t, + ) -> None: dyHx = Dby(h[0]) dzHx = Dbz(h[0]) dxHy = Dbx(h[1]) @@ -175,7 +184,11 @@ def updates_with_cpml( e[2] += dt / epsilon[2] * (dxHy - dyHx + pE[2]) - def update_H(e: fdfield_t, h: fdfield_t, mu: fdfield_t = (1, 1, 1)) -> None: + def update_H( + e: fdfield_t, + h: fdfield_t, + mu: fdfield_t = numpy.ones(3), + ) -> None: dyEx = Dfy(e[0]) dzEx = Dfz(e[0]) dxEy = Dfx(e[1]) diff --git a/meanas/test/test_fdfd.py b/meanas/test/test_fdfd.py index fef80fd..e5a5875 100644 --- a/meanas/test/test_fdfd.py +++ b/meanas/test/test_fdfd.py @@ -99,8 +99,8 @@ class FDResult: dxes: List[List[NDArray[numpy.float64]]] epsilon: NDArray[numpy.float64] omega: complex - j: NDArray[numpy.float64] - e: NDArray[numpy.float64] + j: NDArray[numpy.complex128] + e: NDArray[numpy.complex128] pmc: Optional[NDArray[numpy.float64]] pec: Optional[NDArray[numpy.float64]] @@ -111,7 +111,7 @@ def sim( shape: Tuple[int, ...], epsilon: NDArray[numpy.float64], dxes: List[List[NDArray[numpy.float64]]], - j_distribution: NDArray[numpy.float64], + j_distribution: NDArray[numpy.complex128], omega: float, pec: Optional[NDArray[numpy.float64]], pmc: Optional[NDArray[numpy.float64]], @@ -134,8 +134,13 @@ def sim( 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_fdfd_pml.py b/meanas/test/test_fdfd_pml.py index 30eb32d..ff6e4c2 100644 --- a/meanas/test/test_fdfd_pml.py +++ b/meanas/test/test_fdfd_pml.py @@ -78,7 +78,7 @@ def j_distribution( dxes: dx_lists_mut, omega: float, src_polarity: int, - ) -> Iterable[NDArray[numpy.float64]]: + ) -> Iterable[NDArray[numpy.complex128]]: j = numpy.zeros(shape, dtype=complex) dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis @@ -150,7 +150,7 @@ def sim( shape: Tuple[int, ...], epsilon: NDArray[numpy.float64], dxes: dx_lists_mut, - j_distribution: NDArray[numpy.float64], + j_distribution: NDArray[numpy.complex128], omega: float, pec: Optional[NDArray[numpy.float64]], pmc: Optional[NDArray[numpy.float64]], diff --git a/meanas/test/utils.py b/meanas/test/utils.py index b4cb3ab..8d26d11 100644 --- a/meanas/test/utils.py +++ b/meanas/test/utils.py @@ -1,14 +1,15 @@ from typing import Any + import numpy -from typing import ArrayLike +from numpy.typing import ArrayLike, NDArray PRNG = numpy.random.RandomState(12345) def assert_fields_close( - x: ArrayLike, - y: ArrayLike, + x: NDArray, + y: NDArray, *args: Any, **kwargs: Any, ) -> None: @@ -18,8 +19,8 @@ def assert_fields_close( numpy.rollaxis(y, -1)), *args, **kwargs) def assert_close( - x: ArrayLike, - y: ArrayLike, + x: NDArray, + y: NDArray, *args: Any, **kwargs: Any, ) -> None: