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

Major typing updates

Browse Source 
Split field types to differentiate between complex and purely-real

Fix lots of numpy-related stuff
This commit is contained in:
Jan Petykiewicz 2022-10-04 17:17:44 -07:00
parent 31e6e0ec60
commit 52df24ad98
18 changed files with 191 additions and 151 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
meanas/eigensolvers.py
View File

@ -11,9 +11,9 @@ import scipy.sparse.linalg as spalg # type: ignore
def power_iteration( def power_iteration(
operator: sparse.spmatrix, operator: sparse.spmatrix,
guess_vector: Optional[NDArray[numpy.float64]] = None, guess_vector: Optional[NDArray[numpy.complex128]] = None,
iterations: int = 20, iterations: int = 20,
) -> Tuple[complex, NDArray[numpy.float64]]: ) -> Tuple[complex, NDArray[numpy.complex128]]:
""" """
Use power iteration to estimate the dominant eigenvector of a matrix. Use power iteration to estimate the dominant eigenvector of a matrix.
@ -26,7 +26,7 @@ def power_iteration(
(Largest-magnitude eigenvalue, Corresponding eigenvector estimate) (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
""" """
if numpy.any(numpy.equal(guess_vector, None)): 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: else:
v = guess_vector v = guess_vector
@ -41,11 +41,11 @@ def power_iteration(
def rayleigh_quotient_iteration( def rayleigh_quotient_iteration(
operator: Union[sparse.spmatrix, spalg.LinearOperator], operator: Union[sparse.spmatrix, spalg.LinearOperator],
guess_vector: NDArray[numpy.float64], guess_vector: NDArray[numpy.complex128],
iterations: int = 40, iterations: int = 40,
tolerance: float = 1e-13, tolerance: float = 1e-13,
solver: Optional[Callable[..., NDArray[numpy.float64]]] = None, solver: Optional[Callable[..., NDArray[numpy.complex128]]] = None,
) -> Tuple[complex, NDArray[numpy.float64]]: ) -> Tuple[complex, NDArray[numpy.complex128]]:
""" """
Use Rayleigh quotient iteration to refine an eigenvector guess. Use Rayleigh quotient iteration to refine an eigenvector guess.
@ -78,7 +78,7 @@ def rayleigh_quotient_iteration(
matvec=lambda v: eigval * v, matvec=lambda v: eigval * v,
) )
if solver is None: 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] return spalg.bicgstab(A, b)[0]
assert(solver is not None) assert(solver is not None)
@ -99,7 +99,7 @@ def signed_eigensolve(
operator: Union[sparse.spmatrix, spalg.LinearOperator], operator: Union[sparse.spmatrix, spalg.LinearOperator],
how_many: int, how_many: int,
negative: bool = False, 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 Find the largest-magnitude positive-only (or negative-only) eigenvalues and
eigenvectors of the provided matrix. eigenvectors of the provided matrix.

50
meanas/fdfd/bloch.py
View File

@ -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 logging
import numpy import numpy
from numpy import pi, real, trace from numpy import pi, real, trace
@ -91,7 +91,8 @@ import scipy.optimize # type: ignore
from scipy.linalg import norm # type: ignore from scipy.linalg import norm # type: ignore
import scipy.sparse.linalg as spalg # 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__) logger = logging.getLogger(__name__)
@ -110,10 +111,10 @@ try:
'planner_effort': 'FFTW_PATIENT', '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) 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) return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)
except ImportError: except ImportError:
@ -124,7 +125,7 @@ except ImportError:
def generate_kmn( def generate_kmn(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
shape: ArrayLike, shape: Sequence[int],
) -> Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]: ) -> 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. 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) k0 = numpy.array(k0)
Gi_grids = numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij') 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_G = k0[None, None, None, :] - Gi
k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2) k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2)
@ -169,7 +170,7 @@ def maxwell_operator(
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None mu: Optional[fdfield_t] = None
) -> Callable[[NDArray[numpy.float64]], NDArray[numpy.float64]]: ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
""" """
Generate the Maxwell operator Generate the Maxwell operator
@ -198,11 +199,11 @@ def maxwell_operator(
shape = epsilon[0].shape + (1,) shape = epsilon[0].shape + (1,)
k_mag, m, n = generate_kmn(k0, G_matrix, shape) 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: 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. Maxwell operator for Bloch eigenmode simulation.
@ -251,7 +252,7 @@ def hmn_2_exyz(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, 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 Generate an operator which converts a vectorized spatial-frequency-space
`h_mn` into an E-field distribution, i.e. `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` Function for converting `h_mn` into `E_xyz`
""" """
shape = epsilon[0].shape + (1,) 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) 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)] hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
d_xyz = (n * hin_m d_xyz = (n * hin_m
- m * hin_n) * k_mag - m * hin_n) * k_mag
@ -291,7 +292,7 @@ def hmn_2_hxyz(
k0: ArrayLike, k0: ArrayLike,
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t 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 Generate an operator which converts a vectorized spatial-frequency-space
`h_mn` into an H-field distribution, i.e. `h_mn` into an H-field distribution, i.e.
@ -314,7 +315,7 @@ def hmn_2_hxyz(
shape = epsilon[0].shape + (1,) shape = epsilon[0].shape + (1,)
_k_mag, m, n = generate_kmn(k0, G_matrix, shape) _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)] hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
h_xyz = (m * hin_m h_xyz = (m * hin_m
+ n * hin_n) + n * hin_n)
@ -328,7 +329,7 @@ def inverse_maxwell_operator_approx(
G_matrix: ArrayLike, G_matrix: ArrayLike,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, 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, 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`. Function which applies the approximate inverse of the maxwell operator to `h_mn`.
""" """
shape = epsilon[0].shape + (1,) 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) k_mag, m, n = generate_kmn(k0, G_matrix, shape)
if mu is not None: def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
mu = numpy.stack(mu, axis=3)
def operator(h: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
""" """
Approximate inverse Maxwell operator for Bloch eigenmode simulation. Approximate inverse Maxwell operator for Bloch eigenmode simulation.
@ -462,7 +462,7 @@ def eigsolve(
tolerance: float = 1e-20, tolerance: float = 1e-20,
max_iters: int = 10000, max_iters: int = 10000,
reset_iters: int = 100, 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 Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
k0 of the specified structure. 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( def _rtrace_AtB(
A: NDArray[numpy.float64], A: NDArray[numpy.complex128],
B: Union[NDArray[numpy.float64], float], B: Union[NDArray[numpy.complex128], float],
) -> float: ) -> float:
return real(numpy.sum(A.conj() * B)) 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 return (A + A.conj().T) * 0.5

10
meanas/fdfd/farfield.py
View File

@ -6,12 +6,12 @@ import numpy
from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
from numpy import pi from numpy import pi
from ..fdmath import fdfield_t from ..fdmath import cfdfield_t
def near_to_farfield( def near_to_farfield(
E_near: fdfield_t, E_near: cfdfield_t,
H_near: fdfield_t, H_near: cfdfield_t,
dx: float, dx: float,
dy: float, dy: float,
padded_size: List[int] = None padded_size: List[int] = None
@ -122,8 +122,8 @@ def near_to_farfield(
def far_to_nearfield( def far_to_nearfield(
E_far: fdfield_t, E_far: cfdfield_t,
H_far: fdfield_t, H_far: cfdfield_t,
dkx: float, dkx: float,
dky: float, dky: float,
padded_size: List[int] = None padded_size: List[int] = None

60
meanas/fdfd/functional.py
View File

@ -2,13 +2,13 @@
Functional versions of many FDFD operators. These can be useful for performing Functional versions of many FDFD operators. These can be useful for performing
FDFD calculations without needing to construct large matrices in memory. FDFD calculations without needing to construct large matrices in memory.
The functions generated here expect `fdfield_t` inputs with shape (3, 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 component has shape (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 from typing import Callable, Tuple, Optional
import numpy 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 from ..fdmath.functional import curl_forward, curl_back
@ -19,8 +19,8 @@ def e_full(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: fdfield_t = None mu: Optional[fdfield_t] = None
) -> fdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Wave operator for use with E-field. See `operators.e_full` for details. Wave operator for use with E-field. See `operators.e_full` for details.
@ -37,13 +37,13 @@ def e_full(
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def op_1(e: fdfield_t) -> fdfield_t: def op_1(e: cfdfield_t) -> cfdfield_t:
curls = ch(ce(e)) 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: def op_mu(e: cfdfield_t) -> cfdfield_t:
curls = ch(mu * ce(e)) curls = ch(mu * ce(e)) # type: ignore # mu = None ok because we don't return the function
return curls - omega ** 2 * epsilon * e # type: ignore # issues with numpy/mypy return curls - omega ** 2 * epsilon * e
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return op_1 return op_1
@ -56,7 +56,7 @@ def eh_full(
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: fdfield_t = None 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. Wave operator for full (both E and H) field representation.
See `operators.eh_full`. See `operators.eh_full`.
@ -74,13 +74,13 @@ def eh_full(
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def op_1(e: 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, 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, 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)): if numpy.any(numpy.equal(mu, None)):
return op_1 return op_1
@ -92,7 +92,7 @@ def e2h(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[fdfield_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Utility operator for converting the `E` field into the `H` field. Utility operator for converting the `E` field into the `H` field.
For use with `e_full` -- assumes that there is no magnetic current `M`. For use with `e_full` -- assumes that there is no magnetic current `M`.
@ -108,11 +108,11 @@ def e2h(
""" """
ce = curl_forward(dxes[0]) ce = curl_forward(dxes[0])
def e2h_1_1(e: fdfield_t) -> fdfield_t: def e2h_1_1(e: cfdfield_t) -> cfdfield_t:
return ce(e) / (-1j * omega) # type: ignore # issues with numpy/mypy return ce(e) / (-1j * omega)
def e2h_mu(e: fdfield_t) -> fdfield_t: def e2h_mu(e: cfdfield_t) -> cfdfield_t:
return ce(e) / (-1j * omega * mu) # type: ignore # issues with numpy/mypy return ce(e) / (-1j * omega * mu) # type: ignore # mu=None ok
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return e2h_1_1 return e2h_1_1
@ -124,7 +124,7 @@ def m2j(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: Optional[fdfield_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Utility operator for converting magnetic current `M` distribution Utility operator for converting magnetic current `M` distribution
into equivalent electric current distribution `J`. into equivalent electric current distribution `J`.
@ -141,13 +141,13 @@ def m2j(
""" """
ch = curl_back(dxes[1]) ch = curl_back(dxes[1])
def m2j_mu(m: fdfield_t) -> fdfield_t: def m2j_mu(m: cfdfield_t) -> cfdfield_t:
J = ch(m / mu) / (-1j * omega) J = ch(m / mu) / (-1j * omega) # type: ignore # mu=None ok
return J # type: ignore # issues with numpy/mypy return J
def m2j_1(m: fdfield_t) -> fdfield_t: def m2j_1(m: cfdfield_t) -> cfdfield_t:
J = ch(m) / (-1j * omega) J = ch(m) / (-1j * omega)
return J # type: ignore # issues with numpy/mypy return J
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
return m2j_1 return m2j_1
@ -161,7 +161,7 @@ def e_tfsf_source(
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_updater_t: ) -> cfdfield_updater_t:
""" """
Operator that turns an E-field distribution into a total-field/scattered-field Operator that turns an E-field distribution into a total-field/scattered-field
(TFSF) source. (TFSF) source.
@ -182,13 +182,13 @@ def e_tfsf_source(
# TODO documentation # TODO documentation
A = e_full(omega, dxes, epsilon, mu) A = e_full(omega, dxes, epsilon, mu)
def op(e: fdfield_t) -> fdfield_t: def op(e: cfdfield_t) -> cfdfield_t:
neg_iwj = A(TF_region * e) - TF_region * A(e) neg_iwj = A(TF_region * e) - TF_region * A(e)
return neg_iwj / (-1j * omega) return neg_iwj / (-1j * omega)
return op 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 Generates a function that takes the single-frequency `E` and `H` fields
and calculates the cross product `E` x `H` = $E \\times H$ as required and calculates the cross product `E` x `H` = $E \\times H$ as required
@ -210,7 +210,7 @@ def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdf
Returns: Returns:
Function `f` that returns E x H as required for the poynting vector. Function `f` that returns E x H as required for the poynting vector.
""" """
def exh(e: fdfield_t, h: fdfield_t) -> fdfield_t: def exh(e: cfdfield_t, h: cfdfield_t) -> cfdfield_t:
s = numpy.empty_like(e) s = numpy.empty_like(e)
ex = e[0] * dxes[0][0][:, None, None] ex = e[0] * dxes[0][0][:, None, None]
ey = e[1] * dxes[0][1][None, :, None] ey = e[1] * dxes[0][1][None, :, None]

12
meanas/fdfd/operators.py
View File

@ -31,7 +31,7 @@ from typing import Tuple, Optional
import numpy import numpy
import scipy.sparse as sparse # type: ignore 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 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)): if numpy.any(numpy.equal(mu, None)):
m_div = sparse.eye(epsilon.size) m_div = sparse.eye(epsilon.size)
else: 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 op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe
return op return op
@ -275,7 +275,7 @@ def e2h(
op = curl_forward(dxes[0]) / (-1j * omega) op = curl_forward(dxes[0]) / (-1j * omega)
if not numpy.any(numpy.equal(mu, None)): 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)): if not numpy.any(numpy.equal(pmc, None)):
op = sparse.diags(numpy.where(pmc, 0, 1)) @ op op = sparse.diags(numpy.where(pmc, 0, 1)) @ op
@ -303,12 +303,12 @@ def m2j(
op = curl_back(dxes[1]) / (1j * omega) op = curl_back(dxes[1]) / (1j * omega)
if not numpy.any(numpy.equal(mu, None)): 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 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 Operator for computing the Poynting vector, containing the
(E x) portion of the Poynting vector. (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 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. Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector.

15
meanas/fdfd/scpml.py
View File

@ -11,7 +11,7 @@ from numpy.typing import ArrayLike, NDArray
__author__ = 'Jan Petykiewicz' __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()`""" """Typedef for s-functions, see `prepare_s_function()`"""
@ -39,8 +39,8 @@ def prepare_s_function(
def uniform_grid_scpml( def uniform_grid_scpml(
shape: ArrayLike, # ints shape: Sequence[int],
thicknesses: ArrayLike, # ints thicknesses: Sequence[int],
omega: float, omega: float,
epsilon_effective: float = 1.0, epsilon_effective: float = 1.0,
s_function: Optional[s_function_t] = None, s_function: Optional[s_function_t] = None,
@ -70,12 +70,11 @@ def uniform_grid_scpml(
if s_function is None: if s_function is None:
s_function = prepare_s_function() s_function = prepare_s_function()
shape = tuple(shape)
thicknesses = tuple(thicknesses)
# Normalized distance to nearest boundary # Normalized distance to nearest boundary
def ll( def ll(u: NDArray[numpy.float64], n: int, t: int) -> NDArray[numpy.float64]:
u: NDArray[numpy.float64],
n: NDArray[numpy.float64],
t: NDArray[numpy.float64],
) -> NDArray[numpy.float64]:
return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t
dx_a = [numpy.array(numpy.inf)] * 3 dx_a = [numpy.array(numpy.inf)] * 3

6
meanas/fdfd/solvers.py
View File

@ -10,7 +10,7 @@ from numpy.typing import ArrayLike, NDArray
from numpy.linalg import norm from numpy.linalg import norm
import scipy.sparse.linalg # type: ignore 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 from . import operators
@ -65,7 +65,7 @@ def _scipy_qmr(
def generic( def generic(
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
J: vfdfield_t, J: vcfdfield_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: vfdfield_t = None,
pec: vfdfield_t = None, pec: vfdfield_t = None,
@ -73,7 +73,7 @@ def generic(
adjoint: bool = False, adjoint: bool = False,
matrix_solver: Callable[..., ArrayLike] = _scipy_qmr, matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
matrix_solver_opts: Optional[Dict[str, Any]] = None, matrix_solver_opts: Optional[Dict[str, Any]] = None,
) -> vfdfield_t: ) -> vcfdfield_t:
""" """
Conjugate gradient FDFD solver using CSR sparse matrices. Conjugate gradient FDFD solver using CSR sparse matrices.

26
meanas/fdfd/waveguide_2d.py
View File

@ -185,7 +185,7 @@ from numpy.linalg import norm
import scipy.sparse as sparse # type: ignore import scipy.sparse as sparse # type: ignore
from ..fdmath.operators import deriv_forward, deriv_back, cross 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 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
@ -335,7 +335,7 @@ def normalized_fields_e(
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
prop_phase: float = 0, 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, Given a vector `e_xy` containing the vectorized E_x and E_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -370,7 +370,7 @@ def normalized_fields_h(
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
prop_phase: float = 0, 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, Given a vector `h_xy` containing the vectorized H_x and H_y fields,
returns normalized, vectorized E and H fields for the system. returns normalized, vectorized E and H fields for the system.
@ -398,14 +398,14 @@ def normalized_fields_h(
def _normalized_fields( def _normalized_fields(
e: ArrayLike, e: vcfdfield_t,
h: ArrayLike, h: vcfdfield_t,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: Optional[vfdfield_t] = None, mu: Optional[vfdfield_t] = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]: ) -> Tuple[vcfdfield_t, vcfdfield_t]:
# TODO documentation # TODO documentation
shape = [s.size for s in dxes[0]] shape = [s.size for s in dxes[0]]
dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)] 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) op = curl_e(wavenumber, dxes) / (-1j * omega)
if not numpy.any(numpy.equal(mu, None)): 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 return op
@ -649,7 +649,7 @@ def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
def h_err( def h_err(
h: vfdfield_t, h: vcfdfield_t,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
@ -684,12 +684,12 @@ def h_err(
def e_err( def e_err(
e: vfdfield_t, e: vcfdfield_t,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = Optional[None] mu: Optional[vfdfield_t] = None,
) -> float: ) -> float:
""" """
Calculates the relative error in the E field Calculates the relative error in the E field
@ -711,7 +711,7 @@ def e_err(
if numpy.any(numpy.equal(mu, None)): if numpy.any(numpy.equal(mu, None)):
op = ch @ ce @ e - omega ** 2 * (epsilon * e) op = ch @ ce @ e - omega ** 2 * (epsilon * e)
else: else:
mu_inv = sparse.diags(1 / mu) # type: ignore # checked that mu is not None mu_inv = sparse.diags(1 / mu)
op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e) op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
return norm(op) / norm(e) return norm(op) / norm(e)
@ -724,7 +724,7 @@ def solve_modes(
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: vfdfield_t = None,
mode_margin: int = 2, 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. 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, mode_number: int,
*args: Any, *args: Any,
**kwargs: Any, **kwargs: Any,
) -> Tuple[vfdfield_t, complex]: ) -> Tuple[vcfdfield_t, complex]:
""" """
Wrapper around `solve_modes()` that solves for a single mode. Wrapper around `solve_modes()` that solves for a single mode.

14
meanas/fdfd/waveguide_3d.py
View File

@ -8,7 +8,7 @@ from typing import Dict, Optional, Sequence, Union, Any
import numpy import numpy
from numpy.typing import NDArray 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 from . import operators, waveguide_2d
@ -106,7 +106,7 @@ def solve_mode(
def compute_source( def compute_source(
E: fdfield_t, E: cfdfield_t,
wavenumber: complex, wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
@ -115,7 +115,7 @@ def compute_source(
slices: Sequence[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield_t,
mu: Optional[fdfield_t] = None, mu: Optional[fdfield_t] = None,
) -> fdfield_t: ) -> cfdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, returns the current source distribution Given an eigenmode obtained by `solve_mode`, returns the current source distribution
necessary to position a unidirectional source at the slice location. necessary to position a unidirectional source at the slice location.
@ -152,13 +152,13 @@ def compute_source(
def compute_overlap_e( def compute_overlap_e(
E: fdfield_t, E: cfdfield_t,
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: Sequence[slice], slices: Sequence[slice],
) -> fdfield_t: # TODO DOCS ) -> cfdfield_t: # TODO DOCS
""" """
Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the 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) 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( def expand_e(
E: fdfield_t, E: cfdfield_t,
wavenumber: complex, wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: Sequence[slice], slices: Sequence[slice],
) -> fdfield_t: ) -> cfdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D 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 slice where the mode was calculated to the entire domain (along the propagation

8
meanas/fdfd/waveguide_cyl.py
View File

@ -12,7 +12,7 @@ from typing import Dict, Union
import numpy import numpy
import scipy.sparse as sparse # type: ignore 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 ..fdmath.operators import deriv_forward, deriv_back
from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
@ -85,7 +85,7 @@ def solve_mode(
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
r0: float, r0: float,
) -> Dict[str, Union[complex, fdfield_t]]: ) -> Dict[str, Union[complex, cfdfield_t]]:
""" """
TODO: fixup TODO: fixup
Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@ -103,8 +103,8 @@ def solve_mode(
Returns: Returns:
``` ```
{ {
'E': List[NDArray[numpy.float_]], 'E': List[NDArray[numpy.complex_]],
'H': List[NDArray[numpy.float_]], 'H': List[NDArray[numpy.complex_]],
'wavenumber': complex, 'wavenumber': complex,
} }
``` ```

3
meanas/fdmath/__init__.py
View File

@ -741,7 +741,8 @@ the true values can be multiplied back in after the simulation is complete if no
normalized results are needed. 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 .vectorization import vec, unvec
from . import operators, functional, types, vectorization from . import operators, functional, types, vectorization

8
meanas/fdmath/functional.py
View File

@ -24,7 +24,7 @@ def deriv_forward(
Returns: Returns:
List of functions for taking forward derivatives along each axis. 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], 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=1) - f) / dx_e[1][None, :, None],
lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :]) lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :])
@ -48,7 +48,7 @@ def deriv_back(
Returns: Returns:
List of functions for taking forward derivatives along each axis. 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], 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=1)) / dx_h[1][None, :, None],
lambda f: (f - numpy.roll(f, 1, axis=2)) / dx_h[2][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: ) -> Callable:
Dx, Dy, Dz = deriv_forward(dx_e) 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])), return ((-Dz(e[1]), Dy(e[2])),
( Dz(e[0]), -Dx(e[2])), ( Dz(e[0]), -Dx(e[2])),
(-Dy(e[0]), Dx(e[1]))) (-Dy(e[0]), Dx(e[1])))
@ -135,7 +135,7 @@ def curl_back_parts(
) -> Callable: ) -> Callable:
Dx, Dy, Dz = deriv_back(dx_h) 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])), return ((-Dz(h[1]), Dy(h[2])),
( Dz(h[0]), -Dx(h[2])), ( Dz(h[0]), -Dx(h[2])),
(-Dy(h[0]), Dx(h[1]))) (-Dy(h[0]), Dx(h[1])))

9
meanas/fdmath/types.py
View File

@ -13,6 +13,12 @@ fdfield_t = NDArray[numpy.float_]
vfdfield_t = NDArray[numpy.float_] vfdfield_t = NDArray[numpy.float_]
"""Linearized vector field (single vector of length 3*X*Y*Z)""" """Linearized vector field (single vector of length 3*X*Y*Z)"""
cfdfield_t = NDArray[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_]]] 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] fdfield_updater_t = Callable[..., fdfield_t]
"""Convenience type for functions which take and return an fdfield_t""" """Convenience type for functions which take and return an fdfield_t"""
cfdfield_updater_t = Callable[..., cfdfield_t]
"""Convenience type for functions which take and return an cfdfield_t"""

28
meanas/fdmath/vectorization.py
View File

@ -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. 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 import numpy
from numpy.typing import ArrayLike from numpy.typing import ArrayLike
from .types import fdfield_t, vfdfield_t from .types import fdfield_t, vfdfield_t, cfdfield_t, vcfdfield_t
@overload @overload
@ -16,10 +16,18 @@ def vec(f: None) -> None:
pass pass
@overload @overload
def vec(f: Union[fdfield_t, List[ArrayLike]]) -> vfdfield_t: def vec(f: fdfield_t) -> vfdfield_t:
pass 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. 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 @overload
def unvec(v: None, shape: ArrayLike) -> None: def unvec(v: None, shape: Sequence[int]) -> None:
pass pass
@overload @overload
def unvec(v: vfdfield_t, shape: ArrayLike) -> fdfield_t: def unvec(v: vfdfield_t, shape: Sequence[int]) -> fdfield_t:
pass 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 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 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)): if numpy.any(numpy.equal(v, None)):
return 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')

47
meanas/fdtd/pml.py
View File

@ -8,8 +8,9 @@ PML implementations
# TODO retest pmls! # TODO retest pmls!
from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional from typing import List, Callable, Tuple, Dict, Sequence, Any, Optional
from copy import deepcopy
import numpy import numpy
from typing import NDArray from numpy.typing import NDArray, DTypeLike
from ..fdmath import fdfield_t, dx_lists_t from ..fdmath import fdfield_t, dx_lists_t
from ..fdmath.functional import deriv_forward, deriv_back from ..fdmath.functional import deriv_forward, deriv_back
@ -97,34 +98,38 @@ def updates_with_cpml(
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: fdfield_t, epsilon: fdfield_t,
*, *,
dtype: numpy.dtype = numpy.float32, dtype: DTypeLike = numpy.float32,
) -> Tuple[Callable[[fdfield_t, fdfield_t], None], ) -> Tuple[Callable[[fdfield_t, fdfield_t, fdfield_t], None],
Callable[[fdfield_t, fdfield_t], None]]: Callable[[fdfield_t, fdfield_t, fdfield_t], None]]:
Dfx, Dfy, Dfz = deriv_forward(dxes[1]) Dfx, Dfy, Dfz = deriv_forward(dxes[1])
Dbx, Dby, Dbz = deriv_back(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]] psi_E: List[List[Tuple[Any, Any]]] = [[(None, None) for _ in range(2)] for _ in range(3)]
params_E = [[None, None], [None, None], [None, None]] psi_H: List[List[Tuple[Any, Any]]] = deepcopy(psi_E)
params_H = [[None, None], [None, None], [None, None]] 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 axis in range(3):
for pp, polarity in enumerate((-1, 1)): 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_E[axis][pp] = (None, None)
psi_H[axis][pp] = (None, None) psi_H[axis][pp] = (None, None)
continue continue
cpml_param = cpml_params[axis][pp]
region = cpml_param['region'] region = cpml_param['region']
region_shape = epsilon[0][region].shape region_shape = epsilon[0][region].shape
psi_E[axis][pp] = (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),
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_E[axis][pp] = cpml_param['param_e'] + (region,)
params_H[axis][pp] = cpml_param['param_h'] + (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) pE = numpy.empty_like(epsilon, dtype=dtype)
pH = 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]) dyHx = Dby(h[0])
dzHx = Dbz(h[0]) dzHx = Dbz(h[0])
dxHy = Dbx(h[1]) dxHy = Dbx(h[1])
@ -175,7 +184,11 @@ def updates_with_cpml(
e[2] += dt / epsilon[2] * (dxHy - dyHx + pE[2]) 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]) dyEx = Dfy(e[0])
dzEx = Dfz(e[0]) dzEx = Dfz(e[0])
dxEy = Dfx(e[1]) dxEy = Dfx(e[1])

15
meanas/test/test_fdfd.py
View File

@ -99,8 +99,8 @@ class FDResult:
dxes: List[List[NDArray[numpy.float64]]] dxes: List[List[NDArray[numpy.float64]]]
epsilon: NDArray[numpy.float64] epsilon: NDArray[numpy.float64]
omega: complex omega: complex
j: NDArray[numpy.float64] j: NDArray[numpy.complex128]
e: NDArray[numpy.float64] e: NDArray[numpy.complex128]
pmc: Optional[NDArray[numpy.float64]] pmc: Optional[NDArray[numpy.float64]]
pec: Optional[NDArray[numpy.float64]] pec: Optional[NDArray[numpy.float64]]
@ -111,7 +111,7 @@ def sim(
shape: Tuple[int, ...], shape: Tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: NDArray[numpy.float64],
dxes: List[List[NDArray[numpy.float64]]], dxes: List[List[NDArray[numpy.float64]]],
j_distribution: NDArray[numpy.float64], j_distribution: NDArray[numpy.complex128],
omega: float, omega: float,
pec: Optional[NDArray[numpy.float64]], pec: Optional[NDArray[numpy.float64]],
pmc: Optional[NDArray[numpy.float64]], pmc: Optional[NDArray[numpy.float64]],
@ -134,8 +134,13 @@ def sim(
j_vec = vec(j_distribution) j_vec = vec(j_distribution)
eps_vec = vec(epsilon) eps_vec = vec(epsilon)
e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec, e_vec = fdfd.solvers.generic(
matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11}) 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:]) e = unvec(e_vec, shape[1:])
sim = FDResult( sim = FDResult(

4
meanas/test/test_fdfd_pml.py
View File

@ -78,7 +78,7 @@ def j_distribution(
dxes: dx_lists_mut, dxes: dx_lists_mut,
omega: float, omega: float,
src_polarity: int, src_polarity: int,
) -> Iterable[NDArray[numpy.float64]]: ) -> Iterable[NDArray[numpy.complex128]]:
j = numpy.zeros(shape, dtype=complex) j = numpy.zeros(shape, dtype=complex)
dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0] # Propagation axis
@ -150,7 +150,7 @@ def sim(
shape: Tuple[int, ...], shape: Tuple[int, ...],
epsilon: NDArray[numpy.float64], epsilon: NDArray[numpy.float64],
dxes: dx_lists_mut, dxes: dx_lists_mut,
j_distribution: NDArray[numpy.float64], j_distribution: NDArray[numpy.complex128],
omega: float, omega: float,
pec: Optional[NDArray[numpy.float64]], pec: Optional[NDArray[numpy.float64]],
pmc: Optional[NDArray[numpy.float64]], pmc: Optional[NDArray[numpy.float64]],

11
meanas/test/utils.py
View File

@ -1,14 +1,15 @@
from typing import Any from typing import Any
import numpy import numpy
from typing import ArrayLike from numpy.typing import ArrayLike, NDArray
PRNG = numpy.random.RandomState(12345) PRNG = numpy.random.RandomState(12345)
def assert_fields_close( def assert_fields_close(
x: ArrayLike, x: NDArray,
y: ArrayLike, y: NDArray,
*args: Any, *args: Any,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None:
@ -18,8 +19,8 @@ def assert_fields_close(
numpy.rollaxis(y, -1)), *args, **kwargs) numpy.rollaxis(y, -1)), *args, **kwargs)
def assert_close( def assert_close(
x: ArrayLike, x: NDArray,
y: ArrayLike, y: NDArray,
*args: Any, *args: Any,
**kwargs: Any, **kwargs: Any,
) -> None: ) -> None: