fixup! add sensitivity calculation

examples/fdtd.py

@@ -157,7 +157,8 @@ def main():
         e[1][tuple(grid.shape//2)] += field_source(t)
         update_H(e, h)
 
-        print('iteration {}: average {} iterations per sec'.format(t, (t+1)/(time.perf_counter()-start)))
+        avg_rate = (t + 1)/(time.perf_counter() - start))
+        print(f'iteration {t}: average {avg_rate} iterations per sec')
         sys.stdout.flush()
 
         if t % 20 == 0:

examples/tcyl.py

@@ -3,7 +3,7 @@ import numpy
 from numpy.linalg import norm
 
 from meanas.fdmath import vec, unvec
-from meanas.fdfd import waveguide_mode, functional, scpml
+from meanas.fdfd import waveguide_cyl, functional, scpml
 from meanas.fdfd.solvers import generic as generic_solver
 
 import gridlock
@@ -37,29 +37,34 @@ def test1(solver=generic_solver):
     xyz_max = numpy.array([800, y_max, z_max]) + (pml_thickness + 2) * dx
 
     # Coordinates of the edges of the cells.
-    half_edge_coords = [numpy.arange(dx/2, m + dx/2, step=dx) for m in xyz_max]
+    half_edge_coords = [numpy.arange(dx / 2, m + dx / 2, step=dx) for m in xyz_max]
     edge_coords = [numpy.hstack((-h[::-1], h)) for h in half_edge_coords]
     edge_coords[0] = numpy.array([-dx, dx])
 
     # #### Create the grid and draw the device ####
     grid = gridlock.Grid(edge_coords)
     epsilon = grid.allocate(n_air**2, dtype=numpy.float32)
-    grid.draw_cuboid(epsilon, center=center, dimensions=[8e3, w, th], eps=n_wg**2)
+    grid.draw_cuboid(epsilon, center=center, dimensions=[8e3, w, th], foreground=n_wg**2)
 
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (1, 2):
         for p in (-1, 1):
-            dxes = scmpl.stretch_with_scpml(dxes, omega=omega, axis=a, polarity=p,
+            dxes = scpml.stretch_with_scpml(
-                                            thickness=pml_thickness)
+                dxes,
+                omega=omega,
+                axis=a,
+                polarity=p,
+                thickness=pml_thickness,
+                )
 
     wg_args = {
         'omega': omega,
         'dxes': [(d[1], d[2]) for d in dxes],
-        'epsilon': vec(g.transpose([1, 2, 0]) for g in epsilon),
+        'epsilon': vec(epsilon.transpose([0, 2, 3, 1])),
         'r0': r0,
     }
 
-    wg_results = waveguide_mode.solve_waveguide_mode_cylindrical(mode_number=0, **wg_args)
+    wg_results = waveguide_cyl.solve_mode(mode_number=0, **wg_args)
 
     E = wg_results['E']
 
@@ -70,20 +75,17 @@ def test1(solver=generic_solver):
     '''
     Plot results
     '''
-    def pcolor(v):
+    def pcolor(fig, ax, v, title):
         vmax = numpy.max(numpy.abs(v))
-        pyplot.pcolor(v.T, cmap='seismic', vmin=-vmax, vmax=vmax)
+        mappable = ax.pcolormesh(v.T, cmap='seismic', vmin=-vmax, vmax=vmax)
-        pyplot.axis('equal')
+        ax.set_aspect('equal', adjustable='box')
-        pyplot.colorbar()
+        ax.set_title(title)
+        ax.figure.colorbar(mappable)
 
-    pyplot.figure()
+    fig, axes = pyplot.subplots(2, 2)
-    pyplot.subplot(2, 2, 1)
+    pcolor(fig, axes[0][0], numpy.real(E[0]), 'Ex')
-    pcolor(numpy.real(E[0][:, :]))
+    pcolor(fig, axes[0][1], numpy.real(E[1]), 'Ey')
-    pyplot.subplot(2, 2, 2)
+    pcolor(fig, axes[1][0], numpy.real(E[2]), 'Ez')
-    pcolor(numpy.real(E[1][:, :]))
-    pyplot.subplot(2, 2, 3)
-    pcolor(numpy.real(E[2][:, :]))
-    pyplot.subplot(2, 2, 4)
     pyplot.show()
 
 

meanas/fdfd/bloch.py

@@ -684,11 +684,11 @@ def eigsolve(
                 Qi = Qi_func(theta)
                 c2 = numpy.cos(2 * theta)
                 s2 = numpy.sin(2 * theta)
-                F = -0.5*s2 *  (ZtAZ - DtAD) + c2 * symZtAD
+                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
+                H = -0.5 * s2 * (ZtZ - DtD) + c2 * symZtD
                 trace_deriv -= _rtrace_AtB(G, H)
 
                 trace_deriv *= 2
@@ -696,12 +696,12 @@ def eigsolve(
 
             U_sZtD = U @ symZtD
 
-            dE = 2.0 * (_rtrace_AtB(U, symZtAD) -
+            dE = 2.0 * (_rtrace_AtB(U, symZtAD)
-                        _rtrace_AtB(ZtAZU, U_sZtD))
+                        - _rtrace_AtB(ZtAZU, U_sZtD))
 
-            d2E = 2 * (_rtrace_AtB(U, DtAD) -
+            d2E = 2 * (_rtrace_AtB(U, DtAD)
-                       _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD)) -
+                       - _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD))
-                   4 * _rtrace_AtB(U, symZtAD @ U_sZtD))
+                       - 4 * _rtrace_AtB(U, symZtAD @ U_sZtD))
 
             # Newton-Raphson to find a root of the first derivative:
             theta = -dE / d2E
@@ -781,7 +781,7 @@ def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_to
                                           x_min, x_max, isave, dsave)
         for i in range(int(1e6)):
             if task != 'F':
-                logging.info('search converged in {} iterations'.format(i))
+                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,

meanas/fdfd/solvers.py

@@ -43,7 +43,8 @@ def _scipy_qmr(
         nonlocal ii
         ii += 1
         if ii % 100 == 0:
-            logger.info('Solver residual at iteration {} : {}'.format(ii, norm(A @ xk - b)))
+            cur_norm = norm(A @ xk - b)
+            logger.info(f'Solver residual at iteration {ii} : {cur_norm}')
 
     if 'callback' in kwargs:
         def augmented_callback(xk: ArrayLike) -> None:

meanas/fdfd/waveguide_2d.py

@@ -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, vcfdfield_t
+from ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, vcfdfield_t
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
 
@@ -253,7 +253,8 @@ def operator_e(
     mu_yx = sparse.diags(numpy.hstack((mu_parts[1], mu_parts[0])))
     mu_z_inv = sparse.diags(1 / mu_parts[2])
 
-    op = (omega * omega * mu_yx @ eps_xy
+    op = (
+        omega * omega * mu_yx @ eps_xy
         + mu_yx @ sparse.vstack((-Dby, Dbx)) @ mu_z_inv @ sparse.hstack((-Dfy, Dfx))
         + sparse.vstack((Dfx, Dfy)) @ eps_z_inv @ sparse.hstack((Dbx, Dby)) @ eps_xy
         )
@@ -321,7 +322,8 @@ def operator_h(
     mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1])))
     mu_z_inv = sparse.diags(1 / mu_parts[2])
 
-    op = (omega * omega * eps_yx @ mu_xy
+    op = (
+        omega * omega * eps_yx @ mu_xy
         + eps_yx @ sparse.vstack((-Dfy, Dfx)) @ eps_z_inv @ sparse.hstack((-Dby, Dbx))
         + sparse.vstack((Dbx, Dby)) @ mu_z_inv @ sparse.hstack((Dfx, Dfy)) @ mu_xy
         )
@@ -420,7 +422,7 @@ def _normalized_fields(
     Sz_a = E[0] * numpy.conj(H[1] * phase) * dxes_real[0][1] * dxes_real[1][0]
     Sz_b = E[1] * numpy.conj(H[0] * phase) * dxes_real[0][0] * dxes_real[1][1]
     Sz_tavg = numpy.real(Sz_a.sum() - Sz_b.sum()) * 0.5       # 0.5 since E, H are assumed to be peak (not RMS) amplitudes
-    assert Sz_tavg > 0, 'Found a mode propagating in the wrong direction! Sz_tavg={}'.format(Sz_tavg)
+    assert Sz_tavg > 0, f'Found a mode propagating in the wrong direction! {Sz_tavg=}'
 
     energy = epsilon * e.conj() * e
 
@@ -718,6 +720,109 @@ def e_err(
     return float(norm(op) / norm(e))
 
 
+def sensitivity(
+        e_norm: vcfdfield_t,
+        h_norm: vcfdfield_t,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfdfield_t,
+        mu: vfdfield_t | None = None,
+        ) -> vcfdfield_t:
+    r"""
+      Given a waveguide structure (`dxes`, `epsilon`, `mu`) and mode fields
+    (`e_norm`, `h_norm`, `wavenumber`, `omega`), calculates the sensitivity of the wavenumber
+    $\beta$ to changes in the dielectric structure $\epsilon$.
+
+    The output is a vector of the same size as `vec(epsilon)`, with each element specifying the
+    sensitivity of `wavenumber` to changes in the corresponding element in `vec(epsilon)`, i.e.
+
+    $$sens_{i} = \frac{\partial\beta}{\partial\epsilon_i}$$
+
+    An adjoint approach is used to calculate the sensitivity; the derivation is provided here:
+
+    Starting with the eigenvalue equation
+
+    $$\beta^2 E_{xy} = A_E E_{xy}$$
+
+    where $A_E$ is the waveguide operator from `operator_e()`, and $E_{xy} = \begin{bmatrix} E_x \\
+                                                                                             E_y \end{bmatrix}$,
+    we can differentiate with respect to one of the $\epsilon$ elements (i.e. at one Yee grid point), $\epsilon_i$:
+
+    $$
+    (2 \beta) \partial_{\epsilon_i}(\beta) E_{xy} + \beta^2 \partial_{\epsilon_i} E_{xy}
+        = \partial_{\epsilon_i}(A_E) E_{xy} + A_E \partial_{\epsilon_i} E_{xy}
+    $$
+
+    We then multiply by $H_{yx}^\star = \begin{bmatrix}H_y^\star \\ -H_x^\star \end{bmatrix}$ from the left:
+
+    $$
+    (2 \beta) \partial_{\epsilon_i}(\beta) H_{yx}^\star E_{xy} + \beta^2 H_{yx}^\star \partial_{\epsilon_i} E_{xy}
+        = H_{yx}^\star \partial_{\epsilon_i}(A_E) E_{xy} + H_{yx}^\star A_E \partial_{\epsilon_i} E_{xy}
+    $$
+
+    However, $H_{yx}^\star$ is actually a left-eigenvector of $A_E$. This can be verified by inspecting
+    the form of `operator_h` ($A_H$) and comparing its conjugate transpose to `operator_e` ($A_E$). Also, note
+    $H_{yx}^\star \cdot E_{xy} = H^\star \times E$ recalls the mode orthogonality relation. See doi:10.5194/ars-9-85-201
+    for a similar approach. Therefore,
+
+    $$
+    H_{yx}^\star A_E \partial_{\epsilon_i} E_{xy} = \beta^2 H_{yx}^\star \partial_{\epsilon_i} E_{xy}
+    $$
+
+    and we can simplify to
+
+    $$
+    \partial_{\epsilon_i}(\beta)
+        = \frac{1}{2 \beta} \frac{H_{yx}^\star \partial_{\epsilon_i}(A_E) E_{xy} }{H_{yx}^\star E_{xy}}
+    $$
+
+    This expression can be quickly calculated for all $i$ by writing out the various terms of
+    $\partial_{\epsilon_i} A_E$ and recognizing that the vector-matrix-vector products (i.e. scalars)
+    $sens_i = \vec{v}_{left} \partial_{\epsilon_i} (\epsilon_{xyz}) \vec{v}_{right}$, indexed by $i$, can be expressed as
+    elementwise multiplications $\vec{sens} = \vec{v}_{left} \star \vec{v}_{right}$
+
+
+    Args:
+        e_norm: Normalized, vectorized E_xyz field for the mode. E.g. as returned by `normalized_fields_e`.
+        h_norm: Normalized, vectorized H_xyz field for the mode. E.g. as returned by `normalized_fields_e`.
+        wavenumber: Propagation constant for the mode. The z-axis is assumed to be continuous (i.e. without numerical dispersion).
+        omega: The angular frequency of the system.
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representation of the operator.
+    """
+    if mu is None:
+        mu = numpy.ones_like(epsilon)
+
+    Dfx, Dfy = deriv_forward(dxes[0])
+    Dbx, Dby = deriv_back(dxes[1])
+
+    eps_x, eps_y, eps_z = numpy.split(epsilon, 3)
+    eps_xy = sparse.diags(numpy.hstack((eps_x, eps_y)))
+    eps_z_inv = sparse.diags(1 / eps_z)
+
+    mu_x, mu_y, _mu_z = numpy.split(mu, 3)
+    mu_yx = sparse.diags(numpy.hstack((mu_y, mu_x)))
+
+    da_exxhyy = vec(dxes[1][0][:, None] * dxes[0][1][None, :])
+    da_eyyhxx = vec(dxes[1][1][None, :] * dxes[0][0][:, None])
+    ev_xy = numpy.concatenate(numpy.split(e_norm, 3)[:2]) * numpy.concatenate([da_exxhyy, da_eyyhxx])
+    hx, hy, hz = numpy.split(h_norm, 3)
+    hv_yx_conj = numpy.conj(numpy.concatenate([hy, -hx]))
+
+    sens_xy1 = (hv_yx_conj @ (omega * omega * mu_yx)) * ev_xy
+    sens_xy2 = (hv_yx_conj @ sparse.vstack((Dfx, Dfy)) @ eps_z_inv @ sparse.hstack((Dbx, Dby))) * ev_xy
+    sens_z   = (hv_yx_conj @ sparse.vstack((Dfx, Dfy)) @ (-eps_z_inv * eps_z_inv)) * (sparse.hstack((Dbx, Dby)) @ eps_xy @ ev_xy)
+    norm = hv_yx_conj @ ev_xy
+
+    sens_tot = numpy.concatenate([sens_xy1 + sens_xy2, sens_z]) / (2 * wavenumber * norm)
+    return sens_tot
+
+
 def solve_modes(
         mode_numbers: list[int],
         omega: complex,

meanas/fdfd/waveguide_cyl.py

@@ -25,6 +25,9 @@ def cylindrical_operator(
     """
     Cylindrical coordinate waveguide operator of the form
 
+    (NOTE: See 10.1364/OL.33.001848)
+    TODO: consider 10.1364/OE.20.021583
+
     TODO
 
     for use with a field vector of the form `[E_r, E_y]`.

meanas/fdmath/operators.py

@@ -29,9 +29,9 @@ def shift_circ(
         Sparse matrix for performing the circular shift.
     """
     if len(shape) not in (2, 3):
-        raise Exception('Invalid shape: {}'.format(shape))
+        raise Exception(f'Invalid shape: {shape}')
     if axis not in range(len(shape)):
-        raise Exception('Invalid direction: {}, shape is {}'.format(axis, shape))
+        raise Exception(f'Invalid direction: {axis}, shape is {shape}')
 
     shifts = [abs(shift_distance) if a == axis else 0 for a in range(3)]
     shifted_diags = [(numpy.arange(n) + s) % n for n, s in zip(shape, shifts)]
@@ -69,12 +69,11 @@ def shift_with_mirror(
         Sparse matrix for performing the shift-with-mirror.
     """
     if len(shape) not in (2, 3):
-        raise Exception('Invalid shape: {}'.format(shape))
+        raise Exception(f'Invalid shape: {shape}')
     if axis not in range(len(shape)):
-        raise Exception('Invalid direction: {}, shape is {}'.format(axis, shape))
+        raise Exception(f'Invalid direction: {axis}, shape is {shape}')
     if shift_distance >= shape[axis]:
-        raise Exception('Shift ({}) is too large for axis {} of size {}'.format(
+        raise Exception(f'Shift ({shift_distance}) is too large for axis {axis} of size {shape[axis]}')
-                        shift_distance, axis, shape[axis]))
 
     def mirrored_range(n: int, s: int) -> NDArray[numpy.int_]:
         v = numpy.arange(n) + s
@@ -198,7 +197,7 @@ def avg_forward(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
         Sparse matrix for forward average operation.
     """
     if len(shape) not in (2, 3):
-        raise Exception('Invalid shape: {}'.format(shape))
+        raise Exception(f'Invalid shape: {shape}')
 
     n = numpy.prod(shape)
     return 0.5 * (sparse.eye(n) + shift_circ(axis, shape))

meanas/fdtd/boundaries.py

@@ -15,13 +15,17 @@ def conducting_boundary(
         ) -> tuple[fdfield_updater_t, fdfield_updater_t]:
     dirs = [0, 1, 2]
     if direction not in dirs:
-        raise Exception('Invalid direction: {}'.format(direction))
+        raise Exception(f'Invalid direction: {direction}')
     dirs.remove(direction)
     u, v = dirs
 
+    boundary_slice: list[Any]
+    shifted1_slice: list[Any]
+    shifted2_slice: list[Any]
+
     if polarity < 0:
-        boundary_slice = [slice(None)] * 3      # type: list[Any]
+        boundary_slice = [slice(None)] * 3
-        shifted1_slice = [slice(None)] * 3      # type: list[Any]
+        shifted1_slice = [slice(None)] * 3
         boundary_slice[direction] = 0
         shifted1_slice[direction] = 1
 
@@ -42,7 +46,7 @@ def conducting_boundary(
     if polarity > 0:
         boundary_slice = [slice(None)] * 3
         shifted1_slice = [slice(None)] * 3
-        shifted2_slice = [slice(None)] * 3      # type: list[Any]
+        shifted2_slice = [slice(None)] * 3
         boundary_slice[direction] = -1
         shifted1_slice[direction] = -2
         shifted2_slice[direction] = -3
@@ -64,4 +68,4 @@ def conducting_boundary(
 
         return ep, hp
 
-    raise Exception('Bad polarity: {}'.format(polarity))
+    raise Exception(f'Bad polarity: {polarity}')

meanas/fdtd/pml.py

@@ -33,10 +33,10 @@ def cpml_params(
         ) -> dict[str, Any]:
 
     if axis not in range(3):
-        raise Exception('Invalid axis: {}'.format(axis))
+        raise Exception(f'Invalid axis: {axis}')
 
     if polarity not in (-1, 1):
-        raise Exception('Invalid polarity: {}'.format(polarity))
+        raise Exception(f'Invalid polarity: {polarity}')
 
     if thickness <= 2:
         raise Exception('It would be wise to have a pml with 4+ cells of thickness')

meanas/test/test_fdtd.py

@@ -101,7 +101,7 @@ def test_poynting_divergence(sim: 'TDResult') -> None:
 def test_poynting_planes(sim: 'TDResult') -> None:
     mask = (sim.js[0] != 0).any(axis=0)
     if mask.sum() > 1:
-        pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum()))
+        pytest.skip(f'test_poynting_planes can only test single point sources, got {mask.sum()}')
 
     args: dict[str, Any] = {
         'dxes': sim.dxes,