|
|
|
@ -94,7 +94,7 @@ This module contains functions for generating and solving the
|
|
|
|
|
|
|
|
|
|
'''
|
|
|
|
|
|
|
|
|
|
from typing import Tuple, Callable, Any, List, Optional, cast, Union, Sequence
|
|
|
|
|
from typing import Callable, Any, cast, Sequence
|
|
|
|
|
import logging
|
|
|
|
|
import numpy
|
|
|
|
|
from numpy import pi, real, trace
|
|
|
|
@ -140,7 +140,7 @@ def generate_kmn(
|
|
|
|
|
k0: ArrayLike,
|
|
|
|
|
G_matrix: 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.
|
|
|
|
|
|
|
|
|
@ -155,8 +155,9 @@ def generate_kmn(
|
|
|
|
|
All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
|
|
|
|
|
"""
|
|
|
|
|
k0 = numpy.array(k0)
|
|
|
|
|
G_matrix = numpy.array(G_matrix, copy=False)
|
|
|
|
|
|
|
|
|
|
Gi_grids = numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij')
|
|
|
|
|
Gi_grids = numpy.array(numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij'))
|
|
|
|
|
Gi = numpy.moveaxis(Gi_grids, 0, -1)
|
|
|
|
|
|
|
|
|
|
k_G = k0[None, None, None, :] - Gi
|
|
|
|
@ -183,7 +184,7 @@ def maxwell_operator(
|
|
|
|
|
k0: ArrayLike,
|
|
|
|
|
G_matrix: ArrayLike,
|
|
|
|
|
epsilon: fdfield_t,
|
|
|
|
|
mu: Optional[fdfield_t] = None
|
|
|
|
|
mu: fdfield_t | None = None
|
|
|
|
|
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
|
|
|
|
|
"""
|
|
|
|
|
Generate the Maxwell operator
|
|
|
|
@ -225,9 +226,11 @@ def maxwell_operator(
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
h: Raveled h_mn; size `2 * epsilon[0].size`.
|
|
|
|
|
Altered in-place.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)).
|
|
|
|
|
Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)), returned
|
|
|
|
|
and overwritten in-place of `h`.
|
|
|
|
|
"""
|
|
|
|
|
hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
|
|
|
|
|
|
|
|
|
@ -235,29 +238,35 @@ def maxwell_operator(
|
|
|
|
|
|
|
|
|
|
# cross product and transform into xyz basis
|
|
|
|
|
d_xyz = (n * hin_m
|
|
|
|
|
- m * hin_n) * k_mag
|
|
|
|
|
- m * hin_n) * k_mag # noqa: E128
|
|
|
|
|
|
|
|
|
|
# divide by epsilon
|
|
|
|
|
e_xyz = fftn(ifftn(d_xyz, axes=range(3)) / epsilon, axes=range(3))
|
|
|
|
|
temp = ifftn(d_xyz, axes=range(3)) # reuses d_xyz if using pyfftw
|
|
|
|
|
temp /= epsilon
|
|
|
|
|
e_xyz = fftn(temp, axes=range(3))
|
|
|
|
|
|
|
|
|
|
# cross product and transform into mn basis
|
|
|
|
|
b_m = numpy.sum(e_xyz * n, axis=3)[:, :, :, None] * -k_mag
|
|
|
|
|
b_n = numpy.sum(e_xyz * m, axis=3)[:, :, :, None] * +k_mag
|
|
|
|
|
b_m = numpy.sum(e_xyz * n, axis=3, keepdims=True) * -k_mag
|
|
|
|
|
b_n = numpy.sum(e_xyz * m, axis=3, keepdims=True) * +k_mag
|
|
|
|
|
|
|
|
|
|
if mu is None:
|
|
|
|
|
h_m, h_n = b_m, b_n
|
|
|
|
|
else:
|
|
|
|
|
# transform from mn to xyz
|
|
|
|
|
b_xyz = (m * b_m[:, :, :, None]
|
|
|
|
|
+ n * b_n[:, :, :, None])
|
|
|
|
|
+ n * b_n[:, :, :, None]) # noqa: E128
|
|
|
|
|
|
|
|
|
|
# divide by mu
|
|
|
|
|
h_xyz = fftn(ifftn(b_xyz, axes=range(3)) / mu, axes=range(3))
|
|
|
|
|
temp = ifftn(b_xyz, axes=range(3))
|
|
|
|
|
temp /= mu
|
|
|
|
|
h_xyz = fftn(temp, axes=range(3))
|
|
|
|
|
|
|
|
|
|
# transform back to mn
|
|
|
|
|
h_m = numpy.sum(h_xyz * m, axis=3)
|
|
|
|
|
h_n = numpy.sum(h_xyz * n, axis=3)
|
|
|
|
|
return numpy.hstack((h_m.ravel(), h_n.ravel()))
|
|
|
|
|
|
|
|
|
|
h = numpy.concatenate((h_m, h_n), axis=None, out=h) # ravel and merge
|
|
|
|
|
return h
|
|
|
|
|
|
|
|
|
|
return operator
|
|
|
|
|
|
|
|
|
@ -294,7 +303,7 @@ def hmn_2_exyz(
|
|
|
|
|
def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
|
|
|
|
|
hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
|
|
|
|
|
d_xyz = (n * hin_m
|
|
|
|
|
- m * hin_n) * k_mag
|
|
|
|
|
- m * hin_n) * k_mag # noqa: E128
|
|
|
|
|
|
|
|
|
|
# divide by epsilon
|
|
|
|
|
return numpy.array([ei for ei in numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)]) # TODO avoid copy
|
|
|
|
@ -332,7 +341,7 @@ 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)
|
|
|
|
|
+ n * hin_n) # noqa: E128
|
|
|
|
|
return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])
|
|
|
|
|
|
|
|
|
|
return operator
|
|
|
|
@ -342,7 +351,7 @@ def inverse_maxwell_operator_approx(
|
|
|
|
|
k0: ArrayLike,
|
|
|
|
|
G_matrix: ArrayLike,
|
|
|
|
|
epsilon: fdfield_t,
|
|
|
|
|
mu: Optional[fdfield_t] = None,
|
|
|
|
|
mu: fdfield_t | None = None,
|
|
|
|
|
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
|
|
|
|
|
"""
|
|
|
|
|
Generate an approximate inverse of the Maxwell operator,
|
|
|
|
@ -392,10 +401,12 @@ def inverse_maxwell_operator_approx(
|
|
|
|
|
else:
|
|
|
|
|
# transform from mn to xyz
|
|
|
|
|
h_xyz = (m * hin_m[:, :, :, None]
|
|
|
|
|
+ n * hin_n[:, :, :, None])
|
|
|
|
|
+ n * hin_n[:, :, :, None]) # noqa: E128
|
|
|
|
|
|
|
|
|
|
# multiply by mu
|
|
|
|
|
b_xyz = fftn(ifftn(h_xyz, axes=range(3)) * mu, axes=range(3))
|
|
|
|
|
temp = ifftn(h_xyz, axes=range(3))
|
|
|
|
|
temp *= mu
|
|
|
|
|
b_xyz = fftn(temp, axes=range(3))
|
|
|
|
|
|
|
|
|
|
# transform back to mn
|
|
|
|
|
b_m = numpy.sum(b_xyz * m, axis=3)
|
|
|
|
@ -403,16 +414,19 @@ def inverse_maxwell_operator_approx(
|
|
|
|
|
|
|
|
|
|
# cross product and transform into xyz basis
|
|
|
|
|
e_xyz = (n * b_m
|
|
|
|
|
- m * b_n) / k_mag
|
|
|
|
|
- m * b_n) / k_mag # noqa: E128
|
|
|
|
|
|
|
|
|
|
# multiply by epsilon
|
|
|
|
|
d_xyz = fftn(ifftn(e_xyz, axes=range(3)) * epsilon, axes=range(3))
|
|
|
|
|
temp = ifftn(e_xyz, axes=range(3))
|
|
|
|
|
temp *= epsilon
|
|
|
|
|
d_xyz = fftn(temp, axes=range(3))
|
|
|
|
|
|
|
|
|
|
# cross product and transform into mn basis crossinv_t2c
|
|
|
|
|
h_m = numpy.sum(d_xyz * n, axis=3)[:, :, :, None] / +k_mag
|
|
|
|
|
h_n = numpy.sum(d_xyz * m, axis=3)[:, :, :, None] / -k_mag
|
|
|
|
|
h_m = numpy.sum(d_xyz * n, axis=3, keepdims=True) / +k_mag
|
|
|
|
|
h_n = numpy.sum(d_xyz * m, axis=3, keepdims=True) / -k_mag
|
|
|
|
|
|
|
|
|
|
return numpy.hstack((h_m.ravel(), h_n.ravel()))
|
|
|
|
|
h = numpy.concatenate((h_m, h_n), axis=None, out=h)
|
|
|
|
|
return h
|
|
|
|
|
|
|
|
|
|
return operator
|
|
|
|
|
|
|
|
|
@ -423,13 +437,13 @@ def find_k(
|
|
|
|
|
direction: ArrayLike,
|
|
|
|
|
G_matrix: ArrayLike,
|
|
|
|
|
epsilon: fdfield_t,
|
|
|
|
|
mu: Optional[fdfield_t] = None,
|
|
|
|
|
mu: fdfield_t | None = None,
|
|
|
|
|
band: int = 0,
|
|
|
|
|
k_bounds: Tuple[float, float] = (0, 0.5),
|
|
|
|
|
k_guess: Optional[float] = None,
|
|
|
|
|
solve_callback: Optional[Callable[..., None]] = None,
|
|
|
|
|
iter_callback: Optional[Callable[..., None]] = None,
|
|
|
|
|
) -> Tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
|
|
|
|
|
k_bounds: tuple[float, float] = (0, 0.5),
|
|
|
|
|
k_guess: float | None = None,
|
|
|
|
|
solve_callback: Callable[..., None] | None = None,
|
|
|
|
|
iter_callback: Callable[..., None] | None = None,
|
|
|
|
|
) -> tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
|
|
|
|
|
"""
|
|
|
|
|
Search for a bloch vector that has a given frequency.
|
|
|
|
|
|
|
|
|
@ -490,13 +504,13 @@ def eigsolve(
|
|
|
|
|
k0: ArrayLike,
|
|
|
|
|
G_matrix: ArrayLike,
|
|
|
|
|
epsilon: fdfield_t,
|
|
|
|
|
mu: Optional[fdfield_t] = None,
|
|
|
|
|
mu: fdfield_t | None = None,
|
|
|
|
|
tolerance: float = 1e-20,
|
|
|
|
|
max_iters: int = 10000,
|
|
|
|
|
reset_iters: int = 100,
|
|
|
|
|
y0: Optional[ArrayLike] = None,
|
|
|
|
|
callback: Optional[Callable[..., None]] = None,
|
|
|
|
|
) -> Tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
|
|
|
|
|
y0: ArrayLike | None = None,
|
|
|
|
|
callback: Callable[..., None] | None = None,
|
|
|
|
|
) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
|
|
|
|
|
"""
|
|
|
|
|
Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
|
|
|
|
|
k0 of the specified structure.
|
|
|
|
@ -542,6 +556,7 @@ def eigsolve(
|
|
|
|
|
#prev_theta = 0.5
|
|
|
|
|
D = numpy.zeros(shape=y_shape, dtype=complex)
|
|
|
|
|
|
|
|
|
|
Z: NDArray[numpy.complex128]
|
|
|
|
|
if y0 is None:
|
|
|
|
|
Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
|
|
|
|
|
else:
|
|
|
|
@ -563,21 +578,25 @@ def eigsolve(
|
|
|
|
|
continue
|
|
|
|
|
break
|
|
|
|
|
|
|
|
|
|
Zt = numpy.empty(Z.shape[::-1])
|
|
|
|
|
AZ = numpy.empty(Z.shape)
|
|
|
|
|
|
|
|
|
|
for i in range(max_iters):
|
|
|
|
|
Zt = Z.conj().T
|
|
|
|
|
Zt = numpy.conj(Z.T, out=Zt)
|
|
|
|
|
ZtZ = Zt @ Z
|
|
|
|
|
U = numpy.linalg.inv(ZtZ)
|
|
|
|
|
AZ = scipy_op @ Z
|
|
|
|
|
AZ = scipy_op @ Z.copy()
|
|
|
|
|
ZtAZ = Zt @ AZ
|
|
|
|
|
ZtAZU = ZtAZ @ U
|
|
|
|
|
E_signed = real(trace(ZtAZU))
|
|
|
|
|
sgn = numpy.sign(E_signed)
|
|
|
|
|
E = numpy.abs(E_signed)
|
|
|
|
|
G = (AZ @ U - Z @ U @ ZtAZU) * sgn # G = AZU projected onto the space orthonormal to Z
|
|
|
|
|
# via (1 - ZUZt)
|
|
|
|
|
G = (AZ @ U - Z @ U @ ZtAZU) * sgn # G = AZU projected onto the space orthonormal to Z
|
|
|
|
|
# via (1 - ZUZt)
|
|
|
|
|
|
|
|
|
|
if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7):
|
|
|
|
|
logger.info('Optimization succeded: '
|
|
|
|
|
logger.info(
|
|
|
|
|
'Optimization succeded: '
|
|
|
|
|
f'[change in trace] {abs(E - prev_E)} - 5e-8 '
|
|
|
|
|
f'< {tolerance} [tolerance] * {(E + prev_E) / 2} [value of trace]'
|
|
|
|
|
)
|
|
|
|
@ -619,7 +638,7 @@ def eigsolve(
|
|
|
|
|
symZtD = _symmetrize(Zt @ D)
|
|
|
|
|
symZtAD = _symmetrize(Zt @ AD)
|
|
|
|
|
|
|
|
|
|
Qi_memo: List[Optional[float]] = [None, None]
|
|
|
|
|
Qi_memo: list[float | None] = [None, None]
|
|
|
|
|
|
|
|
|
|
def Qi_func(theta: float) -> float:
|
|
|
|
|
nonlocal Qi_memo
|
|
|
|
@ -643,8 +662,8 @@ def eigsolve(
|
|
|
|
|
else:
|
|
|
|
|
raise Exception('Inexplicable singularity in trace_func')
|
|
|
|
|
Qi_memo[0] = theta
|
|
|
|
|
Qi_memo[1] = Qi
|
|
|
|
|
return Qi
|
|
|
|
|
Qi_memo[1] = cast(float, Qi)
|
|
|
|
|
return cast(float, Qi)
|
|
|
|
|
|
|
|
|
|
def trace_func(theta: float) -> float:
|
|
|
|
|
c = numpy.cos(theta)
|
|
|
|
@ -756,7 +775,7 @@ 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.complex128],
|
|
|
|
|
B: Union[NDArray[numpy.complex128], float],
|
|
|
|
|
B: NDArray[numpy.complex128] | float,
|
|
|
|
|
) -> float:
|
|
|
|
|
return real(numpy.sum(A.conj() * B))
|
|
|
|
|
|
|
|
|
|