From b486fa325b11976b6463251cbb42eeaf346ef22f Mon Sep 17 00:00:00 2001 From: jan Date: Wed, 10 Dec 2025 02:14:20 -0800 Subject: [PATCH] Rework field types, use sparse arrays instead of matrices, rework eme arg naming, improve type annotations and linter cleanup --- meanas/eigensolvers.py | 8 +- meanas/fdfd/bloch.py | 49 +++++--- meanas/fdfd/eme.py | 44 ++++--- meanas/fdfd/farfield.py | 10 +- meanas/fdfd/functional.py | 68 +++++------ meanas/fdfd/operators.py | 142 +++++++++++----------- meanas/fdfd/solvers.py | 20 ++-- meanas/fdfd/waveguide_2d.py | 212 ++++++++++++++++----------------- meanas/fdfd/waveguide_3d.py | 33 ++--- meanas/fdfd/waveguide_cyl.py | 79 ++++++------ meanas/fdmath/__init__.py | 34 +++++- meanas/fdmath/operators.py | 51 ++++---- meanas/fdmath/types.py | 61 +++++++++- meanas/fdmath/vectorization.py | 58 +++++++-- meanas/fdtd/energy.py | 84 ++++++------- meanas/fdtd/misc.py | 7 +- meanas/fdtd/pml.py | 4 +- meanas/test/test_fdfd.py | 26 ++-- meanas/test/test_fdfd_pml.py | 12 +- meanas/test/test_fdtd.py | 2 +- 20 files changed, 571 insertions(+), 433 deletions(-) diff --git a/meanas/eigensolvers.py b/meanas/eigensolvers.py index e8630aa..21e2ec0 100644 --- a/meanas/eigensolvers.py +++ b/meanas/eigensolvers.py @@ -64,10 +64,10 @@ def rayleigh_quotient_iteration( (eigenvalues, eigenvectors) """ try: - (operator - sparse.eye(operator.shape[0])) + (operator - sparse.eye_array(operator.shape[0])) - def shift(eigval: float) -> sparse: - return eigval * sparse.eye(operator.shape[0]) + def shift(eigval: float) -> sparse.sparray: + return eigval * sparse.eye_array(operator.shape[0]) if solver is None: solver = spalg.spsolve @@ -130,7 +130,7 @@ def signed_eigensolve( # Try to combine, use general LinearOperator if we fail try: - shifted_operator = operator + shift * sparse.eye(operator.shape[0]) + shifted_operator = operator + shift * sparse.eye_array(operator.shape[0]) except TypeError: shifted_operator = operator + spalg.LinearOperator(shape=operator.shape, matvec=lambda v: shift * v) diff --git a/meanas/fdfd/bloch.py b/meanas/fdfd/bloch.py index deb6ec6..4eedcc4 100644 --- a/meanas/fdfd/bloch.py +++ b/meanas/fdfd/bloch.py @@ -106,7 +106,7 @@ import scipy.optimize from scipy.linalg import norm import scipy.sparse.linalg as spalg -from ..fdmath import fdfield_t, cfdfield_t +from ..fdmath import fdfield, cfdfield, cfdfield_t logger = logging.getLogger(__name__) @@ -183,8 +183,8 @@ def generate_kmn( def maxwell_operator( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield_t, - mu: fdfield_t | None = None + epsilon: fdfield, + mu: fdfield | None = None ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: """ Generate the Maxwell operator @@ -276,7 +276,7 @@ def maxwell_operator( def hmn_2_exyz( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield_t, + epsilon: fdfield, ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: """ Generate an operator which converts a vectorized spatial-frequency-space @@ -308,7 +308,8 @@ def hmn_2_exyz( - m * hin_n) * k_mag # divide by epsilon - return numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0) + exyz = numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0) + return cfdfield_t(exyz) return operator @@ -316,7 +317,7 @@ def hmn_2_exyz( def hmn_2_hxyz( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield_t + epsilon: fdfield, ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: """ Generate an operator which converts a vectorized spatial-frequency-space @@ -343,8 +344,8 @@ def hmn_2_hxyz( def operator(h: NDArray[numpy.complex128]) -> cfdfield_t: hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2)) h_xyz = (m * hin_m - + n * hin_n) # noqa: E128 - return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)]) + + n * hin_n) + return cfdfield_t(numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])) return operator @@ -352,8 +353,8 @@ def hmn_2_hxyz( def inverse_maxwell_operator_approx( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = None, ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: """ Generate an approximate inverse of the Maxwell operator, @@ -440,8 +441,8 @@ def find_k( tolerance: float, direction: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = None, band: int = 0, k_bounds: tuple[float, float] = (0, 0.5), k_guess: float | None = None, @@ -508,8 +509,8 @@ def eigsolve( num_modes: int, k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = None, tolerance: float = 1e-7, max_iters: int = 10000, reset_iters: int = 100, @@ -649,7 +650,7 @@ def eigsolve( def Qi_func(theta: float, Qi_memo=Qi_memo, ZtZ=ZtZ, DtD=DtD, symZtD=symZtD) -> float: # noqa: ANN001 if Qi_memo[0] == theta: - return cast(float, Qi_memo[1]) + return cast('float', Qi_memo[1]) c = numpy.cos(theta) s = numpy.sin(theta) @@ -668,8 +669,8 @@ def eigsolve( else: raise Exception('Inexplicable singularity in trace_func') from err Qi_memo[0] = theta - Qi_memo[1] = cast(float, Qi) - return cast(float, Qi) + Qi_memo[1] = cast('float', Qi) + return cast('float', Qi) def trace_func(theta: float, ZtAZ=ZtAZ, DtAD=DtAD, symZtAD=symZtAD) -> float: # noqa: ANN001 c = numpy.cos(theta) @@ -801,7 +802,12 @@ def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]: -def inner_product(eL, hL, eR, hR) -> complex: +def inner_product( + eL: cfdfield, + hL: cfdfield, + eR: cfdfield, + hR: cfdfield, + ) -> complex: # assumes x-axis propagation assert numpy.array_equal(eR.shape, hR.shape) @@ -829,7 +835,12 @@ def inner_product(eL, hL, eR, hR) -> complex: return eRxhL_x.sum() - eLxhR_x.sum() -def trq(eI, hI, eO, hO) -> tuple[complex, complex]: +def trq( + eI: cfdfield, + hI: cfdfield, + eO: cfdfield, + hO: cfdfield, + ) -> tuple[complex, complex]: pp = inner_product(eO, hO, eI, hI) pn = inner_product(eO, hO, eI, -hI) np = inner_product(eO, -hO, eI, hI) diff --git a/meanas/fdfd/eme.py b/meanas/fdfd/eme.py index 35e1e90..f834973 100644 --- a/meanas/fdfd/eme.py +++ b/meanas/fdfd/eme.py @@ -1,20 +1,29 @@ +from collections.abc import Sequence import numpy +from numpy.typing import NDArray +from scipy import sparse -from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t +from ..fdmath import dx_lists2_t, vcfdfield2 from .waveguide_2d import inner_product -def get_tr(ehL, wavenumbers_L, ehR, wavenumbers_R, dxes: dx_lists_t): +def get_tr( + ehLs: Sequence[Sequence[vcfdfield2]], + wavenumbers_L: Sequence[complex], + ehRs: Sequence[Sequence[vcfdfield2]], + wavenumbers_R: Sequence[complex], + dxes: dx_lists2_t, + ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: nL = len(wavenumbers_L) nR = len(wavenumbers_R) A12 = numpy.zeros((nL, nR), dtype=complex) A21 = numpy.zeros((nL, nR), dtype=complex) B11 = numpy.zeros((nL,), dtype=complex) for ll in range(nL): - eL, hL = ehL[ll] + eL, hL = ehLs[ll] B11[ll] = inner_product(eL, hL, dxes=dxes, conj_h=False) for rr in range(nR): - eR, hR = ehR[rr] + eR, hR = ehRs[rr] A12[ll, rr] = inner_product(eL, hR, dxes=dxes, conj_h=False) # TODO optimize loop? A21[ll, rr] = inner_product(eR, hL, dxes=dxes, conj_h=False) @@ -32,9 +41,15 @@ def get_tr(ehL, wavenumbers_L, ehR, wavenumbers_R, dxes: dx_lists_t): return tt, rr -def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs): - t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs) - t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs) +def get_abcd( + ehLs: Sequence[Sequence[vcfdfield2]], + wavenumbers_L: Sequence[complex], + ehRs: Sequence[Sequence[vcfdfield2]], + wavenumbers_R: Sequence[complex], + **kwargs, + ) -> sparse.sparray: + t12, r12 = get_tr(ehLs, wavenumbers_L, ehRs, wavenumbers_R, **kwargs) + t21, r21 = get_tr(ehRs, wavenumbers_R, ehLs, wavenumbers_L, **kwargs) t21i = numpy.linalg.pinv(t21) A = t12 - r21 @ t21i @ r12 B = r21 @ t21i @@ -44,15 +59,16 @@ def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs): def get_s( - eL_xys, - wavenumbers_L, - eR_xys, - wavenumbers_R, + ehLs: Sequence[Sequence[vcfdfield2]], + wavenumbers_L: Sequence[complex], + ehRs: Sequence[Sequence[vcfdfield2]], + wavenumbers_R: Sequence[complex], force_nogain: bool = False, force_reciprocal: bool = False, - **kwargs): - t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs) - t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs) + **kwargs, + ) -> NDArray[numpy.complex128]: + t12, r12 = get_tr(ehLs, wavenumbers_L, ehRs, wavenumbers_R, **kwargs) + t21, r21 = get_tr(ehRs, wavenumbers_R, ehLs, wavenumbers_L, **kwargs) ss = numpy.block([[r12, t12], [t21, r21]]) diff --git a/meanas/fdfd/farfield.py b/meanas/fdfd/farfield.py index 4829d86..86ec0d7 100644 --- a/meanas/fdfd/farfield.py +++ b/meanas/fdfd/farfield.py @@ -1,14 +1,16 @@ """ Functions for performing near-to-farfield transformation (and the reverse). """ -from typing import Any, cast -from collections.abc import Sequence +from typing import Any, cast, TYPE_CHECKING import numpy from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift from numpy import pi from ..fdmath import cfdfield_t +if TYPE_CHECKING: + from collections.abc import Sequence + def near_to_farfield( E_near: cfdfield_t, @@ -63,7 +65,7 @@ def near_to_farfield( padded_size = (2**numpy.ceil(numpy.log2(s))).astype(int) if not hasattr(padded_size, '__len__'): padded_size = (padded_size, padded_size) # type: ignore # checked if sequence - padded_shape = cast(Sequence[int], padded_size) + padded_shape = cast('Sequence[int]', padded_size) En_fft = [fftshift(fft2(fftshift(Eni), s=padded_shape)) for Eni in E_near] Hn_fft = [fftshift(fft2(fftshift(Hni), s=padded_shape)) for Hni in H_near] @@ -172,7 +174,7 @@ def far_to_nearfield( padded_size = (2 ** numpy.ceil(numpy.log2(s))).astype(int) if not hasattr(padded_size, '__len__'): padded_size = (padded_size, padded_size) # type: ignore # checked if sequence - padded_shape = cast(Sequence[int], padded_size) + padded_shape = cast('Sequence[int]', padded_size) k = 2 * pi kxs = fftshift(fftfreq(s[0], 1 / (s[0] * dkx))) diff --git a/meanas/fdfd/functional.py b/meanas/fdfd/functional.py index f4a250f..440daf2 100644 --- a/meanas/fdfd/functional.py +++ b/meanas/fdfd/functional.py @@ -8,7 +8,7 @@ e.g. E = [E_x, E_y, E_z] where each (complex) component has shape (X, Y, Z) from collections.abc import Callable import numpy -from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t +from ..fdmath import dx_lists_t, cfdfield_t, fdfield, cfdfield, cfdfield_updater_t from ..fdmath.functional import curl_forward, curl_back @@ -18,8 +18,8 @@ __author__ = 'Jan Petykiewicz' def e_full( omega: complex, dxes: dx_lists_t, - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = 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: cfdfield_t) -> cfdfield_t: + def op_1(e: cfdfield) -> cfdfield_t: curls = ch(ce(e)) - return curls - omega ** 2 * epsilon * e + return cfdfield_t(curls - omega ** 2 * epsilon * e) - def op_mu(e: cfdfield_t) -> cfdfield_t: + def op_mu(e: cfdfield) -> 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 + return cfdfield_t(curls - omega ** 2 * epsilon * e) if mu is None: return op_1 @@ -53,9 +53,9 @@ def e_full( def eh_full( omega: complex, dxes: dx_lists_t, - epsilon: fdfield_t, - mu: fdfield_t | None = None, - ) -> Callable[[cfdfield_t, cfdfield_t], tuple[cfdfield_t, cfdfield_t]]: + epsilon: fdfield, + mu: fdfield | None = None, + ) -> Callable[[cfdfield, cfdfield], tuple[cfdfield_t, cfdfield_t]]: """ Wave operator for full (both E and H) field representation. See `operators.eh_full`. @@ -73,13 +73,13 @@ def eh_full( ch = curl_back(dxes[1]) ce = curl_forward(dxes[0]) - def op_1(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]: - return (ch(h) - 1j * omega * epsilon * e, - ce(e) + 1j * omega * h) + def op_1(e: cfdfield, h: cfdfield) -> tuple[cfdfield_t, cfdfield_t]: + return (cfdfield_t(ch(h) - 1j * omega * epsilon * e), + cfdfield_t(ce(e) + 1j * omega * h)) - def op_mu(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]: - return (ch(h) - 1j * omega * epsilon * e, - ce(e) + 1j * omega * mu * h) # type: ignore # mu=None ok + def op_mu(e: cfdfield, h: cfdfield) -> tuple[cfdfield_t, cfdfield_t]: + return (cfdfield_t(ch(h) - 1j * omega * epsilon * e), + cfdfield_t(ce(e) + 1j * omega * mu * h)) # type: ignore # mu=None ok if mu is None: return op_1 @@ -89,7 +89,7 @@ def eh_full( def e2h( omega: complex, dxes: dx_lists_t, - mu: fdfield_t | None = None, + mu: fdfield | None = None, ) -> cfdfield_updater_t: """ Utility operator for converting the `E` field into the `H` field. @@ -106,11 +106,11 @@ def e2h( """ ce = curl_forward(dxes[0]) - def e2h_1_1(e: cfdfield_t) -> cfdfield_t: - return ce(e) / (-1j * omega) + def e2h_1_1(e: cfdfield) -> cfdfield_t: + return cfdfield_t(ce(e) / (-1j * omega)) - def e2h_mu(e: cfdfield_t) -> cfdfield_t: - return ce(e) / (-1j * omega * mu) # type: ignore # mu=None ok + def e2h_mu(e: cfdfield) -> cfdfield_t: + return cfdfield_t(ce(e) / (-1j * omega * mu)) # type: ignore # mu=None ok if mu is None: return e2h_1_1 @@ -120,7 +120,7 @@ def e2h( def m2j( omega: complex, dxes: dx_lists_t, - mu: fdfield_t | None = None, + mu: fdfield | None = None, ) -> cfdfield_updater_t: """ Utility operator for converting magnetic current `M` distribution @@ -138,13 +138,13 @@ def m2j( """ ch = curl_back(dxes[1]) - def m2j_mu(m: cfdfield_t) -> cfdfield_t: + def m2j_mu(m: cfdfield) -> cfdfield_t: J = ch(m / mu) / (-1j * omega) # type: ignore # mu=None ok - return J + return cfdfield_t(J) - def m2j_1(m: cfdfield_t) -> cfdfield_t: + def m2j_1(m: cfdfield) -> cfdfield_t: J = ch(m) / (-1j * omega) - return J + return cfdfield_t(J) if mu is None: return m2j_1 @@ -152,11 +152,11 @@ def m2j( def e_tfsf_source( - TF_region: fdfield_t, + TF_region: fdfield, omega: complex, dxes: dx_lists_t, - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = None, ) -> cfdfield_updater_t: """ Operator that turns an E-field distribution into a total-field/scattered-field @@ -178,13 +178,13 @@ def e_tfsf_source( # TODO documentation A = e_full(omega, dxes, epsilon, mu) - def op(e: cfdfield_t) -> cfdfield_t: + def op(e: cfdfield) -> cfdfield_t: neg_iwj = A(TF_region * e) - TF_region * A(e) - return neg_iwj / (-1j * omega) + return cfdfield_t(neg_iwj / (-1j * omega)) return op -def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield_t, cfdfield_t], cfdfield_t]: +def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield, cfdfield], cfdfield_t]: r""" 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 @@ -206,7 +206,7 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield_t, cfdfield_t], c Returns: Function `f` that returns E x H as required for the poynting vector. """ - def exh(e: cfdfield_t, h: cfdfield_t) -> cfdfield_t: + def exh(e: cfdfield, h: cfdfield) -> cfdfield_t: s = numpy.empty_like(e) ex = e[0] * dxes[0][0][:, None, None] ey = e[1] * dxes[0][1][None, :, None] @@ -217,5 +217,5 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield_t, cfdfield_t], c s[0] = numpy.roll(ey, -1, axis=0) * hz - numpy.roll(ez, -1, axis=0) * hy s[1] = numpy.roll(ez, -1, axis=1) * hx - numpy.roll(ex, -1, axis=1) * hz s[2] = numpy.roll(ex, -1, axis=2) * hy - numpy.roll(ey, -1, axis=2) * hx - return s + return cfdfield_t(s) return exh diff --git a/meanas/fdfd/operators.py b/meanas/fdfd/operators.py index 8c16ef7..829f43e 100644 --- a/meanas/fdfd/operators.py +++ b/meanas/fdfd/operators.py @@ -1,7 +1,7 @@ """ Sparse matrix operators for use with electromagnetic wave equations. -These functions return sparse-matrix (`scipy.sparse.spmatrix`) representations of +These functions return sparse-matrix (`scipy.sparse.sparray`) representations of a variety of operators, intended for use with E and H fields vectorized using the `meanas.fdmath.vectorization.vec()` and `meanas.fdmath.vectorization.unvec()` functions. @@ -30,7 +30,7 @@ The following operators are included: import numpy from scipy import sparse -from ..fdmath import vec, dx_lists_t, vfdfield_t, vcfdfield_t +from ..fdmath import vec, dx_lists_t, vfdfield, vcfdfield from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl_back @@ -40,11 +40,11 @@ __author__ = 'Jan Petykiewicz' def e_full( omega: complex, dxes: dx_lists_t, - epsilon: vfdfield_t | vcfdfield_t, - mu: vfdfield_t | None = None, - pec: vfdfield_t | None = None, - pmc: vfdfield_t | None = None, - ) -> sparse.spmatrix: + epsilon: vfdfield | vcfdfield, + mu: vfdfield | None = None, + pec: vfdfield | None = None, + pmc: vfdfield | None = None, + ) -> sparse.sparray: r""" Wave operator $$ \nabla \times (\frac{1}{\mu} \nabla \times) - \Omega^2 \epsilon $$ @@ -77,20 +77,20 @@ def e_full( ce = curl_forward(dxes[0]) if pec is None: - pe = sparse.eye(epsilon.size) + pe = sparse.eye_array(epsilon.size) else: - pe = sparse.diags(numpy.where(pec, 0, 1)) # Set pe to (not PEC) + pe = sparse.diags_array(numpy.where(pec, 0, 1)) # Set pe to (not PEC) if pmc is None: - pm = sparse.eye(epsilon.size) + pm = sparse.eye_array(epsilon.size) else: - pm = sparse.diags(numpy.where(pmc, 0, 1)) # set pm to (not PMC) + pm = sparse.diags_array(numpy.where(pmc, 0, 1)) # set pm to (not PMC) - e = sparse.diags(epsilon) + e = sparse.diags_array(epsilon) if mu is None: - m_div = sparse.eye(epsilon.size) + m_div = sparse.eye_array(epsilon.size) else: - m_div = sparse.diags(1 / mu) + m_div = sparse.diags_array(1 / mu) op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe return op @@ -98,7 +98,7 @@ def e_full( def e_full_preconditioners( dxes: dx_lists_t, - ) -> tuple[sparse.spmatrix, sparse.spmatrix]: + ) -> tuple[sparse.sparray, sparse.sparray]: """ Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator. @@ -118,19 +118,19 @@ def e_full_preconditioners( dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[0][2][None, None, :]] p_vector = numpy.sqrt(vec(p_squared)) - P_left = sparse.diags(p_vector) - P_right = sparse.diags(1 / p_vector) + P_left = sparse.diags_array(p_vector) + P_right = sparse.diags_array(1 / p_vector) return P_left, P_right def h_full( omega: complex, dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, - pec: vfdfield_t | None = None, - pmc: vfdfield_t | None = None, - ) -> sparse.spmatrix: + epsilon: vfdfield, + mu: vfdfield | None = None, + pec: vfdfield | None = None, + pmc: vfdfield | None = None, + ) -> sparse.sparray: r""" Wave operator $$ \nabla \times (\frac{1}{\epsilon} \nabla \times) - \omega^2 \mu $$ @@ -161,20 +161,20 @@ def h_full( ce = curl_forward(dxes[0]) if pec is None: - pe = sparse.eye(epsilon.size) + pe = sparse.eye_array(epsilon.size) else: - pe = sparse.diags(numpy.where(pec, 0, 1)) # set pe to (not PEC) + pe = sparse.diags_array(numpy.where(pec, 0, 1)) # set pe to (not PEC) if pmc is None: - pm = sparse.eye(epsilon.size) + pm = sparse.eye_array(epsilon.size) else: - pm = sparse.diags(numpy.where(pmc, 0, 1)) # Set pe to (not PMC) + pm = sparse.diags_array(numpy.where(pmc, 0, 1)) # Set pe to (not PMC) - e_div = sparse.diags(1 / epsilon) + e_div = sparse.diags_array(1 / epsilon) if mu is None: - m = sparse.eye(epsilon.size) + m = sparse.eye_array(epsilon.size) else: - m = sparse.diags(mu) + m = sparse.diags_array(mu) A = pm @ (ce @ pe @ e_div @ ch - omega**2 * m) @ pm return A @@ -183,11 +183,11 @@ def h_full( def eh_full( omega: complex, dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, - pec: vfdfield_t | None = None, - pmc: vfdfield_t | None = None, - ) -> sparse.spmatrix: + epsilon: vfdfield, + mu: vfdfield | None = None, + pec: vfdfield | None = None, + pmc: vfdfield | None = None, + ) -> sparse.sparray: r""" Wave operator for `[E, H]` field representation. This operator implements Maxwell's equations without cancelling out either E or H. The operator is @@ -227,35 +227,35 @@ def eh_full( Sparse matrix containing the wave operator. """ if pec is None: - pe = sparse.eye(epsilon.size) + pe = sparse.eye_array(epsilon.size) else: - pe = sparse.diags(numpy.where(pec, 0, 1)) # set pe to (not PEC) + pe = sparse.diags_array(numpy.where(pec, 0, 1)) # set pe to (not PEC) if pmc is None: - pm = sparse.eye(epsilon.size) + pm = sparse.eye_array(epsilon.size) else: - pm = sparse.diags(numpy.where(pmc, 0, 1)) # set pm to (not PMC) + pm = sparse.diags_array(numpy.where(pmc, 0, 1)) # set pm to (not PMC) - iwe = pe @ (1j * omega * sparse.diags(epsilon)) @ pe + iwe = pe @ (1j * omega * sparse.diags_array(epsilon)) @ pe iwm = 1j * omega if mu is not None: - iwm *= sparse.diags(mu) + iwm *= sparse.diags_array(mu) iwm = pm @ iwm @ pm A1 = pe @ curl_back(dxes[1]) @ pm A2 = pm @ curl_forward(dxes[0]) @ pe - A = sparse.bmat([[-iwe, A1], - [A2, iwm]]) + A = sparse.block_array([[-iwe, A1], + [A2, iwm]]) return A def e2h( omega: complex, dxes: dx_lists_t, - mu: vfdfield_t | None = None, - pmc: vfdfield_t | None = None, - ) -> sparse.spmatrix: + mu: vfdfield | None = None, + pmc: vfdfield | None = None, + ) -> sparse.sparray: """ Utility operator for converting the E field into the H field. For use with `e_full()` -- assumes that there is no magnetic current M. @@ -274,10 +274,10 @@ def e2h( op = curl_forward(dxes[0]) / (-1j * omega) if mu is not None: - op = sparse.diags(1 / mu) @ op + op = sparse.diags_array(1 / mu) @ op if pmc is not None: - op = sparse.diags(numpy.where(pmc, 0, 1)) @ op + op = sparse.diags_array(numpy.where(pmc, 0, 1)) @ op return op @@ -285,8 +285,8 @@ def e2h( def m2j( omega: complex, dxes: dx_lists_t, - mu: vfdfield_t | None = None, - ) -> sparse.spmatrix: + mu: vfdfield | None = None, + ) -> sparse.sparray: """ Operator for converting a magnetic current M into an electric current J. For use with eg. `e_full()`. @@ -302,12 +302,12 @@ def m2j( op = curl_back(dxes[1]) / (1j * omega) if mu is not None: - op = op @ sparse.diags(1 / mu) + op = op @ sparse.diags_array(1 / mu) return op -def poynting_e_cross(e: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: +def poynting_e_cross(e: vcfdfield, dxes: dx_lists_t) -> sparse.sparray: """ Operator for computing the Poynting vector, containing the (E x) portion of the Poynting vector. @@ -330,13 +330,13 @@ def poynting_e_cross(e: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: block_diags = [[ None, fx @ -Ez, fx @ Ey], [ fy @ Ez, None, fy @ -Ex], [ fz @ -Ey, fz @ Ex, None]] - block_matrix = sparse.bmat([[sparse.diags(x) if x is not None else None for x in row] - for row in block_diags]) - P = block_matrix @ sparse.diags(numpy.concatenate(dxbg)) + block_matrix = sparse.block_array([[sparse.diags_array(x) if x is not None else None for x in row] + for row in block_diags]) + P = block_matrix @ sparse.diags_array(numpy.concatenate(dxbg)) return P -def poynting_h_cross(h: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: +def poynting_h_cross(h: vcfdfield, dxes: dx_lists_t) -> sparse.sparray: """ Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector. @@ -353,23 +353,23 @@ def poynting_h_cross(h: vcfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix: dxag = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[0], indexing='ij')] dxbg = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[1], indexing='ij')] - Hx, Hy, Hz = (sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg, strict=True)) + Hx, Hy, Hz = (sparse.diags_array(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg, strict=True)) - P = (sparse.bmat( + P = (sparse.block_array( [[ None, -Hz @ fx, Hy @ fx], [ Hz @ fy, None, -Hx @ fy], [-Hy @ fz, Hx @ fz, None]]) - @ sparse.diags(numpy.concatenate(dxag))) + @ sparse.diags_array(numpy.concatenate(dxag))) return P def e_tfsf_source( - TF_region: vfdfield_t, + TF_region: vfdfield, omega: complex, dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, - ) -> sparse.spmatrix: + epsilon: vfdfield, + mu: vfdfield | None = None, + ) -> sparse.sparray: """ Operator that turns a desired E-field distribution into a total-field/scattered-field (TFSF) source. @@ -390,18 +390,18 @@ def e_tfsf_source( """ # TODO documentation A = e_full(omega, dxes, epsilon, mu) - Q = sparse.diags(TF_region) + Q = sparse.diags_array(TF_region) return (A @ Q - Q @ A) / (-1j * omega) def e_boundary_source( - mask: vfdfield_t, + mask: vfdfield, omega: complex, dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + epsilon: vfdfield, + mu: vfdfield | None = None, periodic_mask_edges: bool = False, - ) -> sparse.spmatrix: + ) -> sparse.sparray: """ 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 @@ -424,10 +424,10 @@ def e_boundary_source( shape = [len(dxe) for dxe in dxes[0]] jmask = numpy.zeros_like(mask, dtype=bool) - def shift_rot(axis: int, polarity: int) -> sparse.spmatrix: + def shift_rot(axis: int, polarity: int) -> sparse.sparray: return shift_circ(axis=axis, shape=shape, shift_distance=polarity) - def shift_mir(axis: int, polarity: int) -> sparse.spmatrix: + def shift_mir(axis: int, polarity: int) -> sparse.sparray: return shift_with_mirror(axis=axis, shape=shape, shift_distance=polarity) shift = shift_rot if periodic_mask_edges else shift_mir @@ -436,7 +436,7 @@ def e_boundary_source( if shape[axis] == 1: continue for polarity in (-1, +1): - r = shift(axis, polarity) - sparse.eye(numpy.prod(shape)) # shifted minus original + r = shift(axis, polarity) - sparse.eye_array(numpy.prod(shape)) # shifted minus original r3 = sparse.block_diag((r, r, r)) jmask = numpy.logical_or(jmask, numpy.abs(r3 @ mask)) @@ -447,5 +447,5 @@ def e_boundary_source( # (numpy.roll(mask, -1, axis=2) != mask) | # (numpy.roll(mask, +1, axis=2) != mask)) - return sparse.diags(jmask.astype(int)) @ full + return sparse.diags_array(jmask.astype(int)) @ full diff --git a/meanas/fdfd/solvers.py b/meanas/fdfd/solvers.py index 81d1d09..c0aed44 100644 --- a/meanas/fdfd/solvers.py +++ b/meanas/fdfd/solvers.py @@ -11,7 +11,7 @@ from numpy.typing import ArrayLike, NDArray from numpy.linalg import norm import scipy.sparse.linalg -from ..fdmath import dx_lists_t, vfdfield_t, vcfdfield_t +from ..fdmath import dx_lists_t, vfdfield, vcfdfield, vcfdfield_t from . import operators @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) def _scipy_qmr( - A: scipy.sparse.csr_matrix, + A: scipy.sparse.csr_array, b: ArrayLike, **kwargs: Any, ) -> NDArray[numpy.float64]: @@ -66,16 +66,16 @@ def _scipy_qmr( def generic( omega: complex, dxes: dx_lists_t, - J: vcfdfield_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + J: vcfdfield, + epsilon: vfdfield, + mu: vfdfield | None = None, *, - pec: vfdfield_t | None = None, - pmc: vfdfield_t | None = None, + pec: vfdfield | None = None, + pmc: vfdfield | None = None, adjoint: bool = False, matrix_solver: Callable[..., ArrayLike] = _scipy_qmr, matrix_solver_opts: dict[str, Any] | None = None, - E_guess: vcfdfield_t | None = None, + E_guess: vcfdfield | None = None, ) -> vcfdfield_t: """ Conjugate gradient FDFD solver using CSR sparse matrices. @@ -95,7 +95,7 @@ def generic( (at H-field locations; non-zero value indicates PMC is present) 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`; + where `A`: `scipy.sparse.csr_array`; `b`: `ArrayLike`; `x`: `ArrayLike`; Default is a wrapped version of `scipy.sparse.linalg.qmr()` @@ -138,4 +138,4 @@ def generic( else: x0 = Pr @ x - return x0 + return vcfdfield_t(x0) diff --git a/meanas/fdfd/waveguide_2d.py b/meanas/fdfd/waveguide_2d.py index 5fda683..e8f766b 100644 --- a/meanas/fdfd/waveguide_2d.py +++ b/meanas/fdfd/waveguide_2d.py @@ -180,12 +180,12 @@ if the result is introduced into a space with a discretized z-axis. from typing import Any from collections.abc import Sequence import numpy -from numpy.typing import NDArray, ArrayLike +from numpy.typing import NDArray from numpy.linalg import norm from scipy import sparse from ..fdmath.operators import deriv_forward, deriv_back, cross -from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t +from ..fdmath import vec, unvec, dx_lists2_t, vcfdfield2_t, vcfdslice_t, vcfdfield2, vfdslice, vcfdslice from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration @@ -194,10 +194,10 @@ __author__ = 'Jan Petykiewicz' def operator_e( omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, + ) -> sparse.sparray: r""" Waveguide operator of the form @@ -246,12 +246,12 @@ def operator_e( Dbx, Dby = deriv_back(dxes[1]) eps_parts = numpy.split(epsilon, 3) - eps_xy = sparse.diags(numpy.hstack((eps_parts[0], eps_parts[1]))) - eps_z_inv = sparse.diags(1 / eps_parts[2]) + eps_xy = sparse.diags_array(numpy.hstack((eps_parts[0], eps_parts[1]))) + eps_z_inv = sparse.diags_array(1 / eps_parts[2]) mu_parts = numpy.split(mu, 3) - mu_yx = sparse.diags(numpy.hstack((mu_parts[1], mu_parts[0]))) - mu_z_inv = sparse.diags(1 / mu_parts[2]) + mu_yx = sparse.diags_array(numpy.hstack((mu_parts[1], mu_parts[0]))) + mu_z_inv = sparse.diags_array(1 / mu_parts[2]) op = ( omega * omega * mu_yx @ eps_xy @@ -263,10 +263,10 @@ def operator_e( def operator_h( omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, + ) -> sparse.sparray: r""" Waveguide operator of the form @@ -315,12 +315,12 @@ def operator_h( Dbx, Dby = deriv_back(dxes[1]) eps_parts = numpy.split(epsilon, 3) - eps_yx = sparse.diags(numpy.hstack((eps_parts[1], eps_parts[0]))) - eps_z_inv = sparse.diags(1 / eps_parts[2]) + eps_yx = sparse.diags_array(numpy.hstack((eps_parts[1], eps_parts[0]))) + eps_z_inv = sparse.diags_array(1 / eps_parts[2]) mu_parts = numpy.split(mu, 3) - mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1]))) - mu_z_inv = sparse.diags(1 / mu_parts[2]) + mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1]))) + mu_z_inv = sparse.diags_array(1 / mu_parts[2]) op = ( omega * omega * eps_yx @ mu_xy @@ -331,14 +331,14 @@ def operator_h( def normalized_fields_e( - e_xy: ArrayLike, + e_xy: vcfdfield2, wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, prop_phase: float = 0, - ) -> tuple[vcfdfield_t, vcfdfield_t]: + ) -> tuple[vcfdslice_t, vcfdslice_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. @@ -366,14 +366,14 @@ def normalized_fields_e( def normalized_fields_h( - h_xy: ArrayLike, + h_xy: vcfdfield2, wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, prop_phase: float = 0, - ) -> tuple[vcfdfield_t, vcfdfield_t]: + ) -> tuple[vcfdslice_t, vcfdslice_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. @@ -401,21 +401,19 @@ def normalized_fields_h( def _normalized_fields( - e: vcfdfield_t, - h: vcfdfield_t, + e: vcfdslice, + h: vcfdslice, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, prop_phase: float = 0, - ) -> tuple[vcfdfield_t, vcfdfield_t]: + ) -> tuple[vcfdslice_t, vcfdslice_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)] # Find time-averaged Sz and normalize to it # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting - phase = numpy.exp(-1j * -prop_phase / 2) Sz_tavg = inner_product(e, h, dxes=dxes, prop_phase=prop_phase, conj_h=True).real assert Sz_tavg > 0, f'Found a mode propagating in the wrong direction! {Sz_tavg=}' @@ -436,16 +434,16 @@ def _normalized_fields( e *= norm_factor h *= norm_factor - return e, h + return vcfdslice_t(e), vcfdslice_t(h) def exy2h( wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None + ) -> sparse.sparray: """ 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 @@ -468,10 +466,10 @@ def exy2h( def hxy2e( wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None + ) -> sparse.sparray: """ 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 @@ -493,9 +491,9 @@ def hxy2e( def hxy2h( wavenumber: complex, - dxes: dx_lists_t, - mu: vfdfield_t | None = None - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + mu: vfdslice | None = None + ) -> sparse.sparray: """ 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 @@ -514,22 +512,22 @@ def hxy2h( if mu is not None: mu_parts = numpy.split(mu, 3) - mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1]))) - mu_z_inv = sparse.diags(1 / mu_parts[2]) + mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1]))) + mu_z_inv = sparse.diags_array(1 / mu_parts[2]) hxy2hz = mu_z_inv @ hxy2hz @ mu_xy n_pts = dxes[1][0].size * dxes[1][1].size - op = sparse.vstack((sparse.eye(2 * n_pts), + op = sparse.vstack((sparse.eye_array(2 * n_pts), hxy2hz)) return op def exy2e( wavenumber: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + epsilon: vfdslice, + ) -> sparse.sparray: r""" 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 @@ -575,13 +573,13 @@ def exy2e( if epsilon is not None: epsilon_parts = numpy.split(epsilon, 3) - epsilon_xy = sparse.diags(numpy.hstack((epsilon_parts[0], epsilon_parts[1]))) - epsilon_z_inv = sparse.diags(1 / epsilon_parts[2]) + epsilon_xy = sparse.diags_array(numpy.hstack((epsilon_parts[0], epsilon_parts[1]))) + epsilon_z_inv = sparse.diags_array(1 / epsilon_parts[2]) exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy n_pts = dxes[0][0].size * dxes[0][1].size - op = sparse.vstack((sparse.eye(2 * n_pts), + op = sparse.vstack((sparse.eye_array(2 * n_pts), exy2ez)) return op @@ -589,12 +587,12 @@ def exy2e( def e2h( wavenumber: complex, omega: complex, - dxes: dx_lists_t, - mu: vfdfield_t | None = None - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + mu: vfdslice | None = None + ) -> sparse.sparray: """ Returns an operator which, when applied to a vectorized E eigenfield, produces - the vectorized H eigenfield. + the vectorized H eigenfield slice. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -607,19 +605,19 @@ def e2h( """ op = curl_e(wavenumber, dxes) / (-1j * omega) if mu is not None: - op = sparse.diags(1 / mu) @ op + op = sparse.diags_array(1 / mu) @ op return op def h2e( wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t - ) -> sparse.spmatrix: + dxes: dx_lists2_t, + epsilon: vfdslice, + ) -> sparse.sparray: """ Returns an operator which, when applied to a vectorized H eigenfield, produces - the vectorized E eigenfield. + the vectorized E eigenfield slice. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -630,13 +628,13 @@ def h2e( Returns: Sparse matrix representation of the operator. """ - op = sparse.diags(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes) + op = sparse.diags_array(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes) return op -def curl_e(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: +def curl_e(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray: """ - Discretized curl operator for use with the waveguide E field. + Discretized curl operator for use with the waveguide E field slice. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -645,18 +643,18 @@ def curl_e(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: Returns: Sparse matrix representation of the operator. """ - n = 1 - for d in dxes[0]: - n *= len(d) + nn = 1 + for dd in dxes[0]: + nn *= len(dd) - Bz = -1j * wavenumber * sparse.eye(n) + Bz = -1j * wavenumber * sparse.eye_array(nn) Dfx, Dfy = deriv_forward(dxes[0]) return cross([Dfx, Dfy, Bz]) -def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: +def curl_h(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray: """ - Discretized curl operator for use with the waveguide H field. + Discretized curl operator for use with the waveguide H field slice. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -665,22 +663,22 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: Returns: Sparse matrix representation of the operator. """ - n = 1 - for d in dxes[1]: - n *= len(d) + nn = 1 + for dd in dxes[1]: + nn *= len(dd) - Bz = -1j * wavenumber * sparse.eye(n) + Bz = -1j * wavenumber * sparse.eye_array(nn) Dbx, Dby = deriv_back(dxes[1]) return cross([Dbx, Dby, Bz]) def h_err( - h: vcfdfield_t, + h: vcfdslice, wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None ) -> float: """ Calculates the relative error in the H field @@ -699,7 +697,7 @@ def h_err( ce = curl_e(wavenumber, dxes) ch = curl_h(wavenumber, dxes) - eps_inv = sparse.diags(1 / epsilon) + eps_inv = sparse.diags_array(1 / epsilon) if mu is None: op = ce @ eps_inv @ ch @ h - omega ** 2 * h @@ -710,12 +708,12 @@ def h_err( def e_err( - e: vcfdfield_t, + e: vcfdslice, wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, ) -> float: """ Calculates the relative error in the E field @@ -737,21 +735,21 @@ def e_err( if mu is None: op = ch @ ce @ e - omega ** 2 * (epsilon * e) else: - mu_inv = sparse.diags(1 / mu) + mu_inv = sparse.diags_array(1 / mu) op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e) return float(norm(op) / norm(e)) def sensitivity( - e_norm: vcfdfield_t, - h_norm: vcfdfield_t, + e_norm: vcfdslice, + h_norm: vcfdslice, wavenumber: complex, omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, - ) -> vcfdfield_t: + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, + ) -> vcfdslice_t: r""" Given a waveguide structure (`dxes`, `epsilon`, `mu`) and mode fields (`e_norm`, `h_norm`, `wavenumber`, `omega`), calculates the sensitivity of the wavenumber @@ -825,11 +823,11 @@ def sensitivity( Dbx, Dby = deriv_back(dxes[1]) eps_x, eps_y, eps_z = numpy.split(epsilon, 3) - eps_xy = sparse.diags(numpy.hstack((eps_x, eps_y))) - eps_z_inv = sparse.diags(1 / eps_z) + eps_xy = sparse.diags_array(numpy.hstack((eps_x, eps_y))) + eps_z_inv = sparse.diags_array(1 / eps_z) mu_x, mu_y, _mu_z = numpy.split(mu, 3) - mu_yx = sparse.diags(numpy.hstack((mu_y, mu_x))) + mu_yx = sparse.diags_array(numpy.hstack((mu_y, mu_x))) da_exxhyy = vec(dxes[1][0][:, None] * dxes[0][1][None, :]) da_eyyhxx = vec(dxes[1][1][None, :] * dxes[0][0][:, None]) @@ -843,15 +841,15 @@ def sensitivity( norm = hv_yx_conj @ ev_xy sens_tot = numpy.concatenate([sens_xy1 + sens_xy2, sens_z]) / (2 * wavenumber * norm) - return sens_tot + return vcfdslice_t(sens_tot) def solve_modes( mode_numbers: Sequence[int], omega: complex, - dxes: dx_lists_t, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + dxes: dx_lists2_t, + epsilon: vfdslice, + mu: vfdslice | None = None, mode_margin: int = 2, ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: """ @@ -907,7 +905,7 @@ def solve_mode( mode_number: int, *args: Any, **kwargs: Any, - ) -> tuple[vcfdfield_t, complex]: + ) -> tuple[vcfdfield2_t, complex]: """ Wrapper around `solve_modes()` that solves for a single mode. @@ -921,13 +919,13 @@ def solve_mode( """ kwargs['mode_numbers'] = [mode_number] e_xys, wavenumbers = solve_modes(*args, **kwargs) - return e_xys[0], wavenumbers[0] + return vcfdfield2_t(e_xys[0]), wavenumbers[0] def inner_product( # TODO documentation - e1: vcfdfield_t, - h2: vcfdfield_t, - dxes: dx_lists_t, + e1: vcfdfield2, + h2: vcfdfield2, + dxes: dx_lists2_t, prop_phase: float = 0, conj_h: bool = False, trapezoid: bool = False, diff --git a/meanas/fdfd/waveguide_3d.py b/meanas/fdfd/waveguide_3d.py index 6e2a2db..da50533 100644 --- a/meanas/fdfd/waveguide_3d.py +++ b/meanas/fdfd/waveguide_3d.py @@ -4,13 +4,13 @@ Tools for working with waveguide modes in 3D domains. This module relies heavily on `waveguide_2d` and mostly just transforms its parameters into 2D equivalents and expands the results back into 3D. """ -from typing import Any +from typing import Any, cast from collections.abc import Sequence import numpy from numpy.typing import NDArray from numpy import complexfloating -from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, cfdfield_t +from ..fdmath import vec, unvec, dx_lists_t, cfdfield_t, fdfield, cfdfield from . import operators, waveguide_2d @@ -21,8 +21,8 @@ def solve_mode( axis: int, polarity: int, slices: Sequence[slice], - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = None, ) -> dict[str, complex | NDArray[complexfloating]]: """ Given a 3D grid, selects a slice from the grid and attempts to @@ -95,9 +95,10 @@ def solve_mode( # Expand E, H to full epsilon space we were given E = numpy.zeros_like(epsilon, dtype=complex) H = numpy.zeros_like(epsilon, dtype=complex) - for a, o in enumerate(reverse_order): - E[(a, *slices)] = e[o][:, :, None].transpose(reverse_order) - H[(a, *slices)] = h[o][:, :, None].transpose(reverse_order) + for aa, oo in enumerate(reverse_order): + iii = cast('tuple[slice | int]', (aa, *slices)) + E[iii] = e[oo][:, :, None].transpose(reverse_order) + H[iii] = h[oo][:, :, None].transpose(reverse_order) results = { 'wavenumber': wavenumber, @@ -109,15 +110,15 @@ def solve_mode( def compute_source( - E: cfdfield_t, + E: cfdfield, wavenumber: complex, omega: complex, dxes: dx_lists_t, axis: int, polarity: int, slices: Sequence[slice], - epsilon: fdfield_t, - mu: fdfield_t | None = None, + epsilon: fdfield, + mu: fdfield | None = None, ) -> cfdfield_t: """ Given an eigenmode obtained by `solve_mode`, returns the current source distribution @@ -151,11 +152,11 @@ def compute_source( masked_e2j = operators.e_boundary_source(mask=vec(mask), omega=omega, dxes=dxes, epsilon=vec(epsilon), mu=vec(mu)) J = unvec(masked_e2j @ vec(E_expanded), E.shape[1:]) - return J + return cfdfield_t(J) def compute_overlap_e( - E: cfdfield_t, + E: cfdfield, wavenumber: complex, dxes: dx_lists_t, axis: int, @@ -195,12 +196,12 @@ def compute_overlap_e( Etgt = numpy.zeros_like(Ee) Etgt[slices2] = Ee[slices2] - Etgt /= (Etgt.conj() * Etgt).sum() - return Etgt + Etgt /= (Etgt.conj() * Etgt).sum() # type: ignore + return cfdfield_t(Etgt) def expand_e( - E: cfdfield_t, + E: cfdfield, wavenumber: complex, dxes: dx_lists_t, axis: int, @@ -245,4 +246,4 @@ def expand_e( slices_in = (slice(None), *slices) E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in] - return E_expanded + return cfdfield_t(E_expanded) diff --git a/meanas/fdfd/waveguide_cyl.py b/meanas/fdfd/waveguide_cyl.py index b65e038..597e1cb 100644 --- a/meanas/fdfd/waveguide_cyl.py +++ b/meanas/fdfd/waveguide_cyl.py @@ -68,7 +68,7 @@ import numpy from numpy.typing import NDArray, ArrayLike from scipy import sparse -from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t +from ..fdmath import vec, unvec, dx_lists2_t, vcfdslice_t, vcfdfield2_t, vfdslice, vcfdslice, vcfdfield2 from ..fdmath.operators import deriv_forward, deriv_back from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration from . import waveguide_2d @@ -78,10 +78,10 @@ logger = logging.getLogger(__name__) def cylindrical_operator( omega: float, - dxes: dx_lists_t, - epsilon: vfdfield_t, + dxes: dx_lists2_t, + epsilon: vfdslice, rmin: float, - ) -> sparse.spmatrix: + ) -> sparse.sparray: r""" Cylindrical coordinate waveguide operator of the form @@ -143,11 +143,11 @@ def cylindrical_operator( def solve_modes( mode_numbers: Sequence[int], omega: float, - dxes: dx_lists_t, - epsilon: vfdfield_t, + dxes: dx_lists2_t, + epsilon: vfdslice, rmin: float, mode_margin: int = 2, - ) -> tuple[vcfdfield_t, NDArray[numpy.complex128]]: + ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: """ Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode of the bent waveguide with the specified mode number. @@ -191,7 +191,7 @@ def solve_modes( # Wavenumbers assume the mode is at rmin, which is unlikely # Instead, return the wavenumber in inverse radians - angular_wavenumbers = wavenumbers * cast(complex, rmin) + angular_wavenumbers = wavenumbers * cast('complex', rmin) order = angular_wavenumbers.argsort()[::-1] e_xys = e_xys[order] @@ -204,7 +204,7 @@ def solve_mode( mode_number: int, *args: Any, **kwargs: Any, - ) -> tuple[vcfdfield_t, complex]: + ) -> tuple[vcfdslice, complex]: """ Wrapper around `solve_modes()` that solves for a single mode. @@ -222,10 +222,10 @@ def solve_mode( def linear_wavenumbers( - e_xys: vcfdfield_t, + e_xys: list[vcfdfield2_t], angular_wavenumbers: ArrayLike, - epsilon: vfdfield_t, - dxes: dx_lists_t, + epsilon: vfdslice, + dxes: dx_lists2_t, rmin: float, ) -> NDArray[numpy.complex128]: """ @@ -247,7 +247,6 @@ def linear_wavenumbers( angular_wavenumbers = numpy.asarray(angular_wavenumbers) mode_radii = numpy.empty_like(angular_wavenumbers, dtype=float) - wavenumbers = numpy.empty_like(angular_wavenumbers) shape2d = (len(dxes[0][0]), len(dxes[0][1])) epsilon2d = unvec(epsilon, shape2d)[:2] grid_radii = rmin + numpy.cumsum(dxes[0][0]) @@ -265,11 +264,11 @@ def linear_wavenumbers( def exy2h( angular_wavenumber: complex, omega: float, - dxes: dx_lists_t, + dxes: dx_lists2_t, rmin: float, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None - ) -> sparse.spmatrix: + epsilon: vfdslice, + mu: vfdslice | None = None + ) -> sparse.sparray: """ 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 @@ -294,10 +293,10 @@ def exy2h( def exy2e( angular_wavenumber: complex, omega: float, - dxes: dx_lists_t, + dxes: dx_lists2_t, rmin: float, - epsilon: vfdfield_t, - ) -> sparse.spmatrix: + epsilon: vfdslice, + ) -> sparse.sparray: """ 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 @@ -323,7 +322,7 @@ def exy2e( Ta, Tb = dxes2T(dxes=dxes, rmin=rmin) Tai = sparse.diags_array(1 / Ta.diagonal()) - Tbi = sparse.diags_array(1 / Tb.diagonal()) + #Tbi = sparse.diags_array(1 / Tb.diagonal()) epsilon_parts = numpy.split(epsilon, 3) epsilon_x, epsilon_y = (sparse.diags_array(epsi) for epsi in epsilon_parts[:2]) @@ -331,8 +330,6 @@ def exy2e( n_pts = dxes[0][0].size * dxes[0][1].size zeros = sparse.coo_array((n_pts, n_pts)) - keep_x = sparse.block_array([[sparse.eye_array(n_pts), None], [None, zeros]]) - keep_y = sparse.block_array([[zeros, None], [None, sparse.eye_array(n_pts)]]) mu_z = numpy.ones(n_pts) mu_z_inv = sparse.diags_array(1 / mu_z) @@ -352,10 +349,10 @@ def exy2e( def e2h( angular_wavenumber: complex, omega: float, - dxes: dx_lists_t, + dxes: dx_lists2_t, rmin: float, - mu: vfdfield_t | None = None - ) -> sparse.spmatrix: + mu: vfdslice | None = None + ) -> sparse.sparray: r""" Returns an operator which, when applied to a vectorized E eigenfield, produces the vectorized H eigenfield. @@ -396,7 +393,7 @@ def e2h( def dxes2T( - dxes: dx_lists_t, + dxes: dx_lists2_t, rmin: float, ) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64]]: r""" @@ -421,15 +418,15 @@ def dxes2T( def normalized_fields_e( - e_xy: ArrayLike, + e_xy: vcfdfield2, angular_wavenumber: complex, omega: float, - dxes: dx_lists_t, + dxes: dx_lists2_t, rmin: float, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + epsilon: vfdslice, + mu: vfdslice | None = None, prop_phase: float = 0, - ) -> tuple[vcfdfield_t, vcfdfield_t]: + ) -> tuple[vcfdslice_t, vcfdslice_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. @@ -459,23 +456,21 @@ def normalized_fields_e( def _normalized_fields( - e: vcfdfield_t, - h: vcfdfield_t, + e: vcfdslice, + h: vcfdslice, omega: complex, - dxes: dx_lists_t, - rmin: float, - epsilon: vfdfield_t, - mu: vfdfield_t | None = None, + dxes: dx_lists2_t, + rmin: float, # Currently unused, but may want to use cylindrical poynting + epsilon: vfdslice, + mu: vfdslice | None = None, prop_phase: float = 0, - ) -> tuple[vcfdfield_t, vcfdfield_t]: + ) -> tuple[vcfdslice_t, vcfdslice_t]: h *= -1 # TODO documentation for normalized_fields 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)] # Find time-averaged Sz and normalize to it # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting - phase = numpy.exp(-1j * -prop_phase / 2) Sz_tavg = waveguide_2d.inner_product(e, h, dxes=dxes, prop_phase=prop_phase, conj_h=True).real # Note, using linear poynting vector assert Sz_tavg > 0, f'Found a mode propagating in the wrong direction! {Sz_tavg=}' @@ -495,4 +490,4 @@ def _normalized_fields( e *= norm_factor h *= norm_factor - return e, h + return vcfdslice_t(e), vcfdslice_t(h) diff --git a/meanas/fdmath/__init__.py b/meanas/fdmath/__init__.py index b1d8354..857cf18 100644 --- a/meanas/fdmath/__init__.py +++ b/meanas/fdmath/__init__.py @@ -742,12 +742,34 @@ normalized results are needed. """ from .types import ( - fdfield_t as fdfield_t, - vfdfield_t as vfdfield_t, - cfdfield_t as cfdfield_t, - vcfdfield_t as vcfdfield_t, - dx_lists_t as dx_lists_t, - dx_lists_mut as dx_lists_mut, + fdfield_t as fdfield_t, + vfdfield_t as vfdfield_t, + cfdfield_t as cfdfield_t, + vcfdfield_t as vcfdfield_t, + fdfield2_t as fdfield2_t, + vfdfield2_t as vfdfield2_t, + cfdfield2_t as cfdfield2_t, + vcfdfield2_t as vcfdfield2_t, + fdfield as fdfield, + vfdfield as vfdfield, + cfdfield as cfdfield, + vcfdfield as vcfdfield, + fdfield2 as fdfield2, + vfdfield2 as vfdfield2, + cfdfield2 as cfdfield2, + vcfdfield2 as vcfdfield2, + fdslice_t as fdslice_t, + vfdslice_t as vfdslice_t, + cfdslice_t as cfdslice_t, + vcfdslice_t as vcfdslice_t, + fdslice as fdslice, + vfdslice as vfdslice, + cfdslice as cfdslice, + vcfdslice as vcfdslice, + dx_lists_t as dx_lists_t, + dx_lists2_t as dx_lists2_t, + dx_lists_mut as dx_lists_mut, + dx_lists2_mut as dx_lists2_mut, fdfield_updater_t as fdfield_updater_t, cfdfield_updater_t as cfdfield_updater_t, ) diff --git a/meanas/fdmath/operators.py b/meanas/fdmath/operators.py index 946eb88..19ccb80 100644 --- a/meanas/fdmath/operators.py +++ b/meanas/fdmath/operators.py @@ -16,7 +16,7 @@ def shift_circ( axis: int, shape: Sequence[int], shift_distance: int = 1, - ) -> sparse.spmatrix: + ) -> sparse.sparray: """ Utility operator for performing a circular shift along a specified axis by a specified number of elements. @@ -44,7 +44,7 @@ def shift_circ( vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C'))) - d = sparse.csr_matrix(vij, shape=(n, n)) + d = sparse.csr_array(vij, shape=(n, n)) if shift_distance < 0: d = d.T @@ -56,7 +56,7 @@ def shift_with_mirror( axis: int, shape: Sequence[int], shift_distance: int = 1, - ) -> sparse.spmatrix: + ) -> sparse.sparray: """ Utility operator for performing an n-element shift along a specified axis, with mirror boundary conditions applied to the cells beyond the receding edge. @@ -92,13 +92,13 @@ def shift_with_mirror( vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C'))) - d = sparse.csr_matrix(vij, shape=(n, n)) + d = sparse.csr_array(vij, shape=(n, n)) return d def deriv_forward( dx_e: Sequence[NDArray[floating | complexfloating]], - ) -> list[sparse.spmatrix]: + ) -> list[sparse.sparray]: """ Utility operators for taking discretized derivatives (forward variant). @@ -114,10 +114,10 @@ def deriv_forward( dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij') - def deriv(axis: int) -> sparse.spmatrix: - return shift_circ(axis, shape, 1) - sparse.eye(n) + def deriv(axis: int) -> sparse.sparray: + return shift_circ(axis, shape, 1) - sparse.eye_array(n) - Ds = [sparse.diags(+1 / dx.ravel(order='C')) @ deriv(a) + Ds = [sparse.diags_array(+1 / dx.ravel(order='C')) @ deriv(a) for a, dx in enumerate(dx_e_expanded)] return Ds @@ -125,7 +125,7 @@ def deriv_forward( def deriv_back( dx_h: Sequence[NDArray[floating | complexfloating]], - ) -> list[sparse.spmatrix]: + ) -> list[sparse.sparray]: """ Utility operators for taking discretized derivatives (backward variant). @@ -141,18 +141,18 @@ def deriv_back( dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij') - def deriv(axis: int) -> sparse.spmatrix: - return shift_circ(axis, shape, -1) - sparse.eye(n) + def deriv(axis: int) -> sparse.sparray: + return shift_circ(axis, shape, -1) - sparse.eye_array(n) - Ds = [sparse.diags(-1 / dx.ravel(order='C')) @ deriv(a) + Ds = [sparse.diags_array(-1 / dx.ravel(order='C')) @ deriv(a) for a, dx in enumerate(dx_h_expanded)] return Ds def cross( - B: Sequence[sparse.spmatrix], - ) -> sparse.spmatrix: + B: Sequence[sparse.sparray], + ) -> sparse.sparray: """ Cross product operator @@ -164,13 +164,14 @@ def cross( Sparse matrix corresponding to (B x), where x is the cross product. """ n = B[0].shape[0] - zero = sparse.csr_matrix((n, n)) - return sparse.bmat([[zero, -B[2], B[1]], - [B[2], zero, -B[0]], - [-B[1], B[0], zero]]) + zero = sparse.csr_array((n, n)) + return sparse.block_array([ + [zero, -B[2], B[1]], + [B[2], zero, -B[0]], + [-B[1], B[0], zero]]) -def vec_cross(b: vfdfield_t) -> sparse.spmatrix: +def vec_cross(b: vfdfield_t) -> sparse.sparray: """ Vector cross product operator @@ -182,11 +183,11 @@ def vec_cross(b: vfdfield_t) -> sparse.spmatrix: Sparse matrix corresponding to (b x), where x is the cross product. """ - B = [sparse.diags(c) for c in numpy.split(b, 3)] + B = [sparse.diags_array(c) for c in numpy.split(b, 3)] return cross(B) -def avg_forward(axis: int, shape: Sequence[int]) -> sparse.spmatrix: +def avg_forward(axis: int, shape: Sequence[int]) -> sparse.sparray: """ Forward average operator `(x4 = (x4 + x5) / 2)` @@ -201,10 +202,10 @@ def avg_forward(axis: int, shape: Sequence[int]) -> sparse.spmatrix: raise Exception(f'Invalid shape: {shape}') n = numpy.prod(shape) - return 0.5 * (sparse.eye(n) + shift_circ(axis, shape)) + return 0.5 * (sparse.eye_array(n) + shift_circ(axis, shape)) -def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix: +def avg_back(axis: int, shape: Sequence[int]) -> sparse.sparray: """ Backward average operator `(x4 = (x4 + x3) / 2)` @@ -220,7 +221,7 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix: def curl_forward( dx_e: Sequence[NDArray[floating | complexfloating]], - ) -> sparse.spmatrix: + ) -> sparse.sparray: """ Curl operator for use with the E field. @@ -236,7 +237,7 @@ def curl_forward( def curl_back( dx_h: Sequence[NDArray[floating | complexfloating]], - ) -> sparse.spmatrix: + ) -> sparse.sparray: """ Curl operator for use with the H field. diff --git a/meanas/fdmath/types.py b/meanas/fdmath/types.py index d44b30a..222d18a 100644 --- a/meanas/fdmath/types.py +++ b/meanas/fdmath/types.py @@ -1,25 +1,64 @@ """ Types shared across multiple submodules """ +from typing import NewType from collections.abc import Sequence, Callable, MutableSequence from numpy.typing import NDArray from numpy import floating, complexfloating # Field types -fdfield_t = NDArray[floating] +fdfield_t = NewType('fdfield_t', NDArray[floating]) +type fdfield = fdfield_t | NDArray[floating] """Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)""" -vfdfield_t = NDArray[floating] +vfdfield_t = NewType('vfdfield_t', NDArray[floating]) +type vfdfield = vfdfield_t | NDArray[floating] """Linearized vector field (single vector of length 3*X*Y*Z)""" -cfdfield_t = NDArray[complexfloating] +cfdfield_t = NewType('cfdfield_t', NDArray[complexfloating]) +type cfdfield = cfdfield_t | NDArray[complexfloating] """Complex vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)""" -vcfdfield_t = NDArray[complexfloating] +vcfdfield_t = NewType('vcfdfield_t', NDArray[complexfloating]) +type vcfdfield = vcfdfield_t | NDArray[complexfloating] """Linearized complex vector field (single vector of length 3*X*Y*Z)""" +fdslice_t = NewType('fdslice_t', NDArray[floating]) +type fdslice = fdslice_t | NDArray[floating] +"""Vector field slice with shape (3, X, Y) (e.g. `[E_x, E_y, E_z]` at a single Z position)""" + +vfdslice_t = NewType('vfdslice_t', NDArray[floating]) +type vfdslice = vfdslice_t | NDArray[floating] +"""Linearized vector field slice (single vector of length 3*X*Y)""" + +cfdslice_t = NewType('cfdslice_t', NDArray[complexfloating]) +type cfdslice = cfdslice_t | NDArray[complexfloating] +"""Complex vector field slice with shape (3, X, Y) (e.g. `[E_x, E_y, E_z]` at a single Z position)""" + +vcfdslice_t = NewType('vcfdslice_t', NDArray[complexfloating]) +type vcfdslice = vcfdslice_t | NDArray[complexfloating] +"""Linearized complex vector field slice (single vector of length 3*X*Y)""" + + +fdfield2_t = NewType('fdfield2_t', NDArray[floating]) +type fdfield2 = fdfield2_t | NDArray[floating] +"""2D Vector field with shape (2, X, Y) (e.g. `[E_x, E_y]`)""" + +vfdfield2_t = NewType('vfdfield2_t', NDArray[floating]) +type vfdfield2 = vfdfield2_t | NDArray[floating] +"""2D Linearized vector field (single vector of length 2*X*Y)""" + +cfdfield2_t = NewType('cfdfield2_t', NDArray[complexfloating]) +type cfdfield2 = cfdfield2_t | NDArray[complexfloating] +"""2D Complex vector field with shape (2, X, Y) (e.g. `[E_x, E_y]`)""" + +vcfdfield2_t = NewType('vcfdfield2_t', NDArray[complexfloating]) +type vcfdfield2 = vcfdfield2_t | NDArray[complexfloating] +"""2D Linearized complex vector field (single vector of length 2*X*Y)""" + + dx_lists_t = Sequence[Sequence[NDArray[floating | complexfloating]]] """ 'dxes' datastructure which contains grid cell width information in the following format: @@ -31,9 +70,23 @@ dx_lists_t = Sequence[Sequence[NDArray[floating | complexfloating]]] and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc. """ +dx_lists2_t = Sequence[Sequence[NDArray[floating | complexfloating]]] +""" + 2D 'dxes' datastructure which contains grid cell width information in the following format: + + [[[dx_e[0], dx_e[1], ...], [dy_e[0], ...]], + [[dx_h[0], dx_h[1], ...], [dy_h[0], ...]]] + + 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[NDArray[floating | complexfloating]]] """Mutable version of `dx_lists_t`""" +dx_lists2_mut = MutableSequence[MutableSequence[NDArray[floating | complexfloating]]] +"""Mutable version of `dx_lists2_t`""" + fdfield_updater_t = Callable[..., 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 8f5ff39..f2c01d0 100644 --- a/meanas/fdmath/vectorization.py +++ b/meanas/fdmath/vectorization.py @@ -7,9 +7,13 @@ Vectorized versions of the field use row-major (ie., C-style) ordering. from typing import overload from collections.abc import Sequence import numpy -from numpy.typing import ArrayLike +from numpy.typing import ArrayLike, NDArray -from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t +from .types import ( + fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t, + fdslice_t, vfdslice_t, cfdslice_t, vcfdslice_t, + fdfield2_t, vfdfield2_t, cfdfield2_t, vcfdfield2_t, + ) @overload @@ -25,12 +29,28 @@ def vec(f: cfdfield_t) -> vcfdfield_t: pass @overload -def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t: +def vec(f: fdfield2_t) -> vfdfield2_t: + pass + +@overload +def vec(f: cfdfield2_t) -> vcfdfield2_t: + pass + +@overload +def vec(f: fdslice_t) -> vfdslice_t: + pass + +@overload +def vec(f: cfdslice_t) -> vcfdslice_t: + pass + +@overload +def vec(f: ArrayLike) -> NDArray: pass def vec( - f: fdfield_t | cfdfield_t | ArrayLike | None, - ) -> vfdfield_t | vcfdfield_t | None: + f: fdfield_t | cfdfield_t | fdfield2_t | cfdfield2_t | fdslice_t | cfdslice_t | ArrayLike | None, + ) -> vfdfield_t | vcfdfield_t | vfdfield2_t | vcfdfield2_t | vfdslice_t | vcfdslice_t | NDArray | None: """ Create a 1D ndarray from a vector field which spans a 1-3D region. @@ -45,7 +65,7 @@ def vec( """ if f is None: return None - return numpy.ravel(f, order='C') + return numpy.ravel(f, order='C') # type: ignore @overload @@ -60,11 +80,31 @@ def unvec(v: vfdfield_t, shape: Sequence[int], nvdim: int = 3) -> fdfield_t: def unvec(v: vcfdfield_t, shape: Sequence[int], nvdim: int = 3) -> cfdfield_t: pass +@overload +def unvec(v: vfdfield2_t, shape: Sequence[int], nvdim: int = 3) -> fdfield2_t: + pass + +@overload +def unvec(v: vcfdfield2_t, shape: Sequence[int], nvdim: int = 3) -> cfdfield2_t: + pass + +@overload +def unvec(v: vfdslice_t, shape: Sequence[int], nvdim: int = 3) -> fdslice_t: + pass + +@overload +def unvec(v: vcfdslice_t, shape: Sequence[int], nvdim: int = 3) -> cfdslice_t: + pass + +@overload +def unvec(v: ArrayLike, shape: Sequence[int], nvdim: int = 3) -> NDArray: + pass + def unvec( - v: vfdfield_t | vcfdfield_t | None, + v: vfdfield_t | vcfdfield_t | vfdfield2_t | vcfdfield2_t | vfdslice_t | vcfdslice_t | ArrayLike | None, shape: Sequence[int], nvdim: int = 3, - ) -> fdfield_t | cfdfield_t | None: + ) -> fdfield_t | cfdfield_t | fdfield2_t | cfdfield2_t | fdslice_t | cfdslice_t | NDArray | None: """ Perform the inverse of vec(): take a 1D ndarray and output an `nvdim`-component field of form e.g. `[f_x, f_y, f_z]` (`nvdim=3`) where each of `f_*` is a len(shape)-dimensional @@ -82,5 +122,5 @@ def unvec( """ if v is None: return None - return v.reshape((nvdim, *shape), order='C') + return v.reshape((nvdim, *shape), order='C') # type: ignore diff --git a/meanas/fdtd/energy.py b/meanas/fdtd/energy.py index 43ea3a1..76888ca 100644 --- a/meanas/fdtd/energy.py +++ b/meanas/fdtd/energy.py @@ -1,6 +1,6 @@ import numpy -from ..fdmath import dx_lists_t, fdfield_t +from ..fdmath import dx_lists_t, fdfield_t, fdfield from ..fdmath.functional import deriv_back @@ -8,8 +8,8 @@ from ..fdmath.functional import deriv_back def poynting( - e: fdfield_t, - h: fdfield_t, + e: fdfield, + h: fdfield, dxes: dx_lists_t | None = None, ) -> fdfield_t: r""" @@ -84,14 +84,14 @@ def poynting( s[0] = numpy.roll(ey, -1, axis=0) * hz - numpy.roll(ez, -1, axis=0) * hy s[1] = numpy.roll(ez, -1, axis=1) * hx - numpy.roll(ex, -1, axis=1) * hz s[2] = numpy.roll(ex, -1, axis=2) * hy - numpy.roll(ey, -1, axis=2) * hx - return s + return fdfield_t(s) def poynting_divergence( - s: fdfield_t | None = None, + s: fdfield | None = None, *, - e: fdfield_t | None = None, - h: fdfield_t | None = None, + e: fdfield | None = None, + h: fdfield | None = None, dxes: dx_lists_t | None = None, ) -> fdfield_t: """ @@ -122,15 +122,15 @@ def poynting_divergence( Dx, Dy, Dz = deriv_back() ds = Dx(s[0]) + Dy(s[1]) + Dz(s[2]) - return ds + return fdfield_t(ds) def energy_hstep( - e0: fdfield_t, - h1: fdfield_t, - e2: fdfield_t, - epsilon: fdfield_t | None = None, - mu: fdfield_t | None = None, + e0: fdfield, + h1: fdfield, + e2: fdfield, + epsilon: fdfield | None = None, + mu: fdfield | None = None, dxes: dx_lists_t | None = None, ) -> fdfield_t: """ @@ -150,15 +150,15 @@ def energy_hstep( Energy, at the time of the H-field `h1`. """ u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes) - return u + return fdfield_t(u) def energy_estep( - h0: fdfield_t, - e1: fdfield_t, - h2: fdfield_t, - epsilon: fdfield_t | None = None, - mu: fdfield_t | None = None, + h0: fdfield, + e1: fdfield, + h2: fdfield, + epsilon: fdfield | None = None, + mu: fdfield | None = None, dxes: dx_lists_t | None = None, ) -> fdfield_t: """ @@ -178,17 +178,17 @@ def energy_estep( Energy, at the time of the E-field `e1`. """ u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes) - return u + return fdfield_t(u) def delta_energy_h2e( dt: float, - e0: fdfield_t, - h1: fdfield_t, - e2: fdfield_t, - h3: fdfield_t, - epsilon: fdfield_t | None = None, - mu: fdfield_t | None = None, + e0: fdfield, + h1: fdfield, + e2: fdfield, + h3: fdfield, + epsilon: fdfield | None = None, + mu: fdfield | None = None, dxes: dx_lists_t | None = None, ) -> fdfield_t: """ @@ -211,17 +211,17 @@ def delta_energy_h2e( de = e2 * (e2 - e0) / dt dh = h1 * (h3 - h1) / dt du = dxmul(de, dh, epsilon, mu, dxes) - return du + return fdfield_t(du) def delta_energy_e2h( dt: float, - h0: fdfield_t, - e1: fdfield_t, - h2: fdfield_t, - e3: fdfield_t, - epsilon: fdfield_t | None = None, - mu: fdfield_t | None = None, + h0: fdfield, + e1: fdfield, + h2: fdfield, + e3: fdfield, + epsilon: fdfield | None = None, + mu: fdfield | None = None, dxes: dx_lists_t | None = None, ) -> fdfield_t: """ @@ -244,12 +244,12 @@ def delta_energy_e2h( de = e1 * (e3 - e1) / dt dh = h2 * (h2 - h0) / dt du = dxmul(de, dh, epsilon, mu, dxes) - return du + return fdfield_t(du) def delta_energy_j( - j0: fdfield_t, - e1: fdfield_t, + j0: fdfield, + e1: fdfield, dxes: dx_lists_t | None = None, ) -> fdfield_t: """ @@ -267,14 +267,14 @@ def delta_energy_j( * dxes[0][0][:, None, None] * dxes[0][1][None, :, None] * dxes[0][2][None, None, :]) - return du + return fdfield_t(du) def dxmul( - ee: fdfield_t, - hh: fdfield_t, - epsilon: fdfield_t | float | None = None, - mu: fdfield_t | float | None = None, + ee: fdfield, + hh: fdfield, + epsilon: fdfield | float | None = None, + mu: fdfield | float | None = None, dxes: dx_lists_t | None = None, ) -> fdfield_t: if epsilon is None: @@ -292,4 +292,4 @@ def dxmul( * dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[1][2][None, None, :]) - return result + return fdfield_t(result) diff --git a/meanas/fdtd/misc.py b/meanas/fdtd/misc.py index 160682d..3fb3371 100644 --- a/meanas/fdtd/misc.py +++ b/meanas/fdtd/misc.py @@ -1,5 +1,4 @@ -from typing import Callable -from collections.abc import Sequence +from collections.abc import Callable import logging import numpy @@ -95,10 +94,10 @@ def ricker_pulse( logger.warning('meanas.fdtd.misc functions are still very WIP!') # TODO omega = 2 * pi / wl freq = 1 / wl - r0 = omega / 2 + # r0 = omega / 2 from scipy.optimize import root_scalar - delay_results = root_scalar(lambda xx: (1 - omega * omega * tt * tt / 2) * numpy.exp(-omega * omega / 4 * tt * tt) - turn_on, x0=0, x1=-2 / omega) + delay_results = root_scalar(lambda tt: (1 - omega * omega * tt * tt / 2) * numpy.exp(-omega * omega / 4 * tt * tt) - turn_on, x0=0, x1=-2 / omega) delay = delay_results.root delay = numpy.ceil(delay * freq) / freq # force delay to integer number of periods to maintain phase diff --git a/meanas/fdtd/pml.py b/meanas/fdtd/pml.py index 8098da0..fc456eb 100644 --- a/meanas/fdtd/pml.py +++ b/meanas/fdtd/pml.py @@ -13,7 +13,7 @@ from copy import deepcopy import numpy from numpy.typing import NDArray, DTypeLike -from ..fdmath import fdfield_t, dx_lists_t +from ..fdmath import fdfield, fdfield_t, dx_lists_t from ..fdmath.functional import deriv_forward, deriv_back @@ -97,7 +97,7 @@ def updates_with_cpml( cpml_params: Sequence[Sequence[dict[str, Any] | None]], dt: float, dxes: dx_lists_t, - epsilon: fdfield_t, + epsilon: fdfield, *, dtype: DTypeLike = numpy.float32, ) -> tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None], diff --git a/meanas/test/test_fdfd.py b/meanas/test/test_fdfd.py index 5f2cf11..82587b4 100644 --- a/meanas/test/test_fdfd.py +++ b/meanas/test/test_fdfd.py @@ -6,7 +6,7 @@ from numpy.typing import NDArray #from numpy.testing import assert_allclose, assert_array_equal from .. import fdfd -from ..fdmath import vec, unvec +from ..fdmath import vec, unvec, vcfdfield, vfdfield, dx_lists_t from .utils import assert_close # , assert_fields_close from .conftest import FixtureRequest @@ -102,16 +102,16 @@ def j_distribution( @dataclasses.dataclass() class FDResult: shape: tuple[int, ...] - dxes: list[list[NDArray[numpy.float64]]] - epsilon: NDArray[numpy.float64] + dxes: dx_lists_t + epsilon: vfdfield omega: complex - j: NDArray[numpy.complex128] - e: NDArray[numpy.complex128] - pmc: NDArray[numpy.float64] | None - pec: NDArray[numpy.float64] | None + j: vcfdfield + e: vcfdfield + pmc: vfdfield | None + pec: vfdfield | None -@pytest.fixture() +@pytest.fixture def sim( request: FixtureRequest, shape: tuple[int, ...], @@ -141,11 +141,11 @@ 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, 'rtol': 1e-11}, + J = j_vec, + omega = omega, + dxes = dxes, + epsilon = eps_vec, + matrix_solver_opts = dict(atol=1e-15, rtol=1e-11), ) e = unvec(e_vec, shape[1:]) diff --git a/meanas/test/test_fdfd_pml.py b/meanas/test/test_fdfd_pml.py index 832053d..540a3a0 100644 --- a/meanas/test/test_fdfd_pml.py +++ b/meanas/test/test_fdfd_pml.py @@ -5,7 +5,7 @@ from numpy.typing import NDArray from numpy.testing import assert_allclose from .. import fdfd -from ..fdmath import vec, unvec, dx_lists_mut +from ..fdmath import vec, unvec, dx_lists_mut, vfdfield, cfdfield_t #from .utils import assert_close, assert_fields_close from .test_fdfd import FDResult from .conftest import FixtureRequest @@ -70,15 +70,15 @@ def src_polarity(request: FixtureRequest) -> int: return request.param -@pytest.fixture() +@pytest.fixture def j_distribution( request: FixtureRequest, shape: tuple[int, ...], - epsilon: NDArray[numpy.float64], + epsilon: vfdfield, dxes: dx_lists_mut, omega: float, src_polarity: int, - ) -> NDArray[numpy.complex128]: + ) -> cfdfield_t: j = numpy.zeros(shape, dtype=complex) dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis @@ -109,7 +109,7 @@ def j_distribution( return j -@pytest.fixture() +@pytest.fixture def epsilon( request: FixtureRequest, shape: tuple[int, ...], @@ -144,7 +144,7 @@ def dxes( return dxes -@pytest.fixture() +@pytest.fixture def sim( request: FixtureRequest, shape: tuple[int, ...], diff --git a/meanas/test/test_fdtd.py b/meanas/test/test_fdtd.py index 25ee891..03d2b7e 100644 --- a/meanas/test/test_fdtd.py +++ b/meanas/test/test_fdtd.py @@ -189,7 +189,7 @@ def j_distribution( return j -@pytest.fixture() +@pytest.fixture def sim( request: FixtureRequest, shape: tuple[int, ...],