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