jan/meanas
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Rework field types, use sparse arrays instead of matrices, rework eme arg naming, improve type annotations and linter cleanup

Browse Source
This commit is contained in:
jan 2025-12-10 02:14:20 -08:00
parent b7ad5dea2b
commit b486fa325b
20 changed files with 571 additions and 433 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
meanas/eigensolvers.py
View File

@ -64,10 +64,10 @@ def rayleigh_quotient_iteration(
(eigenvalues, eigenvectors) (eigenvalues, eigenvectors)
""" """
try: try:
(operator - sparse.eye(operator.shape[0])) (operator - sparse.eye_array(operator.shape[0]))
def shift(eigval: float) -> sparse: def shift(eigval: float) -> sparse.sparray:
return eigval * sparse.eye(operator.shape[0]) return eigval * sparse.eye_array(operator.shape[0])
if solver is None: if solver is None:
solver = spalg.spsolve solver = spalg.spsolve
@ -130,7 +130,7 @@ def signed_eigensolve(
# Try to combine, use general LinearOperator if we fail # Try to combine, use general LinearOperator if we fail
try: try:
shifted_operator = operator + shift * sparse.eye(operator.shape[0]) shifted_operator = operator + shift * sparse.eye_array(operator.shape[0])
except TypeError: except TypeError:
shifted_operator = operator + spalg.LinearOperator(shape=operator.shape, shifted_operator = operator + spalg.LinearOperator(shape=operator.shape,
matvec=lambda v: shift * v) matvec=lambda v: shift * v)

49
meanas/fdfd/bloch.py
View File

@ -106,7 +106,7 @@ import scipy.optimize
from scipy.linalg import norm from scipy.linalg import norm
import scipy.sparse.linalg as spalg import scipy.sparse.linalg as spalg
from ..fdmath import fdfield_t, cfdfield_t from ..fdmath import fdfield, cfdfield, cfdfield_t
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -183,8 +183,8 @@ def generate_kmn(
def maxwell_operator( def maxwell_operator(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None mu: fdfield | None = None
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
""" """
Generate the Maxwell operator Generate the Maxwell operator
@ -276,7 +276,7 @@ def maxwell_operator(
def hmn_2_exyz( def hmn_2_exyz(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield,
) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
""" """
Generate an operator which converts a vectorized spatial-frequency-space Generate an operator which converts a vectorized spatial-frequency-space
@ -308,7 +308,8 @@ def hmn_2_exyz(
- m * hin_n) * k_mag - m * hin_n) * k_mag
# divide by epsilon # 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 return operator
@ -316,7 +317,7 @@ def hmn_2_exyz(
def hmn_2_hxyz( def hmn_2_hxyz(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t epsilon: fdfield,
) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]: ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
""" """
Generate an operator which converts a vectorized spatial-frequency-space 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: def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2)) hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
h_xyz = (m * hin_m h_xyz = (m * hin_m
+ n * hin_n) # noqa: E128 + n * hin_n)
return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)]) return cfdfield_t(numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)]))
return operator return operator
@ -352,8 +353,8 @@ def hmn_2_hxyz(
def inverse_maxwell_operator_approx( def inverse_maxwell_operator_approx(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]: ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
""" """
Generate an approximate inverse of the Maxwell operator, Generate an approximate inverse of the Maxwell operator,
@ -440,8 +441,8 @@ def find_k(
tolerance: float, tolerance: float,
direction: ArrayLike, direction: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
band: int = 0, band: int = 0,
k_bounds: tuple[float, float] = (0, 0.5), k_bounds: tuple[float, float] = (0, 0.5),
k_guess: float | None = None, k_guess: float | None = None,
@ -508,8 +509,8 @@ def eigsolve(
num_modes: int, num_modes: int,
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
tolerance: float = 1e-7, tolerance: float = 1e-7,
max_iters: int = 10000, max_iters: int = 10000,
reset_iters: int = 100, 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 def Qi_func(theta: float, Qi_memo=Qi_memo, ZtZ=ZtZ, DtD=DtD, symZtD=symZtD) -> float: # noqa: ANN001
if Qi_memo[0] == theta: if Qi_memo[0] == theta:
return cast(float, Qi_memo[1]) return cast('float', Qi_memo[1])
c = numpy.cos(theta) c = numpy.cos(theta)
s = numpy.sin(theta) s = numpy.sin(theta)
@ -668,8 +669,8 @@ def eigsolve(
else: else:
raise Exception('Inexplicable singularity in trace_func') from err raise Exception('Inexplicable singularity in trace_func') from err
Qi_memo[0] = theta Qi_memo[0] = theta
Qi_memo[1] = cast(float, Qi) Qi_memo[1] = cast('float', Qi)
return cast(float, Qi) return cast('float', Qi)
def trace_func(theta: float, ZtAZ=ZtAZ, DtAD=DtAD, symZtAD=symZtAD) -> float: # noqa: ANN001 def trace_func(theta: float, ZtAZ=ZtAZ, DtAD=DtAD, symZtAD=symZtAD) -> float: # noqa: ANN001
c = numpy.cos(theta) 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 # assumes x-axis propagation
assert numpy.array_equal(eR.shape, hR.shape) 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() 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) pp = inner_product(eO, hO, eI, hI)
pn = inner_product(eO, hO, eI, -hI) pn = inner_product(eO, hO, eI, -hI)
np = inner_product(eO, -hO, eI, hI) np = inner_product(eO, -hO, eI, hI)

44
meanas/fdfd/eme.py
View File

@ -1,20 +1,29 @@
from collections.abc import Sequence
import numpy 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 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) nL = len(wavenumbers_L)
nR = len(wavenumbers_R) nR = len(wavenumbers_R)
A12 = numpy.zeros((nL, nR), dtype=complex) A12 = numpy.zeros((nL, nR), dtype=complex)
A21 = numpy.zeros((nL, nR), dtype=complex) A21 = numpy.zeros((nL, nR), dtype=complex)
B11 = numpy.zeros((nL,), dtype=complex) B11 = numpy.zeros((nL,), dtype=complex)
for ll in range(nL): for ll in range(nL):
eL, hL = ehL[ll] eL, hL = ehLs[ll]
B11[ll] = inner_product(eL, hL, dxes=dxes, conj_h=False) B11[ll] = inner_product(eL, hL, dxes=dxes, conj_h=False)
for rr in range(nR): 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? 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) 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 return tt, rr
def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs): def get_abcd(
t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs) ehLs: Sequence[Sequence[vcfdfield2]],
t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs) 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) t21i = numpy.linalg.pinv(t21)
A = t12 - r21 @ t21i @ r12 A = t12 - r21 @ t21i @ r12
B = r21 @ t21i B = r21 @ t21i
@ -44,15 +59,16 @@ def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs):
def get_s( def get_s(
eL_xys, ehLs: Sequence[Sequence[vcfdfield2]],
wavenumbers_L, wavenumbers_L: Sequence[complex],
eR_xys, ehRs: Sequence[Sequence[vcfdfield2]],
wavenumbers_R, wavenumbers_R: Sequence[complex],
force_nogain: bool = False, force_nogain: bool = False,
force_reciprocal: bool = False, force_reciprocal: bool = False,
**kwargs): **kwargs,
t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs) ) -> NDArray[numpy.complex128]:
t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs) 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], ss = numpy.block([[r12, t12],
[t21, r21]]) [t21, r21]])

10
meanas/fdfd/farfield.py
View File

@ -1,14 +1,16 @@
""" """
Functions for performing near-to-farfield transformation (and the reverse). Functions for performing near-to-farfield transformation (and the reverse).
""" """
from typing import Any, cast from typing import Any, cast, TYPE_CHECKING
from collections.abc import Sequence
import numpy import numpy
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
from numpy import pi from numpy import pi
from ..fdmath import cfdfield_t from ..fdmath import cfdfield_t
if TYPE_CHECKING:
from collections.abc import Sequence
def near_to_farfield( def near_to_farfield(
E_near: cfdfield_t, E_near: cfdfield_t,
@ -63,7 +65,7 @@ def near_to_farfield(
padded_size = (2**numpy.ceil(numpy.log2(s))).astype(int) padded_size = (2**numpy.ceil(numpy.log2(s))).astype(int)
if not hasattr(padded_size, '__len__'): if not hasattr(padded_size, '__len__'):
padded_size = (padded_size, padded_size) # type: ignore # checked if sequence padded_size = (padded_size, padded_size) # type: ignore # checked if sequence
padded_shape = cast(Sequence[int], padded_size) padded_shape = cast('Sequence[int]', padded_size)
En_fft = [fftshift(fft2(fftshift(Eni), s=padded_shape)) for Eni in E_near] En_fft = [fftshift(fft2(fftshift(Eni), s=padded_shape)) for Eni in E_near]
Hn_fft = [fftshift(fft2(fftshift(Hni), s=padded_shape)) for Hni in H_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) padded_size = (2 ** numpy.ceil(numpy.log2(s))).astype(int)
if not hasattr(padded_size, '__len__'): if not hasattr(padded_size, '__len__'):
padded_size = (padded_size, padded_size) # type: ignore # checked if sequence padded_size = (padded_size, padded_size) # type: ignore # checked if sequence
padded_shape = cast(Sequence[int], padded_size) padded_shape = cast('Sequence[int]', padded_size)
k = 2 * pi k = 2 * pi
kxs = fftshift(fftfreq(s[0], 1 / (s[0] * dkx))) kxs = fftshift(fftfreq(s[0], 1 / (s[0] * dkx)))

68
meanas/fdfd/functional.py
View File

@ -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 from collections.abc import Callable
import numpy import numpy
from ..fdmath import dx_lists_t, fdfield_t, cfdfield_t, cfdfield_updater_t from ..fdmath import dx_lists_t, cfdfield_t, fdfield, cfdfield, cfdfield_updater_t
from ..fdmath.functional import curl_forward, curl_back from ..fdmath.functional import curl_forward, curl_back
@ -18,8 +18,8 @@ __author__ = 'Jan Petykiewicz'
def e_full( def e_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Wave operator for use with E-field. See `operators.e_full` for details. Wave operator for use with E-field. See `operators.e_full` for details.
@ -37,13 +37,13 @@ def e_full(
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def op_1(e: cfdfield_t) -> cfdfield_t: def op_1(e: cfdfield) -> cfdfield_t:
curls = ch(ce(e)) 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 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: if mu is None:
return op_1 return op_1
@ -53,9 +53,9 @@ def e_full(
def eh_full( def eh_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> Callable[[cfdfield_t, cfdfield_t], tuple[cfdfield_t, cfdfield_t]]: ) -> Callable[[cfdfield, cfdfield], tuple[cfdfield_t, cfdfield_t]]:
""" """
Wave operator for full (both E and H) field representation. Wave operator for full (both E and H) field representation.
See `operators.eh_full`. See `operators.eh_full`.
@ -73,13 +73,13 @@ def eh_full(
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def op_1(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]: def op_1(e: cfdfield, h: cfdfield) -> tuple[cfdfield_t, cfdfield_t]:
return (ch(h) - 1j * omega * epsilon * e, return (cfdfield_t(ch(h) - 1j * omega * epsilon * e),
ce(e) + 1j * omega * h) cfdfield_t(ce(e) + 1j * omega * h))
def op_mu(e: cfdfield_t, h: cfdfield_t) -> tuple[cfdfield_t, cfdfield_t]: def op_mu(e: cfdfield, h: cfdfield) -> tuple[cfdfield_t, cfdfield_t]:
return (ch(h) - 1j * omega * epsilon * e, return (cfdfield_t(ch(h) - 1j * omega * epsilon * e),
ce(e) + 1j * omega * mu * h) # type: ignore # mu=None ok cfdfield_t(ce(e) + 1j * omega * mu * h)) # type: ignore # mu=None ok
if mu is None: if mu is None:
return op_1 return op_1
@ -89,7 +89,7 @@ def eh_full(
def e2h( def e2h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Utility operator for converting the `E` field into the `H` field. Utility operator for converting the `E` field into the `H` field.
@ -106,11 +106,11 @@ def e2h(
""" """
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def e2h_1_1(e: cfdfield_t) -> cfdfield_t: def e2h_1_1(e: cfdfield) -> cfdfield_t:
return ce(e) / (-1j * omega) return cfdfield_t(ce(e) / (-1j * omega))
def e2h_mu(e: cfdfield_t) -> cfdfield_t: def e2h_mu(e: cfdfield) -> cfdfield_t:
return ce(e) / (-1j * omega * mu) # type: ignore # mu=None ok return cfdfield_t(ce(e) / (-1j * omega * mu)) # type: ignore # mu=None ok
if mu is None: if mu is None:
return e2h_1_1 return e2h_1_1
@ -120,7 +120,7 @@ def e2h(
def m2j( def m2j(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Utility operator for converting magnetic current `M` distribution Utility operator for converting magnetic current `M` distribution
@ -138,13 +138,13 @@ def m2j(
""" """
ch = curl_back(dxes[1]) 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 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) J = ch(m) / (-1j * omega)
return J return cfdfield_t(J)
if mu is None: if mu is None:
return m2j_1 return m2j_1
@ -152,11 +152,11 @@ def m2j(
def e_tfsf_source( def e_tfsf_source(
TF_region: fdfield_t, TF_region: fdfield,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> cfdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Operator that turns an E-field distribution into a total-field/scattered-field Operator that turns an E-field distribution into a total-field/scattered-field
@ -178,13 +178,13 @@ def e_tfsf_source(
# TODO documentation # TODO documentation
A = e_full(omega, dxes, epsilon, mu) 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) neg_iwj = A(TF_region * e) - TF_region * A(e)
return neg_iwj / (-1j * omega) return cfdfield_t(neg_iwj / (-1j * omega))
return op 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""" r"""
Generates a function that takes the single-frequency `E` and `H` fields 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 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: Returns:
Function `f` that returns E x H as required for the poynting vector. 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) s = numpy.empty_like(e)
ex = e[0] * dxes[0][0][:, None, None] ex = e[0] * dxes[0][0][:, None, None]
ey = e[1] * dxes[0][1][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[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[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 s[2] = numpy.roll(ex, -1, axis=2) * hy - numpy.roll(ey, -1, axis=2) * hx
return s return cfdfield_t(s)
return exh return exh

142
meanas/fdfd/operators.py
View File

@ -1,7 +1,7 @@
""" """
Sparse matrix operators for use with electromagnetic wave equations. 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 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. `meanas.fdmath.vectorization.vec()` and `meanas.fdmath.vectorization.unvec()` functions.
@ -30,7 +30,7 @@ The following operators are included:
import numpy import numpy
from scipy import sparse 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 from ..fdmath.operators import shift_with_mirror, shift_circ, curl_forward, curl_back
@ -40,11 +40,11 @@ __author__ = 'Jan Petykiewicz'
def e_full( def e_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t | vcfdfield_t, epsilon: vfdfield | vcfdfield,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
pec: vfdfield_t | None = None, pec: vfdfield | None = None,
pmc: vfdfield_t | None = None, pmc: vfdfield | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Wave operator Wave operator
$$ \nabla \times (\frac{1}{\mu} \nabla \times) - \Omega^2 \epsilon $$ $$ \nabla \times (\frac{1}{\mu} \nabla \times) - \Omega^2 \epsilon $$
@ -77,20 +77,20 @@ def e_full(
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
if pec is None: if pec is None:
pe = sparse.eye(epsilon.size) pe = sparse.eye_array(epsilon.size)
else: 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: if pmc is None:
pm = sparse.eye(epsilon.size) pm = sparse.eye_array(epsilon.size)
else: 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: if mu is None:
m_div = sparse.eye(epsilon.size) m_div = sparse.eye_array(epsilon.size)
else: 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 op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe
return op return op
@ -98,7 +98,7 @@ def e_full(
def e_full_preconditioners( def e_full_preconditioners(
dxes: dx_lists_t, dxes: dx_lists_t,
) -> tuple[sparse.spmatrix, sparse.spmatrix]: ) -> tuple[sparse.sparray, sparse.sparray]:
""" """
Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator. Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator.
@ -118,19 +118,19 @@ def e_full_preconditioners(
dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[0][2][None, None, :]] dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[0][2][None, None, :]]
p_vector = numpy.sqrt(vec(p_squared)) p_vector = numpy.sqrt(vec(p_squared))
P_left = sparse.diags(p_vector) P_left = sparse.diags_array(p_vector)
P_right = sparse.diags(1 / p_vector) P_right = sparse.diags_array(1 / p_vector)
return P_left, P_right return P_left, P_right
def h_full( def h_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
pec: vfdfield_t | None = None, pec: vfdfield | None = None,
pmc: vfdfield_t | None = None, pmc: vfdfield | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Wave operator Wave operator
$$ \nabla \times (\frac{1}{\epsilon} \nabla \times) - \omega^2 \mu $$ $$ \nabla \times (\frac{1}{\epsilon} \nabla \times) - \omega^2 \mu $$
@ -161,20 +161,20 @@ def h_full(
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
if pec is None: if pec is None:
pe = sparse.eye(epsilon.size) pe = sparse.eye_array(epsilon.size)
else: 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: if pmc is None:
pm = sparse.eye(epsilon.size) pm = sparse.eye_array(epsilon.size)
else: 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: if mu is None:
m = sparse.eye(epsilon.size) m = sparse.eye_array(epsilon.size)
else: else:
m = sparse.diags(mu) m = sparse.diags_array(mu)
A = pm @ (ce @ pe @ e_div @ ch - omega**2 * m) @ pm A = pm @ (ce @ pe @ e_div @ ch - omega**2 * m) @ pm
return A return A
@ -183,11 +183,11 @@ def h_full(
def eh_full( def eh_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
pec: vfdfield_t | None = None, pec: vfdfield | None = None,
pmc: vfdfield_t | None = None, pmc: vfdfield | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Wave operator for `[E, H]` field representation. This operator implements Maxwell's Wave operator for `[E, H]` field representation. This operator implements Maxwell's
equations without cancelling out either E or H. The operator is equations without cancelling out either E or H. The operator is
@ -227,35 +227,35 @@ def eh_full(
Sparse matrix containing the wave operator. Sparse matrix containing the wave operator.
""" """
if pec is None: if pec is None:
pe = sparse.eye(epsilon.size) pe = sparse.eye_array(epsilon.size)
else: 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: if pmc is None:
pm = sparse.eye(epsilon.size) pm = sparse.eye_array(epsilon.size)
else: 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 iwm = 1j * omega
if mu is not None: if mu is not None:
iwm *= sparse.diags(mu) iwm *= sparse.diags_array(mu)
iwm = pm @ iwm @ pm iwm = pm @ iwm @ pm
A1 = pe @ curl_back(dxes[1]) @ pm A1 = pe @ curl_back(dxes[1]) @ pm
A2 = pm @ curl_forward(dxes[0]) @ pe A2 = pm @ curl_forward(dxes[0]) @ pe
A = sparse.bmat([[-iwe, A1], A = sparse.block_array([[-iwe, A1],
[A2, iwm]]) [A2, iwm]])
return A return A
def e2h( def e2h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
pmc: vfdfield_t | None = None, pmc: vfdfield | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Utility operator for converting the E field into the H field. Utility operator for converting the E field into the H field.
For use with `e_full()` -- assumes that there is no magnetic current M. 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) op = curl_forward(dxes[0]) / (-1j * omega)
if mu is not None: if mu is not None:
op = sparse.diags(1 / mu) @ op op = sparse.diags_array(1 / mu) @ op
if pmc is not None: 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 return op
@ -285,8 +285,8 @@ def e2h(
def m2j( def m2j(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator for converting a magnetic current M into an electric current J. Operator for converting a magnetic current M into an electric current J.
For use with eg. `e_full()`. For use with eg. `e_full()`.
@ -302,12 +302,12 @@ def m2j(
op = curl_back(dxes[1]) / (1j * omega) op = curl_back(dxes[1]) / (1j * omega)
if mu is not None: if mu is not None:
op = op @ sparse.diags(1 / mu) op = op @ sparse.diags_array(1 / mu)
return op 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 Operator for computing the Poynting vector, containing the
(E x) portion of the Poynting vector. (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], block_diags = [[ None, fx @ -Ez, fx @ Ey],
[ fy @ Ez, None, fy @ -Ex], [ fy @ Ez, None, fy @ -Ex],
[ fz @ -Ey, fz @ Ex, None]] [ fz @ -Ey, fz @ Ex, None]]
block_matrix = sparse.bmat([[sparse.diags(x) if x is not None else None for x in row] 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]) for row in block_diags])
P = block_matrix @ sparse.diags(numpy.concatenate(dxbg)) P = block_matrix @ sparse.diags_array(numpy.concatenate(dxbg))
return P 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. 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')] 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')] 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], [[ None, -Hz @ fx, Hy @ fx],
[ Hz @ fy, None, -Hx @ fy], [ Hz @ fy, None, -Hx @ fy],
[-Hy @ fz, Hx @ fz, None]]) [-Hy @ fz, Hx @ fz, None]])
@ sparse.diags(numpy.concatenate(dxag))) @ sparse.diags_array(numpy.concatenate(dxag)))
return P return P
def e_tfsf_source( def e_tfsf_source(
TF_region: vfdfield_t, TF_region: vfdfield,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator that turns a desired E-field distribution into a Operator that turns a desired E-field distribution into a
total-field/scattered-field (TFSF) source. total-field/scattered-field (TFSF) source.
@ -390,18 +390,18 @@ def e_tfsf_source(
""" """
# TODO documentation # TODO documentation
A = e_full(omega, dxes, epsilon, mu) 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) return (A @ Q - Q @ A) / (-1j * omega)
def e_boundary_source( def e_boundary_source(
mask: vfdfield_t, mask: vfdfield,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
periodic_mask_edges: bool = False, periodic_mask_edges: bool = False,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator that turns an E-field distrubtion into a current (J) distribution 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 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]] shape = [len(dxe) for dxe in dxes[0]]
jmask = numpy.zeros_like(mask, dtype=bool) 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) 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) return shift_with_mirror(axis=axis, shape=shape, shift_distance=polarity)
shift = shift_rot if periodic_mask_edges else shift_mir shift = shift_rot if periodic_mask_edges else shift_mir
@ -436,7 +436,7 @@ def e_boundary_source(
if shape[axis] == 1: if shape[axis] == 1:
continue continue
for polarity in (-1, +1): 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)) r3 = sparse.block_diag((r, r, r))
jmask = numpy.logical_or(jmask, numpy.abs(r3 @ mask)) 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) |
# (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

20
meanas/fdfd/solvers.py
View File

@ -11,7 +11,7 @@ from numpy.typing import ArrayLike, NDArray
from numpy.linalg import norm from numpy.linalg import norm
import scipy.sparse.linalg 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 from . import operators
@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
def _scipy_qmr( def _scipy_qmr(
A: scipy.sparse.csr_matrix, A: scipy.sparse.csr_array,
b: ArrayLike, b: ArrayLike,
**kwargs: Any, **kwargs: Any,
) -> NDArray[numpy.float64]: ) -> NDArray[numpy.float64]:
@ -66,16 +66,16 @@ def _scipy_qmr(
def generic( def generic(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
J: vcfdfield_t, J: vcfdfield,
epsilon: vfdfield_t, epsilon: vfdfield,
mu: vfdfield_t | None = None, mu: vfdfield | None = None,
*, *,
pec: vfdfield_t | None = None, pec: vfdfield | None = None,
pmc: vfdfield_t | None = None, pmc: vfdfield | None = None,
adjoint: bool = False, adjoint: bool = False,
matrix_solver: Callable[..., ArrayLike] = _scipy_qmr, matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
matrix_solver_opts: dict[str, Any] | None = None, matrix_solver_opts: dict[str, Any] | None = None,
E_guess: vcfdfield_t | None = None, E_guess: vcfdfield | None = None,
) -> vcfdfield_t: ) -> vcfdfield_t:
""" """
Conjugate gradient FDFD solver using CSR sparse matrices. 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) (at H-field locations; non-zero value indicates PMC is present)
adjoint: If true, solves the adjoint problem. adjoint: If true, solves the adjoint problem.
matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`, 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`; `b`: `ArrayLike`;
`x`: `ArrayLike`; `x`: `ArrayLike`;
Default is a wrapped version of `scipy.sparse.linalg.qmr()` Default is a wrapped version of `scipy.sparse.linalg.qmr()`
@ -138,4 +138,4 @@ def generic(
else: else:
x0 = Pr @ x x0 = Pr @ x
return x0 return vcfdfield_t(x0)

212
meanas/fdfd/waveguide_2d.py
View File

@ -180,12 +180,12 @@ if the result is introduced into a space with a discretized z-axis.
from typing import Any from typing import Any
from collections.abc import Sequence from collections.abc import Sequence
import numpy import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray
from numpy.linalg import norm from numpy.linalg import norm
from scipy import sparse from scipy import sparse
from ..fdmath.operators import deriv_forward, deriv_back, cross 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 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
@ -194,10 +194,10 @@ __author__ = 'Jan Petykiewicz'
def operator_e( def operator_e(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Waveguide operator of the form Waveguide operator of the form
@ -246,12 +246,12 @@ def operator_e(
Dbx, Dby = deriv_back(dxes[1]) Dbx, Dby = deriv_back(dxes[1])
eps_parts = numpy.split(epsilon, 3) eps_parts = numpy.split(epsilon, 3)
eps_xy = sparse.diags(numpy.hstack((eps_parts[0], eps_parts[1]))) eps_xy = sparse.diags_array(numpy.hstack((eps_parts[0], eps_parts[1])))
eps_z_inv = sparse.diags(1 / eps_parts[2]) eps_z_inv = sparse.diags_array(1 / eps_parts[2])
mu_parts = numpy.split(mu, 3) mu_parts = numpy.split(mu, 3)
mu_yx = sparse.diags(numpy.hstack((mu_parts[1], mu_parts[0]))) mu_yx = sparse.diags_array(numpy.hstack((mu_parts[1], mu_parts[0])))
mu_z_inv = sparse.diags(1 / mu_parts[2]) mu_z_inv = sparse.diags_array(1 / mu_parts[2])
op = ( op = (
omega * omega * mu_yx @ eps_xy omega * omega * mu_yx @ eps_xy
@ -263,10 +263,10 @@ def operator_e(
def operator_h( def operator_h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Waveguide operator of the form Waveguide operator of the form
@ -315,12 +315,12 @@ def operator_h(
Dbx, Dby = deriv_back(dxes[1]) Dbx, Dby = deriv_back(dxes[1])
eps_parts = numpy.split(epsilon, 3) eps_parts = numpy.split(epsilon, 3)
eps_yx = sparse.diags(numpy.hstack((eps_parts[1], eps_parts[0]))) eps_yx = sparse.diags_array(numpy.hstack((eps_parts[1], eps_parts[0])))
eps_z_inv = sparse.diags(1 / eps_parts[2]) eps_z_inv = sparse.diags_array(1 / eps_parts[2])
mu_parts = numpy.split(mu, 3) mu_parts = numpy.split(mu, 3)
mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1]))) mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1])))
mu_z_inv = sparse.diags(1 / mu_parts[2]) mu_z_inv = sparse.diags_array(1 / mu_parts[2])
op = ( op = (
omega * omega * eps_yx @ mu_xy omega * omega * eps_yx @ mu_xy
@ -331,14 +331,14 @@ def operator_h(
def normalized_fields_e( def normalized_fields_e(
e_xy: ArrayLike, e_xy: vcfdfield2,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
prop_phase: float = 0, 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, Given a vector `e_xy` containing the vectorized E_x and E_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -366,14 +366,14 @@ def normalized_fields_e(
def normalized_fields_h( def normalized_fields_h(
h_xy: ArrayLike, h_xy: vcfdfield2,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
prop_phase: float = 0, 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, Given a vector `h_xy` containing the vectorized H_x and H_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -401,21 +401,19 @@ def normalized_fields_h(
def _normalized_fields( def _normalized_fields(
e: vcfdfield_t, e: vcfdslice,
h: vcfdfield_t, h: vcfdslice,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> tuple[vcfdfield_t, vcfdfield_t]: ) -> tuple[vcfdslice_t, vcfdslice_t]:
# TODO documentation # TODO documentation
shape = [s.size for s in dxes[0]] shape = [s.size for s in dxes[0]]
dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)]
# Find time-averaged Sz and normalize to it # 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 # 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 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=}' 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 e *= norm_factor
h *= norm_factor h *= norm_factor
return e, h return vcfdslice_t(e), vcfdslice_t(h)
def exy2h( def exy2h(
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
into a vectorized H containing all three H components into a vectorized H containing all three H components
@ -468,10 +466,10 @@ def exy2h(
def hxy2e( def hxy2e(
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
into a vectorized E containing all three E components into a vectorized E containing all three E components
@ -493,9 +491,9 @@ def hxy2e(
def hxy2h( def hxy2h(
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
into a vectorized H containing all three H components into a vectorized H containing all three H components
@ -514,22 +512,22 @@ def hxy2h(
if mu is not None: if mu is not None:
mu_parts = numpy.split(mu, 3) mu_parts = numpy.split(mu, 3)
mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1]))) mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1])))
mu_z_inv = sparse.diags(1 / mu_parts[2]) mu_z_inv = sparse.diags_array(1 / mu_parts[2])
hxy2hz = mu_z_inv @ hxy2hz @ mu_xy hxy2hz = mu_z_inv @ hxy2hz @ mu_xy
n_pts = dxes[1][0].size * dxes[1][1].size 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)) hxy2hz))
return op return op
def exy2e( def exy2e(
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
into a vectorized E containing all three E components into a vectorized E containing all three E components
@ -575,13 +573,13 @@ def exy2e(
if epsilon is not None: if epsilon is not None:
epsilon_parts = numpy.split(epsilon, 3) epsilon_parts = numpy.split(epsilon, 3)
epsilon_xy = sparse.diags(numpy.hstack((epsilon_parts[0], epsilon_parts[1]))) epsilon_xy = sparse.diags_array(numpy.hstack((epsilon_parts[0], epsilon_parts[1])))
epsilon_z_inv = sparse.diags(1 / epsilon_parts[2]) epsilon_z_inv = sparse.diags_array(1 / epsilon_parts[2])
exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy
n_pts = dxes[0][0].size * dxes[0][1].size 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)) exy2ez))
return op return op
@ -589,12 +587,12 @@ def exy2e(
def e2h( def e2h(
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Returns an operator which, when applied to a vectorized E eigenfield, produces Returns an operator which, when applied to a vectorized E eigenfield, produces
the vectorized H eigenfield. the vectorized H eigenfield slice.
Args: Args:
wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` 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) op = curl_e(wavenumber, dxes) / (-1j * omega)
if mu is not None: if mu is not None:
op = sparse.diags(1 / mu) @ op op = sparse.diags_array(1 / mu) @ op
return op return op
def h2e( def h2e(
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t epsilon: vfdslice,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Returns an operator which, when applied to a vectorized H eigenfield, produces Returns an operator which, when applied to a vectorized H eigenfield, produces
the vectorized E eigenfield. the vectorized E eigenfield slice.
Args: Args:
wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
@ -630,13 +628,13 @@ def h2e(
Returns: Returns:
Sparse matrix representation of the operator. 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 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: Args:
wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` 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: Returns:
Sparse matrix representation of the operator. Sparse matrix representation of the operator.
""" """
n = 1 nn = 1
for d in dxes[0]: for dd in dxes[0]:
n *= len(d) nn *= len(dd)
Bz = -1j * wavenumber * sparse.eye(n) Bz = -1j * wavenumber * sparse.eye_array(nn)
Dfx, Dfy = deriv_forward(dxes[0]) Dfx, Dfy = deriv_forward(dxes[0])
return cross([Dfx, Dfy, Bz]) 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: Args:
wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)` 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: Returns:
Sparse matrix representation of the operator. Sparse matrix representation of the operator.
""" """
n = 1 nn = 1
for d in dxes[1]: for dd in dxes[1]:
n *= len(d) nn *= len(dd)
Bz = -1j * wavenumber * sparse.eye(n) Bz = -1j * wavenumber * sparse.eye_array(nn)
Dbx, Dby = deriv_back(dxes[1]) Dbx, Dby = deriv_back(dxes[1])
return cross([Dbx, Dby, Bz]) return cross([Dbx, Dby, Bz])
def h_err( def h_err(
h: vcfdfield_t, h: vcfdslice,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> float: ) -> float:
""" """
Calculates the relative error in the H field Calculates the relative error in the H field
@ -699,7 +697,7 @@ def h_err(
ce = curl_e(wavenumber, dxes) ce = curl_e(wavenumber, dxes)
ch = curl_h(wavenumber, dxes) ch = curl_h(wavenumber, dxes)
eps_inv = sparse.diags(1 / epsilon) eps_inv = sparse.diags_array(1 / epsilon)
if mu is None: if mu is None:
op = ce @ eps_inv @ ch @ h - omega ** 2 * h op = ce @ eps_inv @ ch @ h - omega ** 2 * h
@ -710,12 +708,12 @@ def h_err(
def e_err( def e_err(
e: vcfdfield_t, e: vcfdslice,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
) -> float: ) -> float:
""" """
Calculates the relative error in the E field Calculates the relative error in the E field
@ -737,21 +735,21 @@ def e_err(
if mu is None: if mu is None:
op = ch @ ce @ e - omega ** 2 * (epsilon * e) op = ch @ ce @ e - omega ** 2 * (epsilon * e)
else: else:
mu_inv = sparse.diags(1 / mu) mu_inv = sparse.diags_array(1 / mu)
op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e) op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
return float(norm(op) / norm(e)) return float(norm(op) / norm(e))
def sensitivity( def sensitivity(
e_norm: vcfdfield_t, e_norm: vcfdslice,
h_norm: vcfdfield_t, h_norm: vcfdslice,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
) -> vcfdfield_t: ) -> vcfdslice_t:
r""" r"""
Given a waveguide structure (`dxes`, `epsilon`, `mu`) and mode fields Given a waveguide structure (`dxes`, `epsilon`, `mu`) and mode fields
(`e_norm`, `h_norm`, `wavenumber`, `omega`), calculates the sensitivity of the wavenumber (`e_norm`, `h_norm`, `wavenumber`, `omega`), calculates the sensitivity of the wavenumber
@ -825,11 +823,11 @@ def sensitivity(
Dbx, Dby = deriv_back(dxes[1]) Dbx, Dby = deriv_back(dxes[1])
eps_x, eps_y, eps_z = numpy.split(epsilon, 3) eps_x, eps_y, eps_z = numpy.split(epsilon, 3)
eps_xy = sparse.diags(numpy.hstack((eps_x, eps_y))) eps_xy = sparse.diags_array(numpy.hstack((eps_x, eps_y)))
eps_z_inv = sparse.diags(1 / eps_z) eps_z_inv = sparse.diags_array(1 / eps_z)
mu_x, mu_y, _mu_z = numpy.split(mu, 3) 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_exxhyy = vec(dxes[1][0][:, None] * dxes[0][1][None, :])
da_eyyhxx = vec(dxes[1][1][None, :] * dxes[0][0][:, None]) da_eyyhxx = vec(dxes[1][1][None, :] * dxes[0][0][:, None])
@ -843,15 +841,15 @@ def sensitivity(
norm = hv_yx_conj @ ev_xy norm = hv_yx_conj @ ev_xy
sens_tot = numpy.concatenate([sens_xy1 + sens_xy2, sens_z]) / (2 * wavenumber * norm) sens_tot = numpy.concatenate([sens_xy1 + sens_xy2, sens_z]) / (2 * wavenumber * norm)
return sens_tot return vcfdslice_t(sens_tot)
def solve_modes( def solve_modes(
mode_numbers: Sequence[int], mode_numbers: Sequence[int],
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
mode_margin: int = 2, mode_margin: int = 2,
) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]: ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
""" """
@ -907,7 +905,7 @@ def solve_mode(
mode_number: int, mode_number: int,
*args: Any, *args: Any,
**kwargs: Any, **kwargs: Any,
) -> tuple[vcfdfield_t, complex]: ) -> tuple[vcfdfield2_t, complex]:
""" """
Wrapper around `solve_modes()` that solves for a single mode. Wrapper around `solve_modes()` that solves for a single mode.
@ -921,13 +919,13 @@ def solve_mode(
""" """
kwargs['mode_numbers'] = [mode_number] kwargs['mode_numbers'] = [mode_number]
e_xys, wavenumbers = solve_modes(*args, **kwargs) 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 def inner_product( # TODO documentation
e1: vcfdfield_t, e1: vcfdfield2,
h2: vcfdfield_t, h2: vcfdfield2,
dxes: dx_lists_t, dxes: dx_lists2_t,
prop_phase: float = 0, prop_phase: float = 0,
conj_h: bool = False, conj_h: bool = False,
trapezoid: bool = False, trapezoid: bool = False,

33
meanas/fdfd/waveguide_3d.py
View File

@ -4,13 +4,13 @@ Tools for working with waveguide modes in 3D domains.
This module relies heavily on `waveguide_2d` and mostly just transforms This module relies heavily on `waveguide_2d` and mostly just transforms
its parameters into 2D equivalents and expands the results back into 3D. its parameters into 2D equivalents and expands the results back into 3D.
""" """
from typing import Any from typing import Any, cast
from collections.abc import Sequence from collections.abc import Sequence
import numpy import numpy
from numpy.typing import NDArray from numpy.typing import NDArray
from numpy import complexfloating 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 from . import operators, waveguide_2d
@ -21,8 +21,8 @@ def solve_mode(
axis: int, axis: int,
polarity: int, polarity: int,
slices: Sequence[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> dict[str, complex | NDArray[complexfloating]]: ) -> dict[str, complex | NDArray[complexfloating]]:
""" """
Given a 3D grid, selects a slice from the grid and attempts to Given a 3D grid, selects a slice from the grid and attempts to
@ -95,9 +95,10 @@ def solve_mode(
# Expand E, H to full epsilon space we were given # Expand E, H to full epsilon space we were given
E = numpy.zeros_like(epsilon, dtype=complex) E = numpy.zeros_like(epsilon, dtype=complex)
H = numpy.zeros_like(epsilon, dtype=complex) H = numpy.zeros_like(epsilon, dtype=complex)
for a, o in enumerate(reverse_order): for aa, oo in enumerate(reverse_order):
E[(a, *slices)] = e[o][:, :, None].transpose(reverse_order) iii = cast('tuple[slice | int]', (aa, *slices))
H[(a, *slices)] = h[o][:, :, None].transpose(reverse_order) E[iii] = e[oo][:, :, None].transpose(reverse_order)
H[iii] = h[oo][:, :, None].transpose(reverse_order)
results = { results = {
'wavenumber': wavenumber, 'wavenumber': wavenumber,
@ -109,15 +110,15 @@ def solve_mode(
def compute_source( def compute_source(
E: cfdfield_t, E: cfdfield,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: Sequence[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield,
mu: fdfield_t | None = None, mu: fdfield | None = None,
) -> cfdfield_t: ) -> cfdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, returns the current source distribution Given an eigenmode obtained by `solve_mode`, returns the current source distribution
@ -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)) 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:]) J = unvec(masked_e2j @ vec(E_expanded), E.shape[1:])
return J return cfdfield_t(J)
def compute_overlap_e( def compute_overlap_e(
E: cfdfield_t, E: cfdfield,
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
@ -195,12 +196,12 @@ def compute_overlap_e(
Etgt = numpy.zeros_like(Ee) Etgt = numpy.zeros_like(Ee)
Etgt[slices2] = Ee[slices2] Etgt[slices2] = Ee[slices2]
Etgt /= (Etgt.conj() * Etgt).sum() Etgt /= (Etgt.conj() * Etgt).sum() # type: ignore
return Etgt return cfdfield_t(Etgt)
def expand_e( def expand_e(
E: cfdfield_t, E: cfdfield,
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
@ -245,4 +246,4 @@ def expand_e(
slices_in = (slice(None), *slices) slices_in = (slice(None), *slices)
E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in] E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in]
return E_expanded return cfdfield_t(E_expanded)

79
meanas/fdfd/waveguide_cyl.py
View File

@ -68,7 +68,7 @@ import numpy
from numpy.typing import NDArray, ArrayLike from numpy.typing import NDArray, ArrayLike
from scipy import sparse 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 ..fdmath.operators import deriv_forward, deriv_back
from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
from . import waveguide_2d from . import waveguide_2d
@ -78,10 +78,10 @@ logger = logging.getLogger(__name__)
def cylindrical_operator( def cylindrical_operator(
omega: float, omega: float,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
rmin: float, rmin: float,
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Cylindrical coordinate waveguide operator of the form Cylindrical coordinate waveguide operator of the form
@ -143,11 +143,11 @@ def cylindrical_operator(
def solve_modes( def solve_modes(
mode_numbers: Sequence[int], mode_numbers: Sequence[int],
omega: float, omega: float,
dxes: dx_lists_t, dxes: dx_lists2_t,
epsilon: vfdfield_t, epsilon: vfdslice,
rmin: float, rmin: float,
mode_margin: int = 2, 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 Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
of the bent waveguide with the specified mode number. 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 # Wavenumbers assume the mode is at rmin, which is unlikely
# Instead, return the wavenumber in inverse radians # Instead, return the wavenumber in inverse radians
angular_wavenumbers = wavenumbers * cast(complex, rmin) angular_wavenumbers = wavenumbers * cast('complex', rmin)
order = angular_wavenumbers.argsort()[::-1] order = angular_wavenumbers.argsort()[::-1]
e_xys = e_xys[order] e_xys = e_xys[order]
@ -204,7 +204,7 @@ def solve_mode(
mode_number: int, mode_number: int,
*args: Any, *args: Any,
**kwargs: Any, **kwargs: Any,
) -> tuple[vcfdfield_t, complex]: ) -> tuple[vcfdslice, complex]:
""" """
Wrapper around `solve_modes()` that solves for a single mode. Wrapper around `solve_modes()` that solves for a single mode.
@ -222,10 +222,10 @@ def solve_mode(
def linear_wavenumbers( def linear_wavenumbers(
e_xys: vcfdfield_t, e_xys: list[vcfdfield2_t],
angular_wavenumbers: ArrayLike, angular_wavenumbers: ArrayLike,
epsilon: vfdfield_t, epsilon: vfdslice,
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float,
) -> NDArray[numpy.complex128]: ) -> NDArray[numpy.complex128]:
""" """
@ -247,7 +247,6 @@ def linear_wavenumbers(
angular_wavenumbers = numpy.asarray(angular_wavenumbers) angular_wavenumbers = numpy.asarray(angular_wavenumbers)
mode_radii = numpy.empty_like(angular_wavenumbers, dtype=float) mode_radii = numpy.empty_like(angular_wavenumbers, dtype=float)
wavenumbers = numpy.empty_like(angular_wavenumbers)
shape2d = (len(dxes[0][0]), len(dxes[0][1])) shape2d = (len(dxes[0][0]), len(dxes[0][1]))
epsilon2d = unvec(epsilon, shape2d)[:2] epsilon2d = unvec(epsilon, shape2d)[:2]
grid_radii = rmin + numpy.cumsum(dxes[0][0]) grid_radii = rmin + numpy.cumsum(dxes[0][0])
@ -265,11 +264,11 @@ def linear_wavenumbers(
def exy2h( def exy2h(
angular_wavenumber: complex, angular_wavenumber: complex,
omega: float, omega: float,
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
into a vectorized H containing all three H components into a vectorized H containing all three H components
@ -294,10 +293,10 @@ def exy2h(
def exy2e( def exy2e(
angular_wavenumber: complex, angular_wavenumber: complex,
omega: float, omega: float,
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float,
epsilon: vfdfield_t, epsilon: vfdslice,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
into a vectorized E containing all three E components into a vectorized E containing all three E components
@ -323,7 +322,7 @@ def exy2e(
Ta, Tb = dxes2T(dxes=dxes, rmin=rmin) Ta, Tb = dxes2T(dxes=dxes, rmin=rmin)
Tai = sparse.diags_array(1 / Ta.diagonal()) 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_parts = numpy.split(epsilon, 3)
epsilon_x, epsilon_y = (sparse.diags_array(epsi) for epsi in epsilon_parts[:2]) 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 n_pts = dxes[0][0].size * dxes[0][1].size
zeros = sparse.coo_array((n_pts, n_pts)) 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 = numpy.ones(n_pts)
mu_z_inv = sparse.diags_array(1 / mu_z) mu_z_inv = sparse.diags_array(1 / mu_z)
@ -352,10 +349,10 @@ def exy2e(
def e2h( def e2h(
angular_wavenumber: complex, angular_wavenumber: complex,
omega: float, omega: float,
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float,
mu: vfdfield_t | None = None mu: vfdslice | None = None
) -> sparse.spmatrix: ) -> sparse.sparray:
r""" r"""
Returns an operator which, when applied to a vectorized E eigenfield, produces Returns an operator which, when applied to a vectorized E eigenfield, produces
the vectorized H eigenfield. the vectorized H eigenfield.
@ -396,7 +393,7 @@ def e2h(
def dxes2T( def dxes2T(
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float,
) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64]]: ) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64]]:
r""" r"""
@ -421,15 +418,15 @@ def dxes2T(
def normalized_fields_e( def normalized_fields_e(
e_xy: ArrayLike, e_xy: vcfdfield2,
angular_wavenumber: complex, angular_wavenumber: complex,
omega: float, omega: float,
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float,
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
prop_phase: float = 0, 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, Given a vector `e_xy` containing the vectorized E_x and E_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -459,23 +456,21 @@ def normalized_fields_e(
def _normalized_fields( def _normalized_fields(
e: vcfdfield_t, e: vcfdslice,
h: vcfdfield_t, h: vcfdslice,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists2_t,
rmin: float, rmin: float, # Currently unused, but may want to use cylindrical poynting
epsilon: vfdfield_t, epsilon: vfdslice,
mu: vfdfield_t | None = None, mu: vfdslice | None = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> tuple[vcfdfield_t, vcfdfield_t]: ) -> tuple[vcfdslice_t, vcfdslice_t]:
h *= -1 h *= -1
# TODO documentation for normalized_fields # TODO documentation for normalized_fields
shape = [s.size for s in dxes[0]] shape = [s.size for s in dxes[0]]
dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)]
# Find time-averaged Sz and normalize to it # 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 # 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 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=}' 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 e *= norm_factor
h *= norm_factor h *= norm_factor
return e, h return vcfdslice_t(e), vcfdslice_t(h)

34
meanas/fdmath/__init__.py
View File

@ -742,12 +742,34 @@ normalized results are needed.
""" """
from .types import ( from .types import (
fdfield_t as fdfield_t, fdfield_t as fdfield_t,
vfdfield_t as vfdfield_t, vfdfield_t as vfdfield_t,
cfdfield_t as cfdfield_t, cfdfield_t as cfdfield_t,
vcfdfield_t as vcfdfield_t, vcfdfield_t as vcfdfield_t,
dx_lists_t as dx_lists_t, fdfield2_t as fdfield2_t,
dx_lists_mut as dx_lists_mut, 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, fdfield_updater_t as fdfield_updater_t,
cfdfield_updater_t as cfdfield_updater_t, cfdfield_updater_t as cfdfield_updater_t,
) )

51
meanas/fdmath/operators.py
View File

@ -16,7 +16,7 @@ def shift_circ(
axis: int, axis: int,
shape: Sequence[int], shape: Sequence[int],
shift_distance: int = 1, shift_distance: int = 1,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Utility operator for performing a circular shift along a specified axis by a Utility operator for performing a circular shift along a specified axis by a
specified number of elements. specified number of elements.
@ -44,7 +44,7 @@ def shift_circ(
vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C'))) 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: if shift_distance < 0:
d = d.T d = d.T
@ -56,7 +56,7 @@ def shift_with_mirror(
axis: int, axis: int,
shape: Sequence[int], shape: Sequence[int],
shift_distance: int = 1, shift_distance: int = 1,
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Utility operator for performing an n-element shift along a specified axis, with mirror Utility operator for performing an n-element shift along a specified axis, with mirror
boundary conditions applied to the cells beyond the receding edge. 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'))) 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 return d
def deriv_forward( def deriv_forward(
dx_e: Sequence[NDArray[floating | complexfloating]], dx_e: Sequence[NDArray[floating | complexfloating]],
) -> list[sparse.spmatrix]: ) -> list[sparse.sparray]:
""" """
Utility operators for taking discretized derivatives (forward variant). Utility operators for taking discretized derivatives (forward variant).
@ -114,10 +114,10 @@ def deriv_forward(
dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij') dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij')
def deriv(axis: int) -> sparse.spmatrix: def deriv(axis: int) -> sparse.sparray:
return shift_circ(axis, shape, 1) - sparse.eye(n) 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)] for a, dx in enumerate(dx_e_expanded)]
return Ds return Ds
@ -125,7 +125,7 @@ def deriv_forward(
def deriv_back( def deriv_back(
dx_h: Sequence[NDArray[floating | complexfloating]], dx_h: Sequence[NDArray[floating | complexfloating]],
) -> list[sparse.spmatrix]: ) -> list[sparse.sparray]:
""" """
Utility operators for taking discretized derivatives (backward variant). Utility operators for taking discretized derivatives (backward variant).
@ -141,18 +141,18 @@ def deriv_back(
dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij') dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij')
def deriv(axis: int) -> sparse.spmatrix: def deriv(axis: int) -> sparse.sparray:
return shift_circ(axis, shape, -1) - sparse.eye(n) 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)] for a, dx in enumerate(dx_h_expanded)]
return Ds return Ds
def cross( def cross(
B: Sequence[sparse.spmatrix], B: Sequence[sparse.sparray],
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Cross product operator Cross product operator
@ -164,13 +164,14 @@ def cross(
Sparse matrix corresponding to (B x), where x is the cross product. Sparse matrix corresponding to (B x), where x is the cross product.
""" """
n = B[0].shape[0] n = B[0].shape[0]
zero = sparse.csr_matrix((n, n)) zero = sparse.csr_array((n, n))
return sparse.bmat([[zero, -B[2], B[1]], return sparse.block_array([
[B[2], zero, -B[0]], [zero, -B[2], B[1]],
[-B[1], B[0], zero]]) [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 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. 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) 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)` 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}') raise Exception(f'Invalid shape: {shape}')
n = numpy.prod(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)` Backward average operator `(x4 = (x4 + x3) / 2)`
@ -220,7 +221,7 @@ def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
def curl_forward( def curl_forward(
dx_e: Sequence[NDArray[floating | complexfloating]], dx_e: Sequence[NDArray[floating | complexfloating]],
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Curl operator for use with the E field. Curl operator for use with the E field.
@ -236,7 +237,7 @@ def curl_forward(
def curl_back( def curl_back(
dx_h: Sequence[NDArray[floating | complexfloating]], dx_h: Sequence[NDArray[floating | complexfloating]],
) -> sparse.spmatrix: ) -> sparse.sparray:
""" """
Curl operator for use with the H field. Curl operator for use with the H field.

61
meanas/fdmath/types.py
View File

@ -1,25 +1,64 @@
""" """
Types shared across multiple submodules Types shared across multiple submodules
""" """
from typing import NewType
from collections.abc import Sequence, Callable, MutableSequence from collections.abc import Sequence, Callable, MutableSequence
from numpy.typing import NDArray from numpy.typing import NDArray
from numpy import floating, complexfloating from numpy import floating, complexfloating
# Field types # 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]`)""" """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)""" """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]`)""" """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)""" """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]]] dx_lists_t = Sequence[Sequence[NDArray[floating | complexfloating]]]
""" """
'dxes' datastructure which contains grid cell width information in the following format: '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. 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]]] dx_lists_mut = MutableSequence[MutableSequence[NDArray[floating | complexfloating]]]
"""Mutable version of `dx_lists_t`""" """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] fdfield_updater_t = Callable[..., fdfield_t]
"""Convenience type for functions which take and return an fdfield_t""" """Convenience type for functions which take and return an fdfield_t"""

58
meanas/fdmath/vectorization.py
View File

@ -7,9 +7,13 @@ Vectorized versions of the field use row-major (ie., C-style) ordering.
from typing import overload from typing import overload
from collections.abc import Sequence from collections.abc import Sequence
import numpy 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 @overload
@ -25,12 +29,28 @@ def vec(f: cfdfield_t) -> vcfdfield_t:
pass pass
@overload @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 pass
def vec( def vec(
f: fdfield_t | cfdfield_t | ArrayLike | None, f: fdfield_t | cfdfield_t | fdfield2_t | cfdfield2_t | fdslice_t | cfdslice_t | ArrayLike | None,
) -> vfdfield_t | vcfdfield_t | 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. Create a 1D ndarray from a vector field which spans a 1-3D region.
@ -45,7 +65,7 @@ def vec(
""" """
if f is None: if f is None:
return None return None
return numpy.ravel(f, order='C') return numpy.ravel(f, order='C') # type: ignore
@overload @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: def unvec(v: vcfdfield_t, shape: Sequence[int], nvdim: int = 3) -> cfdfield_t:
pass 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( 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], shape: Sequence[int],
nvdim: int = 3, 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 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 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: if v is None:
return None return None
return v.reshape((nvdim, *shape), order='C') return v.reshape((nvdim, *shape), order='C') # type: ignore

84
meanas/fdtd/energy.py
View File

@ -1,6 +1,6 @@
import numpy 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 from ..fdmath.functional import deriv_back
@ -8,8 +8,8 @@ from ..fdmath.functional import deriv_back
def poynting( def poynting(
e: fdfield_t, e: fdfield,
h: fdfield_t, h: fdfield,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
r""" r"""
@ -84,14 +84,14 @@ def poynting(
s[0] = numpy.roll(ey, -1, axis=0) * hz - numpy.roll(ez, -1, axis=0) * hy 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[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 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( def poynting_divergence(
s: fdfield_t | None = None, s: fdfield | None = None,
*, *,
e: fdfield_t | None = None, e: fdfield | None = None,
h: fdfield_t | None = None, h: fdfield | None = None,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
@ -122,15 +122,15 @@ def poynting_divergence(
Dx, Dy, Dz = deriv_back() Dx, Dy, Dz = deriv_back()
ds = Dx(s[0]) + Dy(s[1]) + Dz(s[2]) ds = Dx(s[0]) + Dy(s[1]) + Dz(s[2])
return ds return fdfield_t(ds)
def energy_hstep( def energy_hstep(
e0: fdfield_t, e0: fdfield,
h1: fdfield_t, h1: fdfield,
e2: fdfield_t, e2: fdfield,
epsilon: fdfield_t | None = None, epsilon: fdfield | None = None,
mu: fdfield_t | None = None, mu: fdfield | None = None,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
@ -150,15 +150,15 @@ def energy_hstep(
Energy, at the time of the H-field `h1`. Energy, at the time of the H-field `h1`.
""" """
u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes) u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes)
return u return fdfield_t(u)
def energy_estep( def energy_estep(
h0: fdfield_t, h0: fdfield,
e1: fdfield_t, e1: fdfield,
h2: fdfield_t, h2: fdfield,
epsilon: fdfield_t | None = None, epsilon: fdfield | None = None,
mu: fdfield_t | None = None, mu: fdfield | None = None,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
@ -178,17 +178,17 @@ def energy_estep(
Energy, at the time of the E-field `e1`. Energy, at the time of the E-field `e1`.
""" """
u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes) u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes)
return u return fdfield_t(u)
def delta_energy_h2e( def delta_energy_h2e(
dt: float, dt: float,
e0: fdfield_t, e0: fdfield,
h1: fdfield_t, h1: fdfield,
e2: fdfield_t, e2: fdfield,
h3: fdfield_t, h3: fdfield,
epsilon: fdfield_t | None = None, epsilon: fdfield | None = None,
mu: fdfield_t | None = None, mu: fdfield | None = None,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
@ -211,17 +211,17 @@ def delta_energy_h2e(
de = e2 * (e2 - e0) / dt de = e2 * (e2 - e0) / dt
dh = h1 * (h3 - h1) / dt dh = h1 * (h3 - h1) / dt
du = dxmul(de, dh, epsilon, mu, dxes) du = dxmul(de, dh, epsilon, mu, dxes)
return du return fdfield_t(du)
def delta_energy_e2h( def delta_energy_e2h(
dt: float, dt: float,
h0: fdfield_t, h0: fdfield,
e1: fdfield_t, e1: fdfield,
h2: fdfield_t, h2: fdfield,
e3: fdfield_t, e3: fdfield,
epsilon: fdfield_t | None = None, epsilon: fdfield | None = None,
mu: fdfield_t | None = None, mu: fdfield | None = None,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
@ -244,12 +244,12 @@ def delta_energy_e2h(
de = e1 * (e3 - e1) / dt de = e1 * (e3 - e1) / dt
dh = h2 * (h2 - h0) / dt dh = h2 * (h2 - h0) / dt
du = dxmul(de, dh, epsilon, mu, dxes) du = dxmul(de, dh, epsilon, mu, dxes)
return du return fdfield_t(du)
def delta_energy_j( def delta_energy_j(
j0: fdfield_t, j0: fdfield,
e1: fdfield_t, e1: fdfield,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
@ -267,14 +267,14 @@ def delta_energy_j(
* dxes[0][0][:, None, None] * dxes[0][0][:, None, None]
* dxes[0][1][None, :, None] * dxes[0][1][None, :, None]
* dxes[0][2][None, None, :]) * dxes[0][2][None, None, :])
return du return fdfield_t(du)
def dxmul( def dxmul(
ee: fdfield_t, ee: fdfield,
hh: fdfield_t, hh: fdfield,
epsilon: fdfield_t | float | None = None, epsilon: fdfield | float | None = None,
mu: fdfield_t | float | None = None, mu: fdfield | float | None = None,
dxes: dx_lists_t | None = None, dxes: dx_lists_t | None = None,
) -> fdfield_t: ) -> fdfield_t:
if epsilon is None: if epsilon is None:
@ -292,4 +292,4 @@ def dxmul(
* dxes[1][0][:, None, None] * dxes[1][0][:, None, None]
* dxes[1][1][None, :, None] * dxes[1][1][None, :, None]
* dxes[1][2][None, None, :]) * dxes[1][2][None, None, :])
return result return fdfield_t(result)

7
meanas/fdtd/misc.py
View File

@ -1,5 +1,4 @@
from typing import Callable from collections.abc import Callable
from collections.abc import Sequence
import logging import logging
import numpy import numpy
@ -95,10 +94,10 @@ def ricker_pulse(
logger.warning('meanas.fdtd.misc functions are still very WIP!') # TODO logger.warning('meanas.fdtd.misc functions are still very WIP!') # TODO
omega = 2 * pi / wl omega = 2 * pi / wl
freq = 1 / wl freq = 1 / wl
r0 = omega / 2 # r0 = omega / 2
from scipy.optimize import root_scalar 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 = delay_results.root
delay = numpy.ceil(delay * freq) / freq # force delay to integer number of periods to maintain phase delay = numpy.ceil(delay * freq) / freq # force delay to integer number of periods to maintain phase

4
meanas/fdtd/pml.py
View File

@ -13,7 +13,7 @@ from copy import deepcopy
import numpy import numpy
from numpy.typing import NDArray, DTypeLike 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 from ..fdmath.functional import deriv_forward, deriv_back
@ -97,7 +97,7 @@ def updates_with_cpml(
cpml_params: Sequence[Sequence[dict[str, Any] | None]], cpml_params: Sequence[Sequence[dict[str, Any] | None]],
dt: float, dt: float,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield,
*, *,
dtype: DTypeLike = numpy.float32, dtype: DTypeLike = numpy.float32,
) -> tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None], ) -> tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None],

26
meanas/test/test_fdfd.py
View File

@ -6,7 +6,7 @@ from numpy.typing import NDArray
#from numpy.testing import assert_allclose, assert_array_equal #from numpy.testing import assert_allclose, assert_array_equal
from .. import fdfd from .. import fdfd
from ..fdmath import vec, unvec from ..fdmath import vec, unvec, vcfdfield, vfdfield, dx_lists_t
from .utils import assert_close # , assert_fields_close from .utils import assert_close # , assert_fields_close
from .conftest import FixtureRequest from .conftest import FixtureRequest
@ -102,16 +102,16 @@ def j_distribution(
@dataclasses.dataclass() @dataclasses.dataclass()
class FDResult: class FDResult:
shape: tuple[int, ...] shape: tuple[int, ...]
dxes: list[list[NDArray[numpy.float64]]] dxes: dx_lists_t
epsilon: NDArray[numpy.float64] epsilon: vfdfield
omega: complex omega: complex
j: NDArray[numpy.complex128] j: vcfdfield
e: NDArray[numpy.complex128] e: vcfdfield
pmc: NDArray[numpy.float64] | None pmc: vfdfield | None
pec: NDArray[numpy.float64] | None pec: vfdfield | None
@pytest.fixture() @pytest.fixture
def sim( def sim(
request: FixtureRequest, request: FixtureRequest,
shape: tuple[int, ...], shape: tuple[int, ...],
@ -141,11 +141,11 @@ def sim(
j_vec = vec(j_distribution) j_vec = vec(j_distribution)
eps_vec = vec(epsilon) eps_vec = vec(epsilon)
e_vec = fdfd.solvers.generic( e_vec = fdfd.solvers.generic(
J=j_vec, J = j_vec,
omega=omega, omega = omega,
dxes=dxes, dxes = dxes,
epsilon=eps_vec, epsilon = eps_vec,
matrix_solver_opts={'atol': 1e-15, 'rtol': 1e-11}, matrix_solver_opts = dict(atol=1e-15, rtol=1e-11),
) )
e = unvec(e_vec, shape[1:]) e = unvec(e_vec, shape[1:])

12
meanas/test/test_fdfd_pml.py
View File

@ -5,7 +5,7 @@ from numpy.typing import NDArray
from numpy.testing import assert_allclose from numpy.testing import assert_allclose
from .. import fdfd 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 .utils import assert_close, assert_fields_close
from .test_fdfd import FDResult from .test_fdfd import FDResult
from .conftest import FixtureRequest from .conftest import FixtureRequest
@ -70,15 +70,15 @@ def src_polarity(request: FixtureRequest) -> int:
return request.param return request.param
@pytest.fixture() @pytest.fixture
def j_distribution( def j_distribution(
request: FixtureRequest, request: FixtureRequest,
shape: tuple[int, ...], shape: tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: vfdfield,
dxes: dx_lists_mut, dxes: dx_lists_mut,
omega: float, omega: float,
src_polarity: int, src_polarity: int,
) -> NDArray[numpy.complex128]: ) -> cfdfield_t:
j = numpy.zeros(shape, dtype=complex) j = numpy.zeros(shape, dtype=complex)
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
@ -109,7 +109,7 @@ def j_distribution(
return j return j
@pytest.fixture() @pytest.fixture
def epsilon( def epsilon(
request: FixtureRequest, request: FixtureRequest,
shape: tuple[int, ...], shape: tuple[int, ...],
@ -144,7 +144,7 @@ def dxes(
return dxes return dxes
@pytest.fixture() @pytest.fixture
def sim( def sim(
request: FixtureRequest, request: FixtureRequest,
shape: tuple[int, ...], shape: tuple[int, ...],

2
meanas/test/test_fdtd.py
View File

@ -189,7 +189,7 @@ def j_distribution(
return j return j
@pytest.fixture() @pytest.fixture
def sim( def sim(
request: FixtureRequest, request: FixtureRequest,
shape: tuple[int, ...], shape: tuple[int, ...],