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

49
meanas/fdfd/bloch.py
View File

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

44
meanas/fdfd/eme.py
View File

@ -1,20 +1,29 @@
from collections.abc import Sequence
import numpy
from numpy.typing import NDArray
from scipy import sparse
from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t
from ..fdmath import dx_lists2_t, vcfdfield2
from .waveguide_2d import inner_product
def get_tr(ehL, wavenumbers_L, ehR, wavenumbers_R, dxes: dx_lists_t):
def get_tr(
ehLs: Sequence[Sequence[vcfdfield2]],
wavenumbers_L: Sequence[complex],
ehRs: Sequence[Sequence[vcfdfield2]],
wavenumbers_R: Sequence[complex],
dxes: dx_lists2_t,
) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
nL = len(wavenumbers_L)
nR = len(wavenumbers_R)
A12 = numpy.zeros((nL, nR), dtype=complex)
A21 = numpy.zeros((nL, nR), dtype=complex)
B11 = numpy.zeros((nL,), dtype=complex)
for ll in range(nL):
eL, hL = ehL[ll]
eL, hL = ehLs[ll]
B11[ll] = inner_product(eL, hL, dxes=dxes, conj_h=False)
for rr in range(nR):
eR, hR = ehR[rr]
eR, hR = ehRs[rr]
A12[ll, rr] = inner_product(eL, hR, dxes=dxes, conj_h=False) # TODO optimize loop?
A21[ll, rr] = inner_product(eR, hL, dxes=dxes, conj_h=False)
@ -32,9 +41,15 @@ def get_tr(ehL, wavenumbers_L, ehR, wavenumbers_R, dxes: dx_lists_t):
return tt, rr
def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs):
t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs)
t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs)
def get_abcd(
ehLs: Sequence[Sequence[vcfdfield2]],
wavenumbers_L: Sequence[complex],
ehRs: Sequence[Sequence[vcfdfield2]],
wavenumbers_R: Sequence[complex],
**kwargs,
) -> sparse.sparray:
t12, r12 = get_tr(ehLs, wavenumbers_L, ehRs, wavenumbers_R, **kwargs)
t21, r21 = get_tr(ehRs, wavenumbers_R, ehLs, wavenumbers_L, **kwargs)
t21i = numpy.linalg.pinv(t21)
A = t12 - r21 @ t21i @ r12
B = r21 @ t21i
@ -44,15 +59,16 @@ def get_abcd(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs):
def get_s(
eL_xys,
wavenumbers_L,
eR_xys,
wavenumbers_R,
ehLs: Sequence[Sequence[vcfdfield2]],
wavenumbers_L: Sequence[complex],
ehRs: Sequence[Sequence[vcfdfield2]],
wavenumbers_R: Sequence[complex],
force_nogain: bool = False,
force_reciprocal: bool = False,
**kwargs):
t12, r12 = get_tr(eL_xys, wavenumbers_L, eR_xys, wavenumbers_R, **kwargs)
t21, r21 = get_tr(eR_xys, wavenumbers_R, eL_xys, wavenumbers_L, **kwargs)
**kwargs,
) -> NDArray[numpy.complex128]:
t12, r12 = get_tr(ehLs, wavenumbers_L, ehRs, wavenumbers_R, **kwargs)
t21, r21 = get_tr(ehRs, wavenumbers_R, ehLs, wavenumbers_L, **kwargs)
ss = numpy.block([[r12, t12],
[t21, r21]])

10
meanas/fdfd/farfield.py
View File

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

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

142
meanas/fdfd/operators.py
View File

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

20
meanas/fdfd/solvers.py
View File

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

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

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

79
meanas/fdfd/waveguide_cyl.py
View File

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

34
meanas/fdmath/__init__.py
View File

@ -742,12 +742,34 @@ normalized results are needed.
"""
from .types import (
fdfield_t as fdfield_t,
vfdfield_t as vfdfield_t,
cfdfield_t as cfdfield_t,
vcfdfield_t as vcfdfield_t,
dx_lists_t as dx_lists_t,
dx_lists_mut as dx_lists_mut,
fdfield_t as fdfield_t,
vfdfield_t as vfdfield_t,
cfdfield_t as cfdfield_t,
vcfdfield_t as vcfdfield_t,
fdfield2_t as fdfield2_t,
vfdfield2_t as vfdfield2_t,
cfdfield2_t as cfdfield2_t,
vcfdfield2_t as vcfdfield2_t,
fdfield as fdfield,
vfdfield as vfdfield,
cfdfield as cfdfield,
vcfdfield as vcfdfield,
fdfield2 as fdfield2,
vfdfield2 as vfdfield2,
cfdfield2 as cfdfield2,
vcfdfield2 as vcfdfield2,
fdslice_t as fdslice_t,
vfdslice_t as vfdslice_t,
cfdslice_t as cfdslice_t,
vcfdslice_t as vcfdslice_t,
fdslice as fdslice,
vfdslice as vfdslice,
cfdslice as cfdslice,
vcfdslice as vcfdslice,
dx_lists_t as dx_lists_t,
dx_lists2_t as dx_lists2_t,
dx_lists_mut as dx_lists_mut,
dx_lists2_mut as dx_lists2_mut,
fdfield_updater_t as fdfield_updater_t,
cfdfield_updater_t as cfdfield_updater_t,
)

51
meanas/fdmath/operators.py
View File

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

61
meanas/fdmath/types.py
View File

@ -1,25 +1,64 @@
"""
Types shared across multiple submodules
"""
from typing import NewType
from collections.abc import Sequence, Callable, MutableSequence
from numpy.typing import NDArray
from numpy import floating, complexfloating
# Field types
fdfield_t = NDArray[floating]
fdfield_t = NewType('fdfield_t', NDArray[floating])
type fdfield = fdfield_t | NDArray[floating]
"""Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
vfdfield_t = NDArray[floating]
vfdfield_t = NewType('vfdfield_t', NDArray[floating])
type vfdfield = vfdfield_t | NDArray[floating]
"""Linearized vector field (single vector of length 3*X*Y*Z)"""
cfdfield_t = NDArray[complexfloating]
cfdfield_t = NewType('cfdfield_t', NDArray[complexfloating])
type cfdfield = cfdfield_t | NDArray[complexfloating]
"""Complex vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
vcfdfield_t = NDArray[complexfloating]
vcfdfield_t = NewType('vcfdfield_t', NDArray[complexfloating])
type vcfdfield = vcfdfield_t | NDArray[complexfloating]
"""Linearized complex vector field (single vector of length 3*X*Y*Z)"""
fdslice_t = NewType('fdslice_t', NDArray[floating])
type fdslice = fdslice_t | NDArray[floating]
"""Vector field slice with shape (3, X, Y) (e.g. `[E_x, E_y, E_z]` at a single Z position)"""
vfdslice_t = NewType('vfdslice_t', NDArray[floating])
type vfdslice = vfdslice_t | NDArray[floating]
"""Linearized vector field slice (single vector of length 3*X*Y)"""
cfdslice_t = NewType('cfdslice_t', NDArray[complexfloating])
type cfdslice = cfdslice_t | NDArray[complexfloating]
"""Complex vector field slice with shape (3, X, Y) (e.g. `[E_x, E_y, E_z]` at a single Z position)"""
vcfdslice_t = NewType('vcfdslice_t', NDArray[complexfloating])
type vcfdslice = vcfdslice_t | NDArray[complexfloating]
"""Linearized complex vector field slice (single vector of length 3*X*Y)"""
fdfield2_t = NewType('fdfield2_t', NDArray[floating])
type fdfield2 = fdfield2_t | NDArray[floating]
"""2D Vector field with shape (2, X, Y) (e.g. `[E_x, E_y]`)"""
vfdfield2_t = NewType('vfdfield2_t', NDArray[floating])
type vfdfield2 = vfdfield2_t | NDArray[floating]
"""2D Linearized vector field (single vector of length 2*X*Y)"""
cfdfield2_t = NewType('cfdfield2_t', NDArray[complexfloating])
type cfdfield2 = cfdfield2_t | NDArray[complexfloating]
"""2D Complex vector field with shape (2, X, Y) (e.g. `[E_x, E_y]`)"""
vcfdfield2_t = NewType('vcfdfield2_t', NDArray[complexfloating])
type vcfdfield2 = vcfdfield2_t | NDArray[complexfloating]
"""2D Linearized complex vector field (single vector of length 2*X*Y)"""
dx_lists_t = Sequence[Sequence[NDArray[floating | complexfloating]]]
"""
'dxes' datastructure which contains grid cell width information in the following format:
@ -31,9 +70,23 @@ dx_lists_t = Sequence[Sequence[NDArray[floating | complexfloating]]]
and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
"""
dx_lists2_t = Sequence[Sequence[NDArray[floating | complexfloating]]]
"""
2D 'dxes' datastructure which contains grid cell width information in the following format:
[[[dx_e[0], dx_e[1], ...], [dy_e[0], ...]],
[[dx_h[0], dx_h[1], ...], [dy_h[0], ...]]]
where `dx_e[0]` is the x-width of the `x=0` cells, as used when calculating dE/dx,
and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
"""
dx_lists_mut = MutableSequence[MutableSequence[NDArray[floating | complexfloating]]]
"""Mutable version of `dx_lists_t`"""
dx_lists2_mut = MutableSequence[MutableSequence[NDArray[floating | complexfloating]]]
"""Mutable version of `dx_lists2_t`"""
fdfield_updater_t = Callable[..., fdfield_t]
"""Convenience type for functions which take and return an fdfield_t"""

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 collections.abc import Sequence
import numpy
from numpy.typing import ArrayLike
from numpy.typing import ArrayLike, NDArray
from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t
from .types import (
fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t,
fdslice_t, vfdslice_t, cfdslice_t, vcfdslice_t,
fdfield2_t, vfdfield2_t, cfdfield2_t, vcfdfield2_t,
)
@overload
@ -25,12 +29,28 @@ def vec(f: cfdfield_t) -> vcfdfield_t:
pass
@overload
def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t:
def vec(f: fdfield2_t) -> vfdfield2_t:
pass
@overload
def vec(f: cfdfield2_t) -> vcfdfield2_t:
pass
@overload
def vec(f: fdslice_t) -> vfdslice_t:
pass
@overload
def vec(f: cfdslice_t) -> vcfdslice_t:
pass
@overload
def vec(f: ArrayLike) -> NDArray:
pass
def vec(
f: fdfield_t | cfdfield_t | ArrayLike | None,
) -> vfdfield_t | vcfdfield_t | None:
f: fdfield_t | cfdfield_t | fdfield2_t | cfdfield2_t | fdslice_t | cfdslice_t | ArrayLike | None,
) -> vfdfield_t | vcfdfield_t | vfdfield2_t | vcfdfield2_t | vfdslice_t | vcfdslice_t | NDArray | None:
"""
Create a 1D ndarray from a vector field which spans a 1-3D region.
@ -45,7 +65,7 @@ def vec(
"""
if f is None:
return None
return numpy.ravel(f, order='C')
return numpy.ravel(f, order='C') # type: ignore
@overload
@ -60,11 +80,31 @@ def unvec(v: vfdfield_t, shape: Sequence[int], nvdim: int = 3) -> fdfield_t:
def unvec(v: vcfdfield_t, shape: Sequence[int], nvdim: int = 3) -> cfdfield_t:
pass
@overload
def unvec(v: vfdfield2_t, shape: Sequence[int], nvdim: int = 3) -> fdfield2_t:
pass
@overload
def unvec(v: vcfdfield2_t, shape: Sequence[int], nvdim: int = 3) -> cfdfield2_t:
pass
@overload
def unvec(v: vfdslice_t, shape: Sequence[int], nvdim: int = 3) -> fdslice_t:
pass
@overload
def unvec(v: vcfdslice_t, shape: Sequence[int], nvdim: int = 3) -> cfdslice_t:
pass
@overload
def unvec(v: ArrayLike, shape: Sequence[int], nvdim: int = 3) -> NDArray:
pass
def unvec(
v: vfdfield_t | vcfdfield_t | None,
v: vfdfield_t | vcfdfield_t | vfdfield2_t | vcfdfield2_t | vfdslice_t | vcfdslice_t | ArrayLike | None,
shape: Sequence[int],
nvdim: int = 3,
) -> fdfield_t | cfdfield_t | None:
) -> fdfield_t | cfdfield_t | fdfield2_t | cfdfield2_t | fdslice_t | cfdslice_t | NDArray | None:
"""
Perform the inverse of vec(): take a 1D ndarray and output an `nvdim`-component field
of form e.g. `[f_x, f_y, f_z]` (`nvdim=3`) where each of `f_*` is a len(shape)-dimensional
@ -82,5 +122,5 @@ def unvec(
"""
if v is None:
return None
return v.reshape((nvdim, *shape), order='C')
return v.reshape((nvdim, *shape), order='C') # type: ignore

84
meanas/fdtd/energy.py
View File

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

7
meanas/fdtd/misc.py
View File

@ -1,5 +1,4 @@
from typing import Callable
from collections.abc import Sequence
from collections.abc import Callable
import logging
import numpy
@ -95,10 +94,10 @@ def ricker_pulse(
logger.warning('meanas.fdtd.misc functions are still very WIP!') # TODO
omega = 2 * pi / wl
freq = 1 / wl
r0 = omega / 2
# r0 = omega / 2
from scipy.optimize import root_scalar
delay_results = root_scalar(lambda xx: (1 - omega * omega * tt * tt / 2) * numpy.exp(-omega * omega / 4 * tt * tt) - turn_on, x0=0, x1=-2 / omega)
delay_results = root_scalar(lambda tt: (1 - omega * omega * tt * tt / 2) * numpy.exp(-omega * omega / 4 * tt * tt) - turn_on, x0=0, x1=-2 / omega)
delay = delay_results.root
delay = numpy.ceil(delay * freq) / freq # force delay to integer number of periods to maintain phase

4
meanas/fdtd/pml.py
View File

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

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

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

2
meanas/test/test_fdtd.py
View File

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