diff --git a/meanas/eigensolvers.py b/meanas/eigensolvers.py index 21e2ec0..e8630aa 100644 --- a/meanas/eigensolvers.py +++ b/meanas/eigensolvers.py @@ -64,10 +64,10 @@ def rayleigh_quotient_iteration( (eigenvalues, eigenvectors) """ try: - (operator - sparse.eye_array(operator.shape[0])) + (operator - sparse.eye(operator.shape[0])) - def shift(eigval: float) -> sparse.sparray: - return eigval * sparse.eye_array(operator.shape[0]) + def shift(eigval: float) -> sparse: + return eigval * sparse.eye(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_array(operator.shape[0]) + shifted_operator = operator + shift * sparse.eye(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 4eedcc4..71b2a8b 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, cfdfield, cfdfield_t +from ..fdmath import fdfield_t, cfdfield_t logger = logging.getLogger(__name__) @@ -183,8 +183,8 @@ def generate_kmn( def maxwell_operator( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield, - mu: fdfield | None = None + epsilon: fdfield_t, + mu: fdfield_t | None = None ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: """ Generate the Maxwell operator @@ -238,7 +238,7 @@ def maxwell_operator( # cross product and transform into xyz basis d_xyz = (n * hin_m - - m * hin_n) * k_mag + - m * hin_n) * k_mag # noqa: E128 # divide by epsilon temp = ifftn(d_xyz, axes=range(3)) # reuses d_xyz if using pyfftw @@ -254,7 +254,7 @@ def maxwell_operator( else: # transform from mn to xyz b_xyz = (m * b_m[:, :, :, None] - + n * b_n[:, :, :, None]) + + n * b_n[:, :, :, None]) # noqa: E128 # divide by mu temp = ifftn(b_xyz, axes=range(3)) @@ -276,7 +276,7 @@ def maxwell_operator( def hmn_2_exyz( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield, + epsilon: fdfield_t, ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: """ Generate an operator which converts a vectorized spatial-frequency-space @@ -305,11 +305,10 @@ def hmn_2_exyz( def operator(h: NDArray[numpy.complex128]) -> cfdfield_t: hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2)) d_xyz = (n * hin_m - - m * hin_n) * k_mag + - m * hin_n) * k_mag # noqa: E128 # divide by epsilon - exyz = numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0) - return cfdfield_t(exyz) + return numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0) return operator @@ -317,7 +316,7 @@ def hmn_2_exyz( def hmn_2_hxyz( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield, + epsilon: fdfield_t ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: """ Generate an operator which converts a vectorized spatial-frequency-space @@ -344,8 +343,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) - return cfdfield_t(numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])) + + n * hin_n) # noqa: E128 + return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)]) return operator @@ -353,8 +352,8 @@ def hmn_2_hxyz( def inverse_maxwell_operator_approx( k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | None = None, ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: """ Generate an approximate inverse of the Maxwell operator, @@ -404,7 +403,7 @@ def inverse_maxwell_operator_approx( else: # transform from mn to xyz h_xyz = (m * hin_m[:, :, :, None] - + n * hin_n[:, :, :, None]) + + n * hin_n[:, :, :, None]) # noqa: E128 # multiply by mu temp = ifftn(h_xyz, axes=range(3)) @@ -417,7 +416,7 @@ def inverse_maxwell_operator_approx( # cross product and transform into xyz basis e_xyz = (n * b_m - - m * b_n) / k_mag + - m * b_n) / k_mag # noqa: E128 # multiply by epsilon temp = ifftn(e_xyz, axes=range(3)) @@ -441,8 +440,8 @@ def find_k( tolerance: float, direction: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | None = None, band: int = 0, k_bounds: tuple[float, float] = (0, 0.5), k_guess: float | None = None, @@ -509,8 +508,8 @@ def eigsolve( num_modes: int, k0: ArrayLike, G_matrix: ArrayLike, - epsilon: fdfield, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | None = None, tolerance: float = 1e-7, max_iters: int = 10000, reset_iters: int = 100, @@ -650,7 +649,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) @@ -669,8 +668,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) @@ -802,12 +801,7 @@ def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]: -def inner_product( - eL: cfdfield, - hL: cfdfield, - eR: cfdfield, - hR: cfdfield, - ) -> complex: +def inner_product(eL, hL, eR, hR) -> complex: # assumes x-axis propagation assert numpy.array_equal(eR.shape, hR.shape) @@ -835,12 +829,7 @@ def inner_product( return eRxhL_x.sum() - eLxhR_x.sum() -def trq( - eI: cfdfield, - hI: cfdfield, - eO: cfdfield, - hO: cfdfield, - ) -> tuple[complex, complex]: +def trq(eI, hI, eO, hO) -> 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 f834973..35e1e90 100644 --- a/meanas/fdfd/eme.py +++ b/meanas/fdfd/eme.py @@ -1,29 +1,20 @@ -from collections.abc import Sequence import numpy -from numpy.typing import NDArray -from scipy import sparse -from ..fdmath import dx_lists2_t, vcfdfield2 +from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t from .waveguide_2d import inner_product -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]]: +def get_tr(ehL, wavenumbers_L, ehR, wavenumbers_R, dxes: dx_lists_t): 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 = ehLs[ll] + eL, hL = ehL[ll] B11[ll] = inner_product(eL, hL, dxes=dxes, conj_h=False) for rr in range(nR): - eR, hR = ehRs[rr] + eR, hR = ehR[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) @@ -41,15 +32,9 @@ def get_tr( return tt, rr -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) +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) t21i = numpy.linalg.pinv(t21) A = t12 - r21 @ t21i @ r12 B = r21 @ t21i @@ -59,16 +44,15 @@ def get_abcd( def get_s( - ehLs: Sequence[Sequence[vcfdfield2]], - wavenumbers_L: Sequence[complex], - ehRs: Sequence[Sequence[vcfdfield2]], - wavenumbers_R: Sequence[complex], + eL_xys, + wavenumbers_L, + eR_xys, + wavenumbers_R, force_nogain: bool = False, force_reciprocal: bool = False, - **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) + **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) ss = numpy.block([[r12, t12], [t21, r21]]) diff --git a/meanas/fdfd/farfield.py b/meanas/fdfd/farfield.py index 86ec0d7..4829d86 100644 --- a/meanas/fdfd/farfield.py +++ b/meanas/fdfd/farfield.py @@ -1,16 +1,14 @@ """ Functions for performing near-to-farfield transformation (and the reverse). """ -from typing import Any, cast, TYPE_CHECKING +from typing import Any, cast +from collections.abc import Sequence 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, @@ -65,7 +63,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] @@ -174,7 +172,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 440daf2..f4a250f 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, cfdfield_t, fdfield, cfdfield, cfdfield_updater_t +from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, 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, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | 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) -> cfdfield_t: + def op_1(e: cfdfield_t) -> cfdfield_t: curls = ch(ce(e)) - return cfdfield_t(curls - omega ** 2 * epsilon * e) + return curls - omega ** 2 * epsilon * e - def op_mu(e: cfdfield) -> cfdfield_t: + 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 cfdfield_t(curls - omega ** 2 * epsilon * e) + return 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, - mu: fdfield | None = None, - ) -> Callable[[cfdfield, cfdfield], tuple[cfdfield_t, cfdfield_t]]: + epsilon: fdfield_t, + mu: fdfield_t | None = None, + ) -> Callable[[cfdfield_t, cfdfield_t], tuple[cfdfield_t, cfdfield_t]]: """ Wave operator for full (both E and H) field representation. See `operators.eh_full`. @@ -73,13 +73,13 @@ def eh_full( ch = curl_back(dxes[1]) ce = curl_forward(dxes[0]) - 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_1(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]: + return (ch(h) - 1j * omega * epsilon * e, + ce(e) + 1j * omega * h) - def op_mu(e: cfdfield, 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 + 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 if mu is None: return op_1 @@ -89,7 +89,7 @@ def eh_full( def e2h( omega: complex, dxes: dx_lists_t, - mu: fdfield | None = None, + mu: fdfield_t | 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) -> cfdfield_t: - return cfdfield_t(ce(e) / (-1j * omega)) + def e2h_1_1(e: cfdfield_t) -> cfdfield_t: + return ce(e) / (-1j * omega) - def e2h_mu(e: cfdfield) -> cfdfield_t: - return cfdfield_t(ce(e) / (-1j * omega * mu)) # type: ignore # mu=None ok + def e2h_mu(e: cfdfield_t) -> cfdfield_t: + return 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 | None = None, + mu: fdfield_t | 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) -> cfdfield_t: + def m2j_mu(m: cfdfield_t) -> cfdfield_t: J = ch(m / mu) / (-1j * omega) # type: ignore # mu=None ok - return cfdfield_t(J) + return J - def m2j_1(m: cfdfield) -> cfdfield_t: + def m2j_1(m: cfdfield_t) -> cfdfield_t: J = ch(m) / (-1j * omega) - return cfdfield_t(J) + return J if mu is None: return m2j_1 @@ -152,11 +152,11 @@ def m2j( def e_tfsf_source( - TF_region: fdfield, + TF_region: fdfield_t, omega: complex, dxes: dx_lists_t, - epsilon: fdfield, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | 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) -> cfdfield_t: + def op(e: cfdfield_t) -> cfdfield_t: neg_iwj = A(TF_region * e) - TF_region * A(e) - return cfdfield_t(neg_iwj / (-1j * omega)) + return neg_iwj / (-1j * omega) return op -def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield, cfdfield], cfdfield_t]: +def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield_t, cfdfield_t], 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, cfdfield], cfdfi Returns: Function `f` that returns E x H as required for the poynting vector. """ - def exh(e: cfdfield, h: cfdfield) -> cfdfield_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] @@ -217,5 +217,5 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[cfdfield, cfdfield], cfdfi 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 cfdfield_t(s) + return s return exh diff --git a/meanas/fdfd/operators.py b/meanas/fdfd/operators.py index 829f43e..8c16ef7 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.sparray`) representations of +These functions return sparse-matrix (`scipy.sparse.spmatrix`) 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, vcfdfield +from ..fdmath import vec, dx_lists_t, vfdfield_t, vcfdfield_t 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 | vcfdfield, - mu: vfdfield | None = None, - pec: vfdfield | None = None, - pmc: vfdfield | None = None, - ) -> sparse.sparray: + epsilon: vfdfield_t | vcfdfield_t, + mu: vfdfield_t | None = None, + pec: vfdfield_t | None = None, + pmc: vfdfield_t | None = None, + ) -> sparse.spmatrix: 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_array(epsilon.size) + pe = sparse.eye(epsilon.size) else: - pe = sparse.diags_array(numpy.where(pec, 0, 1)) # Set pe to (not PEC) + pe = sparse.diags(numpy.where(pec, 0, 1)) # Set pe to (not PEC) if pmc is None: - pm = sparse.eye_array(epsilon.size) + pm = sparse.eye(epsilon.size) else: - pm = sparse.diags_array(numpy.where(pmc, 0, 1)) # set pm to (not PMC) + pm = sparse.diags(numpy.where(pmc, 0, 1)) # set pm to (not PMC) - e = sparse.diags_array(epsilon) + e = sparse.diags(epsilon) if mu is None: - m_div = sparse.eye_array(epsilon.size) + m_div = sparse.eye(epsilon.size) else: - m_div = sparse.diags_array(1 / mu) + m_div = sparse.diags(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.sparray, sparse.sparray]: + ) -> tuple[sparse.spmatrix, sparse.spmatrix]: """ 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_array(p_vector) - P_right = sparse.diags_array(1 / p_vector) + P_left = sparse.diags(p_vector) + P_right = sparse.diags(1 / p_vector) return P_left, P_right def h_full( omega: complex, dxes: dx_lists_t, - epsilon: vfdfield, - mu: vfdfield | None = None, - pec: vfdfield | None = None, - pmc: vfdfield | None = None, - ) -> sparse.sparray: + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, + pec: vfdfield_t | None = None, + pmc: vfdfield_t | None = None, + ) -> sparse.spmatrix: 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_array(epsilon.size) + pe = sparse.eye(epsilon.size) else: - pe = sparse.diags_array(numpy.where(pec, 0, 1)) # set pe to (not PEC) + pe = sparse.diags(numpy.where(pec, 0, 1)) # set pe to (not PEC) if pmc is None: - pm = sparse.eye_array(epsilon.size) + pm = sparse.eye(epsilon.size) else: - pm = sparse.diags_array(numpy.where(pmc, 0, 1)) # Set pe to (not PMC) + pm = sparse.diags(numpy.where(pmc, 0, 1)) # Set pe to (not PMC) - e_div = sparse.diags_array(1 / epsilon) + e_div = sparse.diags(1 / epsilon) if mu is None: - m = sparse.eye_array(epsilon.size) + m = sparse.eye(epsilon.size) else: - m = sparse.diags_array(mu) + m = sparse.diags(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, - mu: vfdfield | None = None, - pec: vfdfield | None = None, - pmc: vfdfield | None = None, - ) -> sparse.sparray: + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, + pec: vfdfield_t | None = None, + pmc: vfdfield_t | None = None, + ) -> sparse.spmatrix: 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_array(epsilon.size) + pe = sparse.eye(epsilon.size) else: - pe = sparse.diags_array(numpy.where(pec, 0, 1)) # set pe to (not PEC) + pe = sparse.diags(numpy.where(pec, 0, 1)) # set pe to (not PEC) if pmc is None: - pm = sparse.eye_array(epsilon.size) + pm = sparse.eye(epsilon.size) else: - pm = sparse.diags_array(numpy.where(pmc, 0, 1)) # set pm to (not PMC) + pm = sparse.diags(numpy.where(pmc, 0, 1)) # set pm to (not PMC) - iwe = pe @ (1j * omega * sparse.diags_array(epsilon)) @ pe + iwe = pe @ (1j * omega * sparse.diags(epsilon)) @ pe iwm = 1j * omega if mu is not None: - iwm *= sparse.diags_array(mu) + iwm *= sparse.diags(mu) iwm = pm @ iwm @ pm A1 = pe @ curl_back(dxes[1]) @ pm A2 = pm @ curl_forward(dxes[0]) @ pe - A = sparse.block_array([[-iwe, A1], - [A2, iwm]]) + A = sparse.bmat([[-iwe, A1], + [A2, iwm]]) return A def e2h( omega: complex, dxes: dx_lists_t, - mu: vfdfield | None = None, - pmc: vfdfield | None = None, - ) -> sparse.sparray: + mu: vfdfield_t | None = None, + pmc: vfdfield_t | None = None, + ) -> sparse.spmatrix: """ 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_array(1 / mu) @ op + op = sparse.diags(1 / mu) @ op if pmc is not None: - op = sparse.diags_array(numpy.where(pmc, 0, 1)) @ op + op = sparse.diags(numpy.where(pmc, 0, 1)) @ op return op @@ -285,8 +285,8 @@ def e2h( def m2j( omega: complex, dxes: dx_lists_t, - mu: vfdfield | None = None, - ) -> sparse.sparray: + mu: vfdfield_t | None = None, + ) -> sparse.spmatrix: """ 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_array(1 / mu) + op = op @ sparse.diags(1 / mu) return op -def poynting_e_cross(e: vcfdfield, dxes: dx_lists_t) -> sparse.sparray: +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. @@ -330,13 +330,13 @@ def poynting_e_cross(e: vcfdfield, dxes: dx_lists_t) -> sparse.sparray: block_diags = [[ None, fx @ -Ez, fx @ Ey], [ fy @ Ez, None, fy @ -Ex], [ fz @ -Ey, fz @ Ex, None]] - 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)) + 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)) return P -def poynting_h_cross(h: vcfdfield, dxes: dx_lists_t) -> sparse.sparray: +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. @@ -353,23 +353,23 @@ def poynting_h_cross(h: vcfdfield, dxes: dx_lists_t) -> sparse.sparray: 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_array(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg, strict=True)) + Hx, Hy, Hz = (sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg, strict=True)) - P = (sparse.block_array( + P = (sparse.bmat( [[ None, -Hz @ fx, Hy @ fx], [ Hz @ fy, None, -Hx @ fy], [-Hy @ fz, Hx @ fz, None]]) - @ sparse.diags_array(numpy.concatenate(dxag))) + @ sparse.diags(numpy.concatenate(dxag))) return P def e_tfsf_source( - TF_region: vfdfield, + TF_region: vfdfield_t, omega: complex, dxes: dx_lists_t, - epsilon: vfdfield, - mu: vfdfield | None = None, - ) -> sparse.sparray: + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, + ) -> sparse.spmatrix: """ 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_array(TF_region) + Q = sparse.diags(TF_region) return (A @ Q - Q @ A) / (-1j * omega) def e_boundary_source( - mask: vfdfield, + mask: vfdfield_t, omega: complex, dxes: dx_lists_t, - epsilon: vfdfield, - mu: vfdfield | None = None, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, periodic_mask_edges: bool = False, - ) -> sparse.sparray: + ) -> sparse.spmatrix: """ Operator that turns an E-field distrubtion into a current (J) distribution along the edges (external and internal) of the provided mask. This is just an @@ -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.sparray: + def shift_rot(axis: int, polarity: int) -> sparse.spmatrix: return shift_circ(axis=axis, shape=shape, shift_distance=polarity) - def shift_mir(axis: int, polarity: int) -> sparse.sparray: + def shift_mir(axis: int, polarity: int) -> sparse.spmatrix: 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_array(numpy.prod(shape)) # shifted minus original + r = shift(axis, polarity) - sparse.eye(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_array(jmask.astype(int)) @ full + return sparse.diags(jmask.astype(int)) @ full diff --git a/meanas/fdfd/solvers.py b/meanas/fdfd/solvers.py index c0aed44..81d1d09 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, vcfdfield, vcfdfield_t +from ..fdmath import dx_lists_t, vfdfield_t, vcfdfield_t from . import operators @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) def _scipy_qmr( - A: scipy.sparse.csr_array, + A: scipy.sparse.csr_matrix, b: ArrayLike, **kwargs: Any, ) -> NDArray[numpy.float64]: @@ -66,16 +66,16 @@ def _scipy_qmr( def generic( omega: complex, dxes: dx_lists_t, - J: vcfdfield, - epsilon: vfdfield, - mu: vfdfield | None = None, + J: vcfdfield_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, *, - pec: vfdfield | None = None, - pmc: vfdfield | None = None, + pec: vfdfield_t | None = None, + pmc: vfdfield_t | None = None, adjoint: bool = False, matrix_solver: Callable[..., ArrayLike] = _scipy_qmr, matrix_solver_opts: dict[str, Any] | None = None, - E_guess: vcfdfield | None = None, + E_guess: vcfdfield_t | 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_array`; + where `A`: `scipy.sparse.csr_matrix`; `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 vcfdfield_t(x0) + return x0 diff --git a/meanas/fdfd/waveguide_2d.py b/meanas/fdfd/waveguide_2d.py index e8f766b..5fda683 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 +from numpy.typing import NDArray, ArrayLike from numpy.linalg import norm from scipy import sparse from ..fdmath.operators import deriv_forward, deriv_back, cross -from ..fdmath import vec, unvec, dx_lists2_t, vcfdfield2_t, vcfdslice_t, vcfdfield2, vfdslice, vcfdslice +from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration @@ -194,10 +194,10 @@ __author__ = 'Jan Petykiewicz' def operator_e( omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, - ) -> sparse.sparray: + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, + ) -> sparse.spmatrix: 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_array(numpy.hstack((eps_parts[0], eps_parts[1]))) - eps_z_inv = sparse.diags_array(1 / eps_parts[2]) + eps_xy = sparse.diags(numpy.hstack((eps_parts[0], eps_parts[1]))) + eps_z_inv = sparse.diags(1 / eps_parts[2]) mu_parts = numpy.split(mu, 3) - mu_yx = sparse.diags_array(numpy.hstack((mu_parts[1], mu_parts[0]))) - mu_z_inv = sparse.diags_array(1 / mu_parts[2]) + mu_yx = sparse.diags(numpy.hstack((mu_parts[1], mu_parts[0]))) + mu_z_inv = sparse.diags(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_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, - ) -> sparse.sparray: + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, + ) -> sparse.spmatrix: 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_array(numpy.hstack((eps_parts[1], eps_parts[0]))) - eps_z_inv = sparse.diags_array(1 / eps_parts[2]) + eps_yx = sparse.diags(numpy.hstack((eps_parts[1], eps_parts[0]))) + eps_z_inv = sparse.diags(1 / eps_parts[2]) mu_parts = numpy.split(mu, 3) - mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1]))) - mu_z_inv = sparse.diags_array(1 / mu_parts[2]) + mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1]))) + mu_z_inv = sparse.diags(1 / mu_parts[2]) op = ( omega * omega * eps_yx @ mu_xy @@ -331,14 +331,14 @@ def operator_h( def normalized_fields_e( - e_xy: vcfdfield2, + e_xy: ArrayLike, wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, prop_phase: float = 0, - ) -> tuple[vcfdslice_t, vcfdslice_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. @@ -366,14 +366,14 @@ def normalized_fields_e( def normalized_fields_h( - h_xy: vcfdfield2, + h_xy: ArrayLike, wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, prop_phase: float = 0, - ) -> tuple[vcfdslice_t, vcfdslice_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. @@ -401,19 +401,21 @@ def normalized_fields_h( def _normalized_fields( - e: vcfdslice, - h: vcfdslice, + e: vcfdfield_t, + h: vcfdfield_t, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, prop_phase: float = 0, - ) -> tuple[vcfdslice_t, vcfdslice_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)] # 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=}' @@ -434,16 +436,16 @@ def _normalized_fields( e *= norm_factor h *= norm_factor - return vcfdslice_t(e), vcfdslice_t(h) + return e, h def exy2h( wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None - ) -> sparse.sparray: + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, into a vectorized H containing all three H components @@ -466,10 +468,10 @@ def exy2h( def hxy2e( wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None - ) -> sparse.sparray: + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, into a vectorized E containing all three E components @@ -491,9 +493,9 @@ def hxy2e( def hxy2h( wavenumber: complex, - dxes: dx_lists2_t, - mu: vfdslice | None = None - ) -> sparse.sparray: + dxes: dx_lists_t, + mu: vfdfield_t | None = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, into a vectorized H containing all three H components @@ -512,22 +514,22 @@ def hxy2h( if mu is not None: mu_parts = numpy.split(mu, 3) - mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1]))) - mu_z_inv = sparse.diags_array(1 / mu_parts[2]) + mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1]))) + mu_z_inv = sparse.diags(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_array(2 * n_pts), + op = sparse.vstack((sparse.eye(2 * n_pts), hxy2hz)) return op def exy2e( wavenumber: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - ) -> sparse.sparray: + dxes: dx_lists_t, + epsilon: vfdfield_t, + ) -> sparse.spmatrix: 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 @@ -573,13 +575,13 @@ def exy2e( if epsilon is not None: epsilon_parts = numpy.split(epsilon, 3) - epsilon_xy = sparse.diags_array(numpy.hstack((epsilon_parts[0], epsilon_parts[1]))) - epsilon_z_inv = sparse.diags_array(1 / epsilon_parts[2]) + epsilon_xy = sparse.diags(numpy.hstack((epsilon_parts[0], epsilon_parts[1]))) + epsilon_z_inv = sparse.diags(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_array(2 * n_pts), + op = sparse.vstack((sparse.eye(2 * n_pts), exy2ez)) return op @@ -587,12 +589,12 @@ def exy2e( def e2h( wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - mu: vfdslice | None = None - ) -> sparse.sparray: + dxes: dx_lists_t, + mu: vfdfield_t | None = None + ) -> sparse.spmatrix: """ Returns an operator which, when applied to a vectorized E eigenfield, produces - the vectorized H eigenfield slice. + the vectorized H eigenfield. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -605,19 +607,19 @@ def e2h( """ op = curl_e(wavenumber, dxes) / (-1j * omega) if mu is not None: - op = sparse.diags_array(1 / mu) @ op + op = sparse.diags(1 / mu) @ op return op def h2e( wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - ) -> sparse.sparray: + dxes: dx_lists_t, + epsilon: vfdfield_t + ) -> sparse.spmatrix: """ Returns an operator which, when applied to a vectorized H eigenfield, produces - the vectorized E eigenfield slice. + the vectorized E eigenfield. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -628,13 +630,13 @@ def h2e( Returns: Sparse matrix representation of the operator. """ - op = sparse.diags_array(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes) + op = sparse.diags(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes) return op -def curl_e(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray: +def curl_e(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: """ - Discretized curl operator for use with the waveguide E field slice. + Discretized curl operator for use with the waveguide E field. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -643,18 +645,18 @@ def curl_e(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray: Returns: Sparse matrix representation of the operator. """ - nn = 1 - for dd in dxes[0]: - nn *= len(dd) + n = 1 + for d in dxes[0]: + n *= len(d) - Bz = -1j * wavenumber * sparse.eye_array(nn) + Bz = -1j * wavenumber * sparse.eye(n) Dfx, Dfy = deriv_forward(dxes[0]) return cross([Dfx, Dfy, Bz]) -def curl_h(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray: +def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix: """ - Discretized curl operator for use with the waveguide H field slice. + Discretized curl operator for use with the waveguide H field. Args: wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` @@ -663,22 +665,22 @@ def curl_h(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray: Returns: Sparse matrix representation of the operator. """ - nn = 1 - for dd in dxes[1]: - nn *= len(dd) + n = 1 + for d in dxes[1]: + n *= len(d) - Bz = -1j * wavenumber * sparse.eye_array(nn) + Bz = -1j * wavenumber * sparse.eye(n) Dbx, Dby = deriv_back(dxes[1]) return cross([Dbx, Dby, Bz]) def h_err( - h: vcfdslice, + h: vcfdfield_t, wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None ) -> float: """ Calculates the relative error in the H field @@ -697,7 +699,7 @@ def h_err( ce = curl_e(wavenumber, dxes) ch = curl_h(wavenumber, dxes) - eps_inv = sparse.diags_array(1 / epsilon) + eps_inv = sparse.diags(1 / epsilon) if mu is None: op = ce @ eps_inv @ ch @ h - omega ** 2 * h @@ -708,12 +710,12 @@ def h_err( def e_err( - e: vcfdslice, + e: vcfdfield_t, wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, ) -> float: """ Calculates the relative error in the E field @@ -735,21 +737,21 @@ def e_err( if mu is None: op = ch @ ce @ e - omega ** 2 * (epsilon * e) else: - mu_inv = sparse.diags_array(1 / mu) + mu_inv = sparse.diags(1 / mu) op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e) return float(norm(op) / norm(e)) def sensitivity( - e_norm: vcfdslice, - h_norm: vcfdslice, + e_norm: vcfdfield_t, + h_norm: vcfdfield_t, wavenumber: complex, omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, - ) -> vcfdslice_t: + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, + ) -> vcfdfield_t: r""" Given a waveguide structure (`dxes`, `epsilon`, `mu`) and mode fields (`e_norm`, `h_norm`, `wavenumber`, `omega`), calculates the sensitivity of the wavenumber @@ -823,11 +825,11 @@ def sensitivity( Dbx, Dby = deriv_back(dxes[1]) eps_x, eps_y, eps_z = numpy.split(epsilon, 3) - eps_xy = sparse.diags_array(numpy.hstack((eps_x, eps_y))) - eps_z_inv = sparse.diags_array(1 / eps_z) + eps_xy = sparse.diags(numpy.hstack((eps_x, eps_y))) + eps_z_inv = sparse.diags(1 / eps_z) mu_x, mu_y, _mu_z = numpy.split(mu, 3) - mu_yx = sparse.diags_array(numpy.hstack((mu_y, mu_x))) + mu_yx = sparse.diags(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]) @@ -841,15 +843,15 @@ def sensitivity( norm = hv_yx_conj @ ev_xy sens_tot = numpy.concatenate([sens_xy1 + sens_xy2, sens_z]) / (2 * wavenumber * norm) - return vcfdslice_t(sens_tot) + return sens_tot def solve_modes( mode_numbers: Sequence[int], omega: complex, - dxes: dx_lists2_t, - epsilon: vfdslice, - mu: vfdslice | None = None, + dxes: dx_lists_t, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, mode_margin: int = 2, ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: """ @@ -905,7 +907,7 @@ def solve_mode( mode_number: int, *args: Any, **kwargs: Any, - ) -> tuple[vcfdfield2_t, complex]: + ) -> tuple[vcfdfield_t, complex]: """ Wrapper around `solve_modes()` that solves for a single mode. @@ -919,13 +921,13 @@ def solve_mode( """ kwargs['mode_numbers'] = [mode_number] e_xys, wavenumbers = solve_modes(*args, **kwargs) - return vcfdfield2_t(e_xys[0]), wavenumbers[0] + return e_xys[0], wavenumbers[0] def inner_product( # TODO documentation - e1: vcfdfield2, - h2: vcfdfield2, - dxes: dx_lists2_t, + e1: vcfdfield_t, + h2: vcfdfield_t, + dxes: dx_lists_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 da50533..8bb0513 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, cast +from typing import Any from collections.abc import Sequence import numpy from numpy.typing import NDArray from numpy import complexfloating -from ..fdmath import vec, unvec, dx_lists_t, cfdfield_t, fdfield, cfdfield +from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, cfdfield_t from . import operators, waveguide_2d @@ -21,8 +21,8 @@ def solve_mode( axis: int, polarity: int, slices: Sequence[slice], - epsilon: fdfield, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | None = None, ) -> dict[str, complex | NDArray[complexfloating]]: """ Given a 3D grid, selects a slice from the grid and attempts to @@ -95,10 +95,9 @@ 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 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) + 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) results = { 'wavenumber': wavenumber, @@ -110,15 +109,15 @@ def solve_mode( def compute_source( - E: cfdfield, + E: cfdfield_t, wavenumber: complex, omega: complex, dxes: dx_lists_t, axis: int, polarity: int, slices: Sequence[slice], - epsilon: fdfield, - mu: fdfield | None = None, + epsilon: fdfield_t, + mu: fdfield_t | None = None, ) -> cfdfield_t: """ Given an eigenmode obtained by `solve_mode`, returns the current source distribution @@ -152,32 +151,35 @@ 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 cfdfield_t(J) + return J def compute_overlap_e( - E: cfdfield, + E: cfdfield_t, wavenumber: complex, dxes: dx_lists_t, axis: int, polarity: int, slices: Sequence[slice], - ) -> cfdfield_t: + ) -> 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) [assumes reflection symmetry]. - TODO: add reference or derivation for compute_overlap_e + TODO: add reference Args: E: E-field of the mode + H: H-field of the mode (advanced by half of a Yee cell from E) wavenumber: Wavenumber of the mode + omega: Angular frequency of the simulation dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` axis: Propagation axis (0=x, 1=y, 2=z) polarity: Propagation direction (+1 for +ve, -1 for -ve) slices: `epsilon[tuple(slices)]` is used to select the portion of the grid to use as the waveguide cross-section. slices[axis] should select only one item. + mu: Magnetic permeability (default 1 everywhere) Returns: overlap_e such that `numpy.sum(overlap_e * other_e.conj())` computes the overlap integral @@ -196,12 +198,12 @@ def compute_overlap_e( Etgt = numpy.zeros_like(Ee) Etgt[slices2] = Ee[slices2] - Etgt /= (Etgt.conj() * Etgt).sum() # type: ignore - return cfdfield_t(Etgt) + Etgt /= (Etgt.conj() * Etgt).sum() + return Etgt def expand_e( - E: cfdfield, + E: cfdfield_t, wavenumber: complex, dxes: dx_lists_t, axis: int, @@ -246,4 +248,4 @@ def expand_e( slices_in = (slice(None), *slices) E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in] - return cfdfield_t(E_expanded) + return E_expanded diff --git a/meanas/fdfd/waveguide_cyl.py b/meanas/fdfd/waveguide_cyl.py index 597e1cb..2134806 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_lists2_t, vcfdslice_t, vcfdfield2_t, vfdslice, vcfdslice, vcfdfield2 +from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t 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_lists2_t, - epsilon: vfdslice, + dxes: dx_lists_t, + epsilon: vfdfield_t, rmin: float, - ) -> sparse.sparray: + ) -> sparse.spmatrix: 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_lists2_t, - epsilon: vfdslice, + dxes: dx_lists_t, + epsilon: vfdfield_t, rmin: float, mode_margin: int = 2, - ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: + ) -> tuple[vcfdfield_t, 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[vcfdslice, complex]: + ) -> tuple[vcfdfield_t, complex]: """ Wrapper around `solve_modes()` that solves for a single mode. @@ -222,10 +222,10 @@ def solve_mode( def linear_wavenumbers( - e_xys: list[vcfdfield2_t], + e_xys: vcfdfield_t, angular_wavenumbers: ArrayLike, - epsilon: vfdslice, - dxes: dx_lists2_t, + epsilon: vfdfield_t, + dxes: dx_lists_t, rmin: float, ) -> NDArray[numpy.complex128]: """ @@ -247,6 +247,7 @@ 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]) @@ -264,11 +265,11 @@ def linear_wavenumbers( def exy2h( angular_wavenumber: complex, omega: float, - dxes: dx_lists2_t, + dxes: dx_lists_t, rmin: float, - epsilon: vfdslice, - mu: vfdslice | None = None - ) -> sparse.sparray: + epsilon: vfdfield_t, + mu: vfdfield_t | None = None + ) -> sparse.spmatrix: """ Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, into a vectorized H containing all three H components @@ -293,10 +294,10 @@ def exy2h( def exy2e( angular_wavenumber: complex, omega: float, - dxes: dx_lists2_t, + dxes: dx_lists_t, rmin: float, - epsilon: vfdslice, - ) -> sparse.sparray: + epsilon: vfdfield_t, + ) -> sparse.spmatrix: """ Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, into a vectorized E containing all three E components @@ -322,7 +323,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]) @@ -330,6 +331,8 @@ 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) @@ -349,10 +352,10 @@ def exy2e( def e2h( angular_wavenumber: complex, omega: float, - dxes: dx_lists2_t, + dxes: dx_lists_t, rmin: float, - mu: vfdslice | None = None - ) -> sparse.sparray: + mu: vfdfield_t | None = None + ) -> sparse.spmatrix: r""" Returns an operator which, when applied to a vectorized E eigenfield, produces the vectorized H eigenfield. @@ -393,7 +396,7 @@ def e2h( def dxes2T( - dxes: dx_lists2_t, + dxes: dx_lists_t, rmin: float, ) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64]]: r""" @@ -418,15 +421,15 @@ def dxes2T( def normalized_fields_e( - e_xy: vcfdfield2, + e_xy: ArrayLike, angular_wavenumber: complex, omega: float, - dxes: dx_lists2_t, + dxes: dx_lists_t, rmin: float, - epsilon: vfdslice, - mu: vfdslice | None = None, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, prop_phase: float = 0, - ) -> tuple[vcfdslice_t, vcfdslice_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. @@ -456,21 +459,23 @@ def normalized_fields_e( def _normalized_fields( - e: vcfdslice, - h: vcfdslice, + e: vcfdfield_t, + h: vcfdfield_t, omega: complex, - dxes: dx_lists2_t, - rmin: float, # Currently unused, but may want to use cylindrical poynting - epsilon: vfdslice, - mu: vfdslice | None = None, + dxes: dx_lists_t, + rmin: float, + epsilon: vfdfield_t, + mu: vfdfield_t | None = None, prop_phase: float = 0, - ) -> tuple[vcfdslice_t, vcfdslice_t]: + ) -> tuple[vcfdfield_t, vcfdfield_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=}' @@ -488,6 +493,10 @@ def _normalized_fields( norm_factor = sign * norm_amplitude * numpy.exp(1j * norm_angle) + print('\nAAA\n', waveguide_2d.inner_product(e, h, dxes, prop_phase=prop_phase)) e *= norm_factor h *= norm_factor - return vcfdslice_t(e), vcfdslice_t(h) + print(f'{sign=} {norm_amplitude=} {norm_angle=} {prop_phase=}') + print(waveguide_2d.inner_product(e, h, dxes, prop_phase=prop_phase)) + + return e, h diff --git a/meanas/fdmath/__init__.py b/meanas/fdmath/__init__.py index 857cf18..b1d8354 100644 --- a/meanas/fdmath/__init__.py +++ b/meanas/fdmath/__init__.py @@ -742,34 +742,12 @@ 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, - 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_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_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 19ccb80..946eb88 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.sparray: + ) -> sparse.spmatrix: """ 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_array(vij, shape=(n, n)) + d = sparse.csr_matrix(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.sparray: + ) -> sparse.spmatrix: """ Utility operator for performing an n-element shift along a specified axis, with mirror boundary conditions applied to the cells beyond the receding edge. @@ -92,13 +92,13 @@ def shift_with_mirror( vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C'))) - d = sparse.csr_array(vij, shape=(n, n)) + d = sparse.csr_matrix(vij, shape=(n, n)) return d def deriv_forward( dx_e: Sequence[NDArray[floating | complexfloating]], - ) -> list[sparse.sparray]: + ) -> list[sparse.spmatrix]: """ 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.sparray: - return shift_circ(axis, shape, 1) - sparse.eye_array(n) + def deriv(axis: int) -> sparse.spmatrix: + return shift_circ(axis, shape, 1) - sparse.eye(n) - Ds = [sparse.diags_array(+1 / dx.ravel(order='C')) @ deriv(a) + Ds = [sparse.diags(+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.sparray]: + ) -> list[sparse.spmatrix]: """ 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.sparray: - return shift_circ(axis, shape, -1) - sparse.eye_array(n) + def deriv(axis: int) -> sparse.spmatrix: + return shift_circ(axis, shape, -1) - sparse.eye(n) - Ds = [sparse.diags_array(-1 / dx.ravel(order='C')) @ deriv(a) + Ds = [sparse.diags(-1 / dx.ravel(order='C')) @ deriv(a) for a, dx in enumerate(dx_h_expanded)] return Ds def cross( - B: Sequence[sparse.sparray], - ) -> sparse.sparray: + B: Sequence[sparse.spmatrix], + ) -> sparse.spmatrix: """ Cross product operator @@ -164,14 +164,13 @@ def cross( Sparse matrix corresponding to (B x), where x is the cross product. """ n = B[0].shape[0] - 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]]) + zero = sparse.csr_matrix((n, n)) + return sparse.bmat([[zero, -B[2], B[1]], + [B[2], zero, -B[0]], + [-B[1], B[0], zero]]) -def vec_cross(b: vfdfield_t) -> sparse.sparray: +def vec_cross(b: vfdfield_t) -> sparse.spmatrix: """ Vector cross product operator @@ -183,11 +182,11 @@ def vec_cross(b: vfdfield_t) -> sparse.sparray: Sparse matrix corresponding to (b x), where x is the cross product. """ - B = [sparse.diags_array(c) for c in numpy.split(b, 3)] + B = [sparse.diags(c) for c in numpy.split(b, 3)] return cross(B) -def avg_forward(axis: int, shape: Sequence[int]) -> sparse.sparray: +def avg_forward(axis: int, shape: Sequence[int]) -> sparse.spmatrix: """ Forward average operator `(x4 = (x4 + x5) / 2)` @@ -202,10 +201,10 @@ def avg_forward(axis: int, shape: Sequence[int]) -> sparse.sparray: raise Exception(f'Invalid shape: {shape}') n = numpy.prod(shape) - return 0.5 * (sparse.eye_array(n) + shift_circ(axis, shape)) + return 0.5 * (sparse.eye(n) + shift_circ(axis, shape)) -def avg_back(axis: int, shape: Sequence[int]) -> sparse.sparray: +def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix: """ Backward average operator `(x4 = (x4 + x3) / 2)` @@ -221,7 +220,7 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.sparray: def curl_forward( dx_e: Sequence[NDArray[floating | complexfloating]], - ) -> sparse.sparray: + ) -> sparse.spmatrix: """ Curl operator for use with the E field. @@ -237,7 +236,7 @@ def curl_forward( def curl_back( dx_h: Sequence[NDArray[floating | complexfloating]], - ) -> sparse.sparray: + ) -> sparse.spmatrix: """ Curl operator for use with the H field. diff --git a/meanas/fdmath/types.py b/meanas/fdmath/types.py index 222d18a..d44b30a 100644 --- a/meanas/fdmath/types.py +++ b/meanas/fdmath/types.py @@ -1,64 +1,25 @@ """ 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 = NewType('fdfield_t', NDArray[floating]) -type fdfield = fdfield_t | NDArray[floating] +fdfield_t = NDArray[floating] """Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)""" -vfdfield_t = NewType('vfdfield_t', NDArray[floating]) -type vfdfield = vfdfield_t | NDArray[floating] +vfdfield_t = NDArray[floating] """Linearized vector field (single vector of length 3*X*Y*Z)""" -cfdfield_t = NewType('cfdfield_t', NDArray[complexfloating]) -type cfdfield = cfdfield_t | NDArray[complexfloating] +cfdfield_t = NDArray[complexfloating] """Complex vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)""" -vcfdfield_t = NewType('vcfdfield_t', NDArray[complexfloating]) -type vcfdfield = vcfdfield_t | NDArray[complexfloating] +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: @@ -70,23 +31,9 @@ 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 f2c01d0..8f5ff39 100644 --- a/meanas/fdmath/vectorization.py +++ b/meanas/fdmath/vectorization.py @@ -7,13 +7,9 @@ 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, NDArray +from numpy.typing import ArrayLike -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, - ) +from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t @overload @@ -29,28 +25,12 @@ def vec(f: cfdfield_t) -> vcfdfield_t: pass @overload -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: +def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t: pass def vec( - 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: + f: fdfield_t | cfdfield_t | ArrayLike | None, + ) -> vfdfield_t | vcfdfield_t | None: """ Create a 1D ndarray from a vector field which spans a 1-3D region. @@ -65,7 +45,7 @@ def vec( """ if f is None: return None - return numpy.ravel(f, order='C') # type: ignore + return numpy.ravel(f, order='C') @overload @@ -80,31 +60,11 @@ 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 | vfdfield2_t | vcfdfield2_t | vfdslice_t | vcfdslice_t | ArrayLike | None, + v: vfdfield_t | vcfdfield_t | None, shape: Sequence[int], nvdim: int = 3, - ) -> fdfield_t | cfdfield_t | fdfield2_t | cfdfield2_t | fdslice_t | cfdslice_t | NDArray | None: + ) -> fdfield_t | cfdfield_t | 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 @@ -122,5 +82,5 @@ def unvec( """ if v is None: return None - return v.reshape((nvdim, *shape), order='C') # type: ignore + return v.reshape((nvdim, *shape), order='C') diff --git a/meanas/fdtd/energy.py b/meanas/fdtd/energy.py index 76888ca..43ea3a1 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, fdfield +from ..fdmath import dx_lists_t, fdfield_t from ..fdmath.functional import deriv_back @@ -8,8 +8,8 @@ from ..fdmath.functional import deriv_back def poynting( - e: fdfield, - h: fdfield, + e: fdfield_t, + h: fdfield_t, 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 fdfield_t(s) + return s def poynting_divergence( - s: fdfield | None = None, + s: fdfield_t | None = None, *, - e: fdfield | None = None, - h: fdfield | None = None, + e: fdfield_t | None = None, + h: fdfield_t | 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 fdfield_t(ds) + return ds def energy_hstep( - e0: fdfield, - h1: fdfield, - e2: fdfield, - epsilon: fdfield | None = None, - mu: fdfield | None = None, + e0: fdfield_t, + h1: fdfield_t, + e2: fdfield_t, + epsilon: fdfield_t | None = None, + mu: fdfield_t | 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 fdfield_t(u) + return u def energy_estep( - h0: fdfield, - e1: fdfield, - h2: fdfield, - epsilon: fdfield | None = None, - mu: fdfield | None = None, + h0: fdfield_t, + e1: fdfield_t, + h2: fdfield_t, + epsilon: fdfield_t | None = None, + mu: fdfield_t | 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 fdfield_t(u) + return u def delta_energy_h2e( dt: float, - e0: fdfield, - h1: fdfield, - e2: fdfield, - h3: fdfield, - epsilon: fdfield | None = None, - mu: fdfield | None = None, + e0: fdfield_t, + h1: fdfield_t, + e2: fdfield_t, + h3: fdfield_t, + epsilon: fdfield_t | None = None, + mu: fdfield_t | 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 fdfield_t(du) + return du def delta_energy_e2h( dt: float, - h0: fdfield, - e1: fdfield, - h2: fdfield, - e3: fdfield, - epsilon: fdfield | None = None, - mu: fdfield | None = None, + h0: fdfield_t, + e1: fdfield_t, + h2: fdfield_t, + e3: fdfield_t, + epsilon: fdfield_t | None = None, + mu: fdfield_t | 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 fdfield_t(du) + return du def delta_energy_j( - j0: fdfield, - e1: fdfield, + j0: fdfield_t, + e1: fdfield_t, 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 fdfield_t(du) + return du def dxmul( - ee: fdfield, - hh: fdfield, - epsilon: fdfield | float | None = None, - mu: fdfield | float | None = None, + ee: fdfield_t, + hh: fdfield_t, + epsilon: fdfield_t | float | None = None, + mu: fdfield_t | 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 fdfield_t(result) + return result diff --git a/meanas/fdtd/misc.py b/meanas/fdtd/misc.py index 3fb3371..160682d 100644 --- a/meanas/fdtd/misc.py +++ b/meanas/fdtd/misc.py @@ -1,4 +1,5 @@ -from collections.abc import Callable +from typing import Callable +from collections.abc import Sequence import logging import numpy @@ -94,10 +95,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 tt: (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 xx: (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 fc456eb..8098da0 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, fdfield_t, dx_lists_t +from ..fdmath import 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, + epsilon: fdfield_t, *, 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 82587b4..5f2cf11 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, vcfdfield, vfdfield, dx_lists_t +from ..fdmath import vec, unvec 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: dx_lists_t - epsilon: vfdfield + dxes: list[list[NDArray[numpy.float64]]] + epsilon: NDArray[numpy.float64] omega: complex - j: vcfdfield - e: vcfdfield - pmc: vfdfield | None - pec: vfdfield | None + j: NDArray[numpy.complex128] + e: NDArray[numpy.complex128] + pmc: NDArray[numpy.float64] | None + pec: NDArray[numpy.float64] | 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 = dict(atol=1e-15, rtol=1e-11), + J=j_vec, + omega=omega, + dxes=dxes, + epsilon=eps_vec, + matrix_solver_opts={'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 540a3a0..832053d 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, vfdfield, cfdfield_t +from ..fdmath import vec, unvec, dx_lists_mut #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: vfdfield, + epsilon: NDArray[numpy.float64], dxes: dx_lists_mut, omega: float, src_polarity: int, - ) -> cfdfield_t: + ) -> NDArray[numpy.complex128]: 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 03d2b7e..25ee891 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, ...],