875 lines
32 KiB
Python
875 lines
32 KiB
Python
"""
|
|
Bloch eigenmode solvers and operators in a spatial Fourier basis.
|
|
|
|
The two stored coefficients at each wavevector represent transverse magnetic
|
|
flux B = mu H. They coincide with H when mu is omitted. Using B ensures
|
|
k dot B = 0 even in anisotropic or spatially varying magnetic materials; H
|
|
itself need not be transverse. The historical h_mn function names are retained.
|
|
Pass the same mu to hmn_2_hxyz and hmn_2_exyz when reconstructing these modes.
|
|
|
|
For each Fourier wavevector, generate_kmn constructs unit vectors m and n with
|
|
k cross m = |k| n. The Cartesian Fourier field is B_k = b_m m + b_n n.
|
|
The Maxwell operator applies, in order:
|
|
|
|
H_k = fftn(ifftn(B_k) / mu)
|
|
E_k = fftn(ifftn(k cross H_k) / epsilon)
|
|
A(B_k) = k cross E_k
|
|
|
|
Material division is performed in real space and curls in Fourier space.
|
|
The explicit factors of i are omitted, giving the historical negative squared
|
|
frequency eigenvalues. Wavevectors and frequencies use cycles per length:
|
|
frequency = sqrt(-real(eigenvalue)). Reconstructed electric fields retain the
|
|
historical overall frequency scaling and are not power-normalized.
|
|
|
|
For finite positive real permeability, the flux operator is self-adjoint in
|
|
the projected magnetic energy metric M = P mu^-1 P, where P projects onto the
|
|
transverse Fourier subspace. The block variational solve uses that metric in
|
|
its Rayleigh quotients, gradients, line searches, and final Ritz extraction.
|
|
The inverse-curl preconditioner regularizes zero wavevectors so static modes
|
|
remain accessible at Gamma. Nonzero-frequency modes have no static Fourier
|
|
flux component at Gamma; the two stored zero-wavevector basis vectors choose
|
|
two representatives of the static sector.
|
|
|
|
Typical usage:
|
|
|
|
values, vectors = eigsolve(5, k0, reciprocal_lattice, epsilon, mu=mu)
|
|
frequency = numpy.sqrt(-numpy.real(values[0]))
|
|
electric = hmn_2_exyz(k0, reciprocal_lattice, epsilon, mu=mu)(vectors[0])
|
|
magnetic = hmn_2_hxyz(k0, reciprocal_lattice, epsilon, mu=mu)(vectors[0])
|
|
|
|
The block iteration follows the approach of S. G. Johnson and J. D.
|
|
Joannopoulos, Optics Express 8, 173-190 (2001).
|
|
"""
|
|
|
|
from typing import Any, cast
|
|
from collections.abc import Callable, Sequence
|
|
import logging
|
|
import numpy
|
|
from numpy import pi, real, trace
|
|
from numpy.fft import fftfreq
|
|
from numpy.typing import NDArray, ArrayLike
|
|
import scipy
|
|
import scipy.optimize
|
|
from scipy.linalg import norm
|
|
import scipy.sparse.linalg as spalg
|
|
|
|
from ..fdmath import fdfield, cfdfield, cfdfield_t
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
try:
|
|
import pyfftw.interfaces.numpy_fft # type: ignore
|
|
import pyfftw.interfaces # type: ignore
|
|
logger.info('Using pyfftw')
|
|
|
|
pyfftw.interfaces.cache.enable()
|
|
pyfftw.interfaces.cache.set_keepalive_time(3600)
|
|
fftw_args = {
|
|
#'threads': multiprocessing.cpu_count(),
|
|
'overwrite_input': True,
|
|
#'planner_effort': 'FFTW_PATIENT',
|
|
}
|
|
|
|
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.complex128]:
|
|
return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)
|
|
|
|
except ImportError:
|
|
from numpy.fft import fftn, ifftn
|
|
logger.info('Using numpy fft')
|
|
|
|
|
|
def _assemble_hmn_vector(
|
|
h_m: NDArray[numpy.complex128],
|
|
h_n: NDArray[numpy.complex128],
|
|
) -> NDArray[numpy.complex128]:
|
|
stacked = numpy.concatenate((numpy.ravel(h_m), numpy.ravel(h_n)))
|
|
return stacked[:, None]
|
|
|
|
|
|
def generate_kmn(
|
|
k0: ArrayLike,
|
|
G_matrix: 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.
|
|
|
|
Args:
|
|
k0: [k0x, k0y, k0z], Bloch wavevector, in G basis.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
shape: [nx, ny, nz] shape of the simulation grid.
|
|
|
|
Returns:
|
|
`(|k|, m, n)` where `|k|` has shape `tuple(shape) + (1,)`
|
|
and `m`, `n` have shape `tuple(shape) + (3,)`.
|
|
All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
|
|
"""
|
|
k0 = numpy.array(k0)
|
|
G_matrix = numpy.asarray(G_matrix)
|
|
|
|
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
|
|
k_xyz = numpy.moveaxis(G_matrix @ numpy.moveaxis(k_G, 3, 2), 3, 2)
|
|
|
|
m = numpy.broadcast_to([0, 1, 0], tuple(shape[:3]) + (3,)).astype(float)
|
|
n = numpy.broadcast_to([0, 0, 1], tuple(shape[:3]) + (3,)).astype(float)
|
|
|
|
xy_non0 = numpy.any(k_xyz[:, :, :, :2] != 0, axis=3)
|
|
if numpy.any(xy_non0):
|
|
u = numpy.cross(k_xyz[xy_non0], [0, 0, 1])
|
|
m[xy_non0, :] = u / norm(u, axis=1)[:, None]
|
|
|
|
z_non0 = numpy.any(k_xyz != 0, axis=3)
|
|
if numpy.any(z_non0):
|
|
v = numpy.cross(k_xyz[z_non0], m[z_non0])
|
|
n[z_non0, :] = v / norm(v, axis=1)[:, None]
|
|
|
|
k_mag = norm(k_xyz, axis=3)[:, :, :, None]
|
|
return k_mag, m, n
|
|
|
|
|
|
def maxwell_operator(
|
|
k0: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
epsilon: fdfield,
|
|
mu: fdfield | None = None
|
|
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
|
|
"""
|
|
Generate the Maxwell operator on transverse magnetic-flux coefficients.
|
|
|
|
k x conv(1/eps_k, k x conv(1/mu_k, ___))
|
|
|
|
The corresponding real-space curl-curl problem is
|
|
|
|
curl(1/eps * curl(1/mu * ___))
|
|
|
|
with the overall factor ``-(2*pi)**2`` omitted for the module's
|
|
cycles-per-length frequency convention.
|
|
|
|
The state represents B = mu H, whose Fourier components are transverse.
|
|
It equals H when mu is omitted. Eigenvalues retain the historical negative
|
|
squared-frequency convention (the two explicit factors of i are omitted).
|
|
|
|
The operator is a function that acts on a vector h_mn of size `2 * epsilon[0].size`
|
|
|
|
See the `meanas.fdfd.bloch` docstring for more information.
|
|
|
|
Args:
|
|
k0: Bloch wavevector, `[k0x, k0y, k0z]`.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
epsilon: Dielectric constant distribution for the simulation.
|
|
All fields are sampled at cell centers (i.e., NOT Yee-gridded)
|
|
mu: Magnetic permability distribution for the simulation.
|
|
Default None (1 everywhere).
|
|
|
|
Returns:
|
|
Function which applies the maxwell operator to h_mn.
|
|
"""
|
|
|
|
shape = epsilon[0].shape + (1,)
|
|
k_mag, m, n = generate_kmn(k0, G_matrix, shape)
|
|
|
|
epsilon = numpy.moveaxis(epsilon, 0, -1)
|
|
if mu is not None:
|
|
mu = numpy.moveaxis(mu, 0, -1)
|
|
|
|
def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
|
|
"""
|
|
Maxwell operator for Bloch eigenmode simulation.
|
|
|
|
h contains the two transverse magnetic-flux coefficients.
|
|
|
|
Args:
|
|
h: Raveled b_mn; size `2 * epsilon[0].size`. Not modified.
|
|
|
|
Returns:
|
|
Raveled transverse flux image under the Maxwell operator.
|
|
"""
|
|
hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
|
|
|
|
#{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
|
|
|
|
# Convert transverse B into H before taking either curl. H itself need
|
|
# not be transverse for anisotropic or spatially varying permeability.
|
|
if mu is None:
|
|
h_m, h_n = hin_m, hin_n
|
|
else:
|
|
b_xyz = m * hin_m + n * hin_n
|
|
h_xyz = fftn(ifftn(b_xyz, axes=range(3)) / mu, axes=range(3))
|
|
h_m = numpy.sum(h_xyz * m, axis=3, keepdims=True)
|
|
h_n = numpy.sum(h_xyz * n, axis=3, keepdims=True)
|
|
d_xyz = (n * h_m - m * h_n) * k_mag
|
|
|
|
# divide by epsilon
|
|
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, keepdims=True) * -k_mag
|
|
b_n = numpy.sum(e_xyz * m, axis=3, keepdims=True) * +k_mag
|
|
|
|
return _assemble_hmn_vector(b_m, b_n)
|
|
|
|
return operator
|
|
|
|
|
|
def hmn_2_exyz(
|
|
k0: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
epsilon: fdfield,
|
|
mu: fdfield | None = None,
|
|
) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
|
|
"""
|
|
Convert transverse Fourier flux coefficients into an unnormalized E field.
|
|
|
|
ifft(conv(1/eps_k, k x conv(1/mu_k, b_mn)))
|
|
|
|
The operator is a function that acts on a vector `h_mn` of size `2 * epsilon[0].size`.
|
|
|
|
See the `meanas.fdfd.bloch` docstring for more information.
|
|
|
|
Args:
|
|
k0: Bloch wavevector, `[k0x, k0y, k0z]`.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
epsilon: Dielectric constant distribution for the simulation.
|
|
All fields are sampled at cell centers (i.e., NOT Yee-gridded)
|
|
mu: Permeability used by the solve. The input coefficients describe B;
|
|
divide by mu to recover H before computing the electric field.
|
|
|
|
Returns:
|
|
Function for converting `h_mn` into `E_xyz`
|
|
"""
|
|
shape = epsilon[0].shape + (1,)
|
|
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)
|
|
|
|
def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
|
|
hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
|
|
if mu is not None:
|
|
b_xyz = m * hin_m + n * hin_n
|
|
h_xyz = fftn(ifftn(b_xyz, axes=range(3)) / mu, axes=range(3))
|
|
hin_m = numpy.sum(h_xyz * m, axis=3, keepdims=True)
|
|
hin_n = numpy.sum(h_xyz * n, axis=3, keepdims=True)
|
|
d_xyz = (n * hin_m
|
|
- m * hin_n) * k_mag
|
|
|
|
# divide by epsilon
|
|
exyz = numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)
|
|
return cfdfield_t(exyz)
|
|
|
|
return operator
|
|
|
|
|
|
def hmn_2_hxyz(
|
|
k0: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
epsilon: fdfield,
|
|
mu: fdfield | None = None,
|
|
) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
|
|
"""
|
|
Convert transverse Fourier flux coefficients into a physical H field.
|
|
|
|
ifft(b_mn) / mu
|
|
|
|
The operator is a function that acts on a vector `h_mn` of size `2 * epsilon[0].size`.
|
|
|
|
See the `meanas.fdfd.bloch` docstring for more information.
|
|
|
|
Args:
|
|
k0: Bloch wavevector, `[k0x, k0y, k0z]`.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
epsilon: Dielectric constant distribution for the simulation.
|
|
Only `epsilon[0].shape` is used.
|
|
mu: Permeability used by the solve. Input coefficients describe B;
|
|
the returned physical H is B / mu. Default is unit permeability.
|
|
|
|
Returns:
|
|
Function for converting `h_mn` into `H_xyz`
|
|
"""
|
|
shape = epsilon[0].shape + (1,)
|
|
_k_mag, m, n = generate_kmn(k0, G_matrix, shape)
|
|
|
|
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)
|
|
field = numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])
|
|
if mu is not None:
|
|
field /= mu
|
|
return cfdfield_t(field)
|
|
|
|
return operator
|
|
|
|
|
|
def inverse_maxwell_operator_approx(
|
|
k0: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
epsilon: fdfield,
|
|
mu: fdfield | None = None,
|
|
) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
|
|
"""
|
|
Generate a positive approximate inverse of minus the flux Maxwell operator,
|
|
|
|
project_mn(conv(mu_k, cross_inv(conv(eps_k, cross_inv(___)))))
|
|
|
|
At k = 0, inverse crosses use the smallest nonzero grid wavevector scale
|
|
(or one on a single-point grid). This finite regularization keeps the
|
|
preconditioner positive on static directions, which must remain accessible
|
|
to the eigensolver at Gamma.
|
|
|
|
It is used for the preconditioned block variational solve.
|
|
|
|
See the `meanas.fdfd.bloch` docstring for more information.
|
|
|
|
Args:
|
|
k0: Bloch wavevector, `[k0x, k0y, k0z]`.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
epsilon: Dielectric constant distribution for the simulation.
|
|
All fields are sampled at cell centers (i.e., NOT Yee-gridded)
|
|
mu: Magnetic permability distribution for the simulation.
|
|
Default None (1 everywhere).
|
|
|
|
Returns:
|
|
Function which applies the approximate inverse of the maxwell operator to `h_mn`.
|
|
"""
|
|
shape = epsilon[0].shape + (1,)
|
|
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)
|
|
nonzero = k_mag[k_mag != 0]
|
|
reference = float(nonzero.min()) if nonzero.size else 1.0
|
|
k_inverse = numpy.divide(1, k_mag, out=numpy.full_like(k_mag, 1 / reference), where=k_mag != 0)
|
|
|
|
def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
|
|
"""
|
|
Approximate inverse Maxwell operator for Bloch eigenmode simulation.
|
|
|
|
h is complex 2-field in (m, n) basis, vectorized
|
|
|
|
Args:
|
|
h: Raveled h_mn; size `2 * epsilon[0].size`.
|
|
|
|
Returns:
|
|
Raveled, regularized approximate inverse flux image.
|
|
"""
|
|
hin_m, hin_n = (hi.reshape(shape) for hi in numpy.split(h, 2))
|
|
|
|
#{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
|
|
|
|
# cross product and transform into xyz basis
|
|
e_xyz = (n * hin_m - m * hin_n) * k_inverse
|
|
|
|
# multiply by epsilon
|
|
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, keepdims=True) * k_inverse
|
|
h_n = numpy.sum(d_xyz * m, axis=3, keepdims=True) * -k_inverse
|
|
if mu is not None:
|
|
h_xyz = m * h_m + n * h_n
|
|
b_xyz = fftn(ifftn(h_xyz, axes=range(3)) * mu, axes=range(3))
|
|
h_m = numpy.sum(b_xyz * m, axis=3, keepdims=True)
|
|
h_n = numpy.sum(b_xyz * n, axis=3, keepdims=True)
|
|
|
|
return _assemble_hmn_vector(h_m, h_n)
|
|
|
|
return operator
|
|
|
|
|
|
def _transverse_material_operator(
|
|
k0: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
material: fdfield,
|
|
*,
|
|
inverse: bool,
|
|
) -> Callable[[NDArray], NDArray]:
|
|
"""Apply P material^(+/-1) P on the transverse Fourier subspace."""
|
|
shape = material[0].shape + (1,)
|
|
_k_mag, m, n = generate_kmn(k0, G_matrix, shape)
|
|
values = numpy.moveaxis(material, 0, -1)
|
|
|
|
def apply(h: NDArray) -> NDArray:
|
|
hm, hn = (part.reshape(shape) for part in numpy.split(h, 2))
|
|
field = ifftn(m * hm + n * hn, axes=range(3))
|
|
field = field / values if inverse else field * values
|
|
transformed = fftn(field, axes=range(3))
|
|
return _assemble_hmn_vector(numpy.sum(transformed * m, axis=3), numpy.sum(transformed * n, axis=3))
|
|
|
|
return apply
|
|
|
|
|
|
def find_k(
|
|
frequency: float,
|
|
tolerance: float,
|
|
direction: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
epsilon: fdfield,
|
|
mu: fdfield | None = None,
|
|
band: int = 0,
|
|
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,
|
|
v0: NDArray[numpy.complex128] | None = None,
|
|
) -> tuple[NDArray[numpy.float64], float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
|
|
"""
|
|
Search for a bloch vector that has a given frequency.
|
|
|
|
Args:
|
|
frequency: Target frequency.
|
|
tolerance: Target frequency tolerance.
|
|
direction: k-vector direction to search along.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
epsilon: Dielectric constant distribution for the simulation.
|
|
All fields are sampled at cell centers (i.e., NOT Yee-gridded)
|
|
mu: Magnetic permability distribution for the simulation.
|
|
Default None (1 everywhere).
|
|
band: Which band to search in. Default 0 (lowest frequency).
|
|
k_bounds: Minimum and maximum values for k. Default (0, 0.5).
|
|
k_guess: Initial value for k.
|
|
solve_callback: TODO
|
|
iter_callback: TODO
|
|
|
|
Returns:
|
|
`(k, actual_frequency, eigenvalues, eigenvectors)`
|
|
The found k-vector and its frequency, along with all eigenvalues and eigenvectors.
|
|
"""
|
|
direction = numpy.array(direction) / norm(direction) # type: ignore[operator]
|
|
|
|
k_bounds = tuple(sorted(k_bounds)) # type: ignore # we know the length already...
|
|
assert len(k_bounds) == 2
|
|
|
|
if k_guess is None:
|
|
k_guess = sum(k_bounds) / 2
|
|
|
|
n = None
|
|
v = v0
|
|
|
|
def get_f(k0_mag: float, band: int = 0) -> float:
|
|
nonlocal n, v
|
|
k0 = direction * k0_mag # type: ignore
|
|
n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu, y0=v, callback=iter_callback)
|
|
f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
|
|
if solve_callback:
|
|
solve_callback(k0_mag, n, v, f)
|
|
return f
|
|
|
|
res = scipy.optimize.minimize_scalar(
|
|
lambda x: abs(get_f(x, band) - frequency),
|
|
method='bounded',
|
|
bounds=k_bounds,
|
|
options={'xatol': abs(tolerance)},
|
|
)
|
|
|
|
assert n is not None
|
|
assert v is not None
|
|
actual_frequency = get_f(float(res.x), band)
|
|
return direction * float(res.x), float(actual_frequency), n, v # type: ignore[operator,return-value]
|
|
|
|
|
|
def eigsolve(
|
|
num_modes: int,
|
|
k0: ArrayLike,
|
|
G_matrix: ArrayLike,
|
|
epsilon: fdfield,
|
|
mu: fdfield | None = None,
|
|
tolerance: float = 1e-7,
|
|
max_iters: int = 10000,
|
|
reset_iters: int = 100,
|
|
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.
|
|
|
|
Args:
|
|
k0: Bloch wavevector, `[k0x, k0y, k0z]`.
|
|
G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
|
|
epsilon: Dielectric constant distribution for the simulation.
|
|
All fields are sampled at cell centers (i.e., NOT Yee-gridded)
|
|
mu: Finite positive real diagonal permeability. Default is one.
|
|
tolerance: Relative projected-gradient tolerance in the magnetic energy
|
|
metric. Flat objective traces trigger a conjugate-direction restart.
|
|
max_iters: Maximum iteration count; returns Ritz pairs from the final
|
|
subspace if the limit is reached.
|
|
reset_iters: Periodic conjugate-direction restart interval.
|
|
callback: Optional no-argument callback after each subspace update.
|
|
y0: Initial flux coefficients with shape (num_modes, 2 * grid_size).
|
|
Copied before use.
|
|
|
|
Returns:
|
|
`(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the
|
|
transverse flux vector `eigenvectors[i, :]`. Pass mu to the field
|
|
converters when it was supplied to this solve.
|
|
"""
|
|
k0 = numpy.asarray(k0)
|
|
|
|
h_size = 2 * epsilon[0].size
|
|
|
|
kmag = norm(G_matrix @ k0)
|
|
|
|
#
|
|
# Generate the operators
|
|
#
|
|
mop = maxwell_operator(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
|
|
imop = inverse_maxwell_operator_approx(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
|
|
|
|
scipy_op = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=mop)
|
|
scipy_iop = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=imop)
|
|
|
|
metric = inverse_metric = None
|
|
if mu is not None:
|
|
if numpy.iscomplexobj(mu) or not numpy.isfinite(mu).all() or numpy.any(mu <= 0):
|
|
raise ValueError('Bloch variational solving requires finite positive real permeability')
|
|
metric = spalg.LinearOperator(
|
|
shape=(h_size, h_size), dtype=complex,
|
|
matvec=_transverse_material_operator(k0, G_matrix, mu, inverse=True),
|
|
)
|
|
inverse_metric = spalg.LinearOperator(
|
|
shape=(h_size, h_size), dtype=complex,
|
|
matvec=_transverse_material_operator(k0, G_matrix, mu, inverse=False),
|
|
)
|
|
|
|
def mass(vectors: NDArray) -> NDArray:
|
|
# P mu^-1 P makes the flux operator self-adjoint in the magnetic
|
|
# energy inner product, even though it is not Euclidean-Hermitian.
|
|
return vectors if metric is None else metric @ vectors
|
|
|
|
y_shape = (h_size, num_modes)
|
|
|
|
prev_E = 0.0
|
|
d_scale = 1.0
|
|
prev_traceGtKG = 0.0
|
|
prev_theta = 0.5
|
|
D = numpy.zeros(shape=y_shape, dtype=complex)
|
|
|
|
rng = numpy.random.default_rng()
|
|
Z: NDArray[numpy.complex128]
|
|
if y0 is None:
|
|
Z = rng.random(y_shape) + 1j * rng.random(y_shape)
|
|
else:
|
|
Z = numpy.array(y0, dtype=complex, copy=True).T
|
|
|
|
while True:
|
|
Z *= num_modes / norm(Z)
|
|
ZtZ = Z.conj().T @ mass(Z)
|
|
try:
|
|
U = numpy.linalg.inv(ZtZ)
|
|
except numpy.linalg.LinAlgError:
|
|
Z = rng.random(y_shape) + 1j * rng.random(y_shape)
|
|
continue
|
|
|
|
trace_U = real(trace(U))
|
|
if trace_U > 1e8 * num_modes:
|
|
Z = Z @ scipy.linalg.sqrtm(U).astype(numpy.complex128).conj().T
|
|
prev_traceGtKG = 0
|
|
continue
|
|
break
|
|
|
|
Zt = numpy.empty(Z.shape[::-1], dtype=numpy.complex128)
|
|
AZ = numpy.empty(Z.shape, dtype=numpy.complex128)
|
|
|
|
for i in range(max_iters):
|
|
Zt = numpy.conj(Z.T, out=Zt)
|
|
ZtZ = Zt @ mass(Z)
|
|
U = numpy.linalg.inv(ZtZ)
|
|
AZ = scipy_op @ Z.copy()
|
|
ZtAZ = Zt @ mass(AZ)
|
|
ZtAZU = ZtAZ @ U
|
|
E_signed = real(trace(ZtAZU))
|
|
sgn = numpy.sign(E_signed)
|
|
E = numpy.abs(E_signed)
|
|
|
|
# G = AZU projected onto the space orthonormal to Z via (1 - ZUZt)
|
|
G = mass(AZ @ U - Z @ U @ ZtAZU) * sgn
|
|
|
|
if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7):
|
|
# A flat line search alone does not imply a stationary subspace.
|
|
# Restart with steepest descent if conjugacy has been lost.
|
|
prev_traceGtKG = 0
|
|
|
|
if norm(G) <= tolerance * max(norm(mass(AZ) @ U), 1e-7):
|
|
break
|
|
# Approximate the inverse energy Hessian by a symmetric sandwich of
|
|
# P mu P around the inverse-curl / epsilon preconditioner.
|
|
KG = scipy_iop @ (G if inverse_metric is None else inverse_metric @ G)
|
|
traceGtKG = _rtrace_AtB(G, KG)
|
|
|
|
if prev_traceGtKG == 0 or i % reset_iters == 0:
|
|
logger.info('CG reset')
|
|
gamma = 0.0
|
|
else:
|
|
gamma = traceGtKG / prev_traceGtKG
|
|
prev_traceGtKG = traceGtKG
|
|
|
|
D = gamma / d_scale * D + KG
|
|
d_scale = num_modes / norm(D)
|
|
D *= d_scale
|
|
|
|
# Now know the direction (D), but need to find how far to go (alpha)
|
|
# We are still minimizing E = tr((Z + alpha D)t A (Z + alpha D) U')
|
|
# = tr(ZtAZU' + alpha DtAZU' + alpha ZtADU' + alpha**2 DtADU')
|
|
# = tr((ZtAZ + 2 alpha sym(DtAZ) + alpha**2 DtAD) U')
|
|
# = tr(R U')
|
|
# = tr(R Qi) = tr(R inv(Q))
|
|
|
|
# where
|
|
# R = ZtAZ + 2 alpha sym(DtAZ) + alpha**2 DtAD
|
|
#
|
|
# Q = (Z + alpha D)t (Z + alpha D)
|
|
# = inv(ZtZ + alpha ZtD + alpha DtZ + alpha**2 DtD)
|
|
#
|
|
# Qi = inv(Q) = U'
|
|
|
|
AD = scipy_op @ D.copy()
|
|
DtD = D.conj().T @ mass(D)
|
|
DtAD = D.conj().T @ mass(AD)
|
|
|
|
symZtD = _symmetrize(Zt @ mass(D))
|
|
symZtAD = _symmetrize(Zt @ mass(AD))
|
|
|
|
Qi_memo: list[float | None] = [None, None]
|
|
|
|
def Qi_func(theta: float, Qi_memo=Qi_memo, ZtZ=ZtZ, DtD=DtD, symZtD=symZtD) -> float: # noqa: ANN001
|
|
if Qi_memo[0] == theta:
|
|
return cast('float', Qi_memo[1])
|
|
|
|
c = numpy.cos(theta)
|
|
s = numpy.sin(theta)
|
|
Q = c * c * ZtZ + s * s * DtD + 2 * s * c * symZtD
|
|
try:
|
|
Qi = numpy.linalg.inv(Q)
|
|
except numpy.linalg.LinAlgError as err:
|
|
logger.info('taylor Qi')
|
|
# if c or s small, taylor expand
|
|
if c < 1e-4 * s and c != 0:
|
|
DtDi = numpy.linalg.inv(DtD)
|
|
Qi = DtDi / (s * s) - 2 * c / (s * s * s) * (DtDi @ (DtDi @ symZtD).conj().T)
|
|
elif s < 1e-4 * c and s != 0:
|
|
ZtZi = numpy.linalg.inv(ZtZ)
|
|
Qi = ZtZi / (c * c) - 2 * s / (c * c * c) * (ZtZi @ (ZtZi @ symZtD).conj().T)
|
|
else:
|
|
raise Exception('Inexplicable singularity in trace_func') from err
|
|
Qi_memo[0] = theta
|
|
Qi_memo[1] = cast('float', Qi)
|
|
return cast('float', Qi)
|
|
|
|
def trace_func(theta: float, ZtAZ=ZtAZ, DtAD=DtAD, symZtAD=symZtAD) -> float: # noqa: ANN001
|
|
c = numpy.cos(theta)
|
|
s = numpy.sin(theta)
|
|
Qi = Qi_func(theta)
|
|
R = c * c * ZtAZ + s * s * DtAD + 2 * s * c * symZtAD
|
|
trace = _rtrace_AtB(R, Qi)
|
|
return numpy.abs(trace)
|
|
|
|
if False:
|
|
def trace_deriv(
|
|
theta: float,
|
|
sgn: int = sgn,
|
|
ZtAZ=ZtAZ, # noqa: ANN001
|
|
DtAD=DtAD, # noqa: ANN001
|
|
symZtD=symZtD, # noqa: ANN001
|
|
symZtAD=symZtAD, # noqa: ANN001
|
|
ZtZ=ZtZ, # noqa: ANN001
|
|
DtD=DtD, # noqa: ANN001
|
|
) -> float:
|
|
Qi = Qi_func(theta)
|
|
c2 = numpy.cos(2 * theta)
|
|
s2 = numpy.sin(2 * theta)
|
|
F = -0.5 * s2 * (ZtAZ - DtAD) + c2 * symZtAD
|
|
trace_deriv = _rtrace_AtB(Qi, F)
|
|
|
|
G = Qi @ F.conj().T @ Qi.conj().T
|
|
H = -0.5 * s2 * (ZtZ - DtD) + c2 * symZtD
|
|
trace_deriv -= _rtrace_AtB(G, H)
|
|
|
|
trace_deriv *= 2
|
|
return trace_deriv * sgn
|
|
|
|
U_sZtD = U @ symZtD
|
|
|
|
dE = 2.0 * (_rtrace_AtB(U, symZtAD)
|
|
- _rtrace_AtB(ZtAZU, U_sZtD))
|
|
|
|
d2E = 2 * (_rtrace_AtB(U, DtAD)
|
|
- _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD))
|
|
- 4 * _rtrace_AtB(U, symZtAD @ U_sZtD))
|
|
|
|
# Newton-Raphson to find a root of the first derivative:
|
|
theta = -dE / d2E
|
|
|
|
if d2E < 0 or abs(theta) >= pi:
|
|
theta = -abs(prev_theta) * numpy.sign(dE)
|
|
|
|
# theta, new_E, new_dE = linmin(theta, E, dE, 0.1, min(tolerance, 1e-6), 1e-14, 0, -numpy.sign(dE) * K_PI, trace_func)
|
|
theta, n, _, new_E, _, _new_dE = scipy.optimize.line_search(
|
|
trace_func,
|
|
trace_deriv,
|
|
xk=theta,
|
|
pk=numpy.ones((1, 1)),
|
|
gfk=dE,
|
|
old_fval=E,
|
|
c1=min(tolerance, 1e-6),
|
|
c2=0.1,
|
|
amax=pi,
|
|
)
|
|
|
|
result = scipy.optimize.minimize_scalar(
|
|
trace_func,
|
|
method='bounded',
|
|
# Center the period on zero so the line search can resolve small
|
|
# steps without the relative angular floor near theta = pi.
|
|
bounds=(-pi / 2, pi / 2),
|
|
options={'xatol': min(tolerance, 1e-12)},
|
|
)
|
|
new_E = result.fun
|
|
theta = result.x
|
|
|
|
improvement = numpy.abs(E - new_E) * 2 / numpy.abs(E + new_E)
|
|
logger.info(f'linmin improvement {improvement}')
|
|
Z *= numpy.cos(theta)
|
|
Z += D * numpy.sin(theta)
|
|
|
|
prev_theta = theta
|
|
prev_E = E
|
|
|
|
if callback:
|
|
callback()
|
|
|
|
#
|
|
# Recover eigenvectors from Z
|
|
#
|
|
U = numpy.linalg.inv(Z.conj().T @ mass(Z))
|
|
Y = Z @ scipy.linalg.sqrtm(U).astype(numpy.complex128)
|
|
W = Y.conj().T @ mass(scipy_op @ Y.copy())
|
|
|
|
eigvals, W_eigvecs = numpy.linalg.eig(W)
|
|
eigvecs = Y @ W_eigvecs
|
|
|
|
for i in range(len(eigvals)):
|
|
v = eigvecs[:, i]
|
|
n = eigvals[i]
|
|
v /= norm(v)
|
|
Av = numpy.asarray(scipy_op @ v.copy()).reshape(-1)
|
|
eigness = norm(Av - n * v)
|
|
f = numpy.sqrt(max(0.0, -numpy.real(n)))
|
|
df = numpy.sqrt(max(0.0, -numpy.real(n) + eigness))
|
|
neff_err = kmag * (1 / df - 1 / f) if f > 0 and df > 0 else numpy.inf
|
|
logger.info(f'eigness {i}: {eigness}\n neff_err: {neff_err}')
|
|
|
|
order = numpy.argsort(numpy.abs(eigvals))
|
|
return eigvals[order], eigvecs.T[order]
|
|
|
|
|
|
'''
|
|
def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_tol=1e-14, x_min=0, linmin_func):
|
|
if df0 > 0:
|
|
x0, f0, df0 = linmin(-x_guess, f0, -df0, -x_max, f_tol, df_tol, x_tol, -x_min, lambda q, dq: -linmin_func(q, dq))
|
|
return -x0, f0, -df0
|
|
elif df0 == 0:
|
|
return 0, f0, df0
|
|
else:
|
|
x = x_guess
|
|
fx = f0
|
|
dfx = df0
|
|
|
|
isave = numpy.zeros((2,), numpy.intc)
|
|
dsave = numpy.zeros((13,), float)
|
|
|
|
x, fx, dfx, task = minpack2.dsrch(x, fx, dfx, f_tol, df_tol, x_tol, task,
|
|
x_min, x_max, isave, dsave)
|
|
for i in range(int(1e6)):
|
|
if task != 'F':
|
|
logging.info(f'search converged in {i} iterations')
|
|
break
|
|
fx = f(x, dfx)
|
|
x, fx, dfx, task = minpack2.dsrch(x, fx, dfx, f_tol, df_tol, x_tol, task,
|
|
x_min, x_max, isave, dsave)
|
|
|
|
return x, fx, dfx
|
|
'''
|
|
|
|
def _rtrace_AtB(
|
|
A: NDArray[numpy.complex128],
|
|
B: NDArray[numpy.complex128] | float,
|
|
) -> float:
|
|
return real(numpy.sum(A.conj() * B))
|
|
|
|
def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
|
|
return (A + A.conj().T) * 0.5
|
|
|
|
|
|
|
|
def inner_product(
|
|
eL: cfdfield,
|
|
hL: cfdfield,
|
|
eR: cfdfield,
|
|
hR: cfdfield,
|
|
) -> complex:
|
|
# assumes x-axis propagation
|
|
|
|
assert numpy.array_equal(eR.shape, hR.shape)
|
|
assert numpy.array_equal(eL.shape, hL.shape)
|
|
assert numpy.array_equal(eR.shape, eL.shape)
|
|
|
|
# Cross product, times 2 since it's <p | n>, then divide by 4. # TODO might want to abs() this?
|
|
norm2R = (eR[1] * hR[2] - eR[2] * hR[1]).sum() / 2
|
|
norm2L = (eL[1] * hL[2] - eL[2] * hL[1]).sum() / 2
|
|
|
|
# eRxhR_x = numpy.cross(eR.reshape(3, -1), hR.reshape(3, -1), axis=0).reshape(eR.shape)[0] / normR
|
|
# logger.info(f'power {eRxhR_x.sum() / 2})
|
|
|
|
eR_norm = eR / numpy.sqrt(abs(norm2R))
|
|
hR_norm = hR / numpy.sqrt(abs(norm2R))
|
|
eL_norm = eL / numpy.sqrt(abs(norm2L))
|
|
hL_norm = hL / numpy.sqrt(abs(norm2L))
|
|
|
|
# (eR x hL)[0] and (eL x hR)[0]
|
|
eRxhL_x = eR_norm[1] * hL_norm[2] - eR_norm[2] * hL_norm[1]
|
|
eLxhR_x = eL_norm[1] * hR_norm[2] - eL_norm[2] * hR_norm[1]
|
|
|
|
#return 1j * (eRxhL_x - eLxhR_x).sum() / numpy.sqrt(norm2R * norm2L)
|
|
#return (eRxhL_x.sum() - eLxhR_x.sum()) / numpy.sqrt(norm2R * norm2L)
|
|
return eLxhR_x.sum() - eRxhL_x.sum()
|
|
|
|
|
|
def trq(
|
|
eI: cfdfield,
|
|
hI: cfdfield,
|
|
eO: cfdfield,
|
|
hO: cfdfield,
|
|
) -> tuple[complex, complex]:
|
|
pp = inner_product(eO, hO, eI, hI)
|
|
pn = inner_product(eO, hO, eI, -hI)
|
|
np = inner_product(eO, -hO, eI, hI)
|
|
nn = inner_product(eO, -hO, eI, -hI)
|
|
|
|
assert numpy.allclose(pp, -nn, atol=1e-12, rtol=1e-12)
|
|
assert numpy.allclose(pn, -np, atol=1e-12, rtol=1e-12)
|
|
|
|
logger.info(f'''
|
|
{pp=:4g} {pn=:4g}
|
|
{nn=:4g} {np=:4g}
|
|
{nn * pp / pn=:4g} {-np=:4g}
|
|
''')
|
|
|
|
r = -pp / pn # -<Pp|Bp>/<Pn/Bp> = -(-pp) / (-pn)
|
|
t = (np - nn * pp / pn) / 4
|
|
|
|
return t, r
|