Add linear_wavenumbers() for calculating 1/distance wavenumbers

examples/fdfd.py

@@ -46,20 +46,24 @@ def test0(solver=generic_solver):
     # #### Create the grid, mask, and draw the device ####
     grid = gridlock.Grid(edge_coords)
     epsilon = grid.allocate(n_air**2, dtype=numpy.float32)
-    grid.draw_cylinder(epsilon,
-                       surface_normal=2,
-                       center=center,
-                       radius=max(radii),
-                       thickness=th,
-                       eps=n_ring**2,
-                       num_points=24)
-    grid.draw_cylinder(epsilon,
-                       surface_normal=2,
-                       center=center,
-                       radius=min(radii),
-                       thickness=th*1.1,
-                       eps=n_air ** 2,
-                       num_points=24)
+    grid.draw_cylinder(
+        epsilon,
+        surface_normal=2,
+        center=center,
+        radius=max(radii),
+        thickness=th,
+        foreground=n_ring**2,
+        num_points=24,
+        )
+    grid.draw_cylinder(
+        epsilon,
+        surface_normal=2,
+        center=center,
+        radius=min(radii),
+        thickness=th*1.1,
+        foreground=n_air ** 2,
+        num_points=24,
+        )
 
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (0, 1, 2):
@@ -71,9 +75,9 @@ def test0(solver=generic_solver):
     J[1][15, grid.shape[1]//2, grid.shape[2]//2] = 1
 
 
-    '''
-    Solve!
-    '''
+    #
+    # Solve!
+    #
     sim_args = {
         'omega': omega,
         'dxes': dxes,
@@ -87,9 +91,9 @@ def test0(solver=generic_solver):
 
     E = unvec(x, grid.shape)
 
-    '''
-    Plot results
-    '''
+    #
+    # Plot results
+    #
     pyplot.figure()
     pyplot.pcolor(numpy.real(E[1][:, :, grid.shape[2]//2]), cmap='seismic')
     pyplot.axis('equal')
@@ -122,7 +126,7 @@ def test1(solver=generic_solver):
     # #### 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 (0, 1, 2):
@@ -169,9 +173,9 @@ def test1(solver=generic_solver):
 #    pcolor((numpy.abs(J3).sum(axis=2).sum(axis=0) > 0).astype(float).T)
     pyplot.show(block=True)
 
-    '''
-    Solve!
-    '''
+    #
+    # Solve!
+    #
     sim_args = {
         'omega': omega,
         'dxes': dxes,
@@ -188,9 +192,9 @@ def test1(solver=generic_solver):
 
     E = unvec(x, grid.shape)
 
-    '''
-    Plot results
-    '''
+    #
+    # Plot results
+    #
     center = grid.pos2ind([0, 0, 0], None).astype(int)
     pyplot.figure()
     pyplot.subplot(2, 2, 1)
@@ -232,7 +236,7 @@ def test1(solver=generic_solver):
     pyplot.grid(alpha=0.6)
     pyplot.title('Overlap with mode')
     pyplot.show()
-    print('Average overlap with mode:', sum(q)/len(q))
+    print('Average overlap with mode:', sum(q[8:32])/len(q[8:32]))
 
 
 def module_available(name):

meanas/fdfd/waveguide_2d.py

@@ -414,18 +414,13 @@ def _normalized_fields(
     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)]
 
-    E = unvec(e, shape)
-    H = unvec(h, shape)
-
     # Find time-averaged Sz and normalize to it
     # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting
     phase = numpy.exp(-1j * -prop_phase / 2)
-    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
+    Sz_tavg = inner_product(e, h, dxes=dxes, prop_phase=prop_phase, conj_h=True).real
     assert Sz_tavg > 0, f'Found a mode propagating in the wrong direction! {Sz_tavg=}'
 
-    energy = epsilon * e.conj() * e
+    energy = numpy.real(epsilon * e.conj() * e)
 
     norm_amplitude = 1 / numpy.sqrt(Sz_tavg)
     norm_angle = -numpy.angle(e[energy.argmax()])       # Will randomly add a negative sign when mode is symmetric
@@ -435,6 +430,7 @@ def _normalized_fields(
     sign = numpy.sign(E_weighted[:,
                                  :max(shape[0] // 2, 1),
                                  :max(shape[1] // 2, 1)].real.sum())
+    assert sign != 0
 
     norm_factor = sign * norm_amplitude * numpy.exp(1j * norm_angle)
 
@@ -825,7 +821,7 @@ def sensitivity(
 
 
 def solve_modes(
-        mode_numbers: list[int],
+        mode_numbers: Sequence[int],
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
@@ -858,7 +854,9 @@ def solve_modes(
     A_r = operator_e(numpy.real(omega), dxes_real, numpy.real(epsilon), mu_real)
 
     eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
-    e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 1)]
+    keep_inds = -(numpy.array(mode_numbers) + 1)
+    e_xys = eigvecs[:, keep_inds].T
+    eigvals = eigvals[keep_inds]
 
     #
     # Now solve for the eigenvector of the full operator, using the real operator's
@@ -866,13 +864,17 @@ def solve_modes(
     #
     A = operator_e(omega, dxes, epsilon, mu)
     for nn in range(len(mode_numbers)):
-        eigvals[nn], e_xys[:, nn] = rayleigh_quotient_iteration(A, e_xys[:, nn])
+        eigvals[nn], e_xys[nn, :] = rayleigh_quotient_iteration(A, e_xys[nn, :])
 
     # Calculate the wave-vector (force the real part to be positive)
     wavenumbers = numpy.sqrt(eigvals)
     wavenumbers *= numpy.sign(numpy.real(wavenumbers))
 
-    return e_xys.T, wavenumbers
+    order = wavenumbers.argsort()[::-1]
+    e_xys = e_xys[order]
+    wavenumbers = wavenumbers[order]
+
+    return e_xys, wavenumbers
 
 
 def solve_mode(
@@ -894,3 +896,37 @@ def solve_mode(
     kwargs['mode_numbers'] = [mode_number]
     e_xys, wavenumbers = solve_modes(*args, **kwargs)
     return e_xys[0], wavenumbers[0]
+
+
+def inner_product(    # TODO documentation
+        e1: vcfdfield_t,
+        h2: vcfdfield_t,
+        dxes: dx_lists_t,
+        prop_phase: float = 0,
+        conj_h: bool = False,
+        trapezoid: bool = False,
+        ) -> tuple[vcfdfield_t, vcfdfield_t]:
+
+    shape = [s.size for s in dxes[0]]
+
+    # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting
+    phase = numpy.exp(-1j * -prop_phase / 2)
+
+    E1 = unvec(e1, shape)
+    H2 = unvec(h2, shape) * phase
+
+    if conj_h:
+        H2 = numpy.conj(H2)
+
+    # Find time-averaged Sz and normalize to it
+    dxes_real = [[numpy.real(dxyz) for dxyz in dxeh] for dxeh in dxes]
+    if integrate:
+        Sz_a = numpy.trapezoid(numpy.trapezoid(E1[0] * H2[1], numpy.cumsum(dxes_real[0][1])), numpy.cumsum(dxes_real[1][0]))
+        Sz_b = numpy.trapezoid(numpy.trapezoid(E1[1] * H2[0], numpy.cumsum(dxes_real[0][0])), numpy.cumsum(dxes_real[1][1]))
+    else:
+        Sz_a = E1[0] * H2[1] * dxes_real[1][0][:, None] * dxes_real[0][1][None, :]
+        Sz_b = E1[1] * H2[0] * dxes_real[0][0][:, None] * dxes_real[1][1][None, :]
+    Sz = 0.5 * (Sz_a.sum() - Sz_b.sum())
+    return Sz
+
+

meanas/fdfd/waveguide_cyl.py

@@ -10,6 +10,7 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
 
 from typing import Any
 from collections.abc import Sequence
+import logging
 
 import numpy
 from numpy.typing import NDArray, ArrayLike
@@ -20,11 +21,13 @@ from ..fdmath.operators import deriv_forward, deriv_back
 from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
 
 
+logger = logging.getLogger(__name__)
+
+
 def cylindrical_operator(
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
-        r0: float,
         rmin: float,
         ) -> sparse.spmatrix:
     """
@@ -47,8 +50,7 @@ def cylindrical_operator(
         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
-        r0: Radius of curvature at x=0
-        rmin: Radius at the left edge of the simulation domain
+        rmin: Radius at the left edge of the simulation domain (minimum 'x')
 
     Returns:
         Sparse matrix representation of the operator
@@ -57,13 +59,7 @@ def cylindrical_operator(
     Dfx, Dfy = deriv_forward(dxes[0])
     Dbx, Dby = deriv_back(dxes[1])
 
-    ra = rmin + dxes[0][0] / 2.0 + numpy.cumsum(dxes[1][0])   # Radius at Ex points
-    rb = rmin + numpy.cumsum(dxes[0][0])                      # Radius at Ey points
-    ta = ra / r0
-    tb = rb / r0
-
-    Ta = sparse.diags(vec(ta[:, None].repeat(dxes[0][1].size, axis=1)))
-    Tb = sparse.diags(vec(tb[:, None].repeat(dxes[1][1].size, axis=1)))
+    Ta, Tb = dxes2T(dxes=dxes, rmin=rmin)
 
     eps_parts = numpy.split(epsilon, 3)
     eps_x = sparse.diags(eps_parts[0])
@@ -99,10 +95,9 @@ def solve_modes(
         omega: complex,
         dxes: dx_lists_t,
         epsilon: vfdfield_t,
-        r0: float,
         rmin: float,
         mode_margin: int = 2,
-        ) -> tuple[vcfdfield_t, NDArray[numpy.complex64]]:
+        ) -> tuple[vcfdfield_t, NDArray[numpy.complex128]]:
     """
     TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
@@ -114,12 +109,12 @@ def solve_modes(
         dxes: Grid parameters [dx_e, dx_h] as described in meanas.fdmath.types.
               The first coordinate is assumed to be r, the second is y.
         epsilon: Dielectric constant
-        r0: Radius of curvature for the simulation. This should be the minimum value of
-            r within the simulation domain.
+        rmin: Radius of curvature for the simulation. This should be the minimum value of
+               r within the simulation domain.
 
     Returns:
         e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
-        wavenumbers: list of wavenumbers
+        angular_wavenumbers: list of wavenumbers in 1/rad units.
     """
 
     #
@@ -127,15 +122,17 @@ def solve_modes(
     #
     dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
 
-    A_r = cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0=r0, rmin=rmin)
+    A_r = cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), rmin=rmin)
     eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
-    e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 1)].T
+    keep_inds = -(numpy.array(mode_numbers) + 1)
+    e_xys = eigvecs[:, keep_inds].T
+    eigvals = eigvals[keep_inds]
 
     #
     # Now solve for the eigenvector of the full operator, using the real operator's
     #  eigenvector as an initial guess for Rayleigh quotient iteration.
     #
-    A = cylindrical_operator(omega, dxes, epsilon, r0=r0, rmin=rmin)
+    A = cylindrical_operator(omega, dxes, epsilon, rmin=rmin)
     for nn in range(len(mode_numbers)):
         eigvals[nn], e_xys[nn, :] = rayleigh_quotient_iteration(A, e_xys[nn, :])
 
@@ -143,7 +140,15 @@ def solve_modes(
     wavenumbers = numpy.sqrt(eigvals)
     wavenumbers *= numpy.sign(numpy.real(wavenumbers))
 
-    return e_xys, wavenumbers
+    # Wavenumbers assume the mode is at rmin, which is unlikely
+    # Instead, return the wavenumber in inverse radians
+    angular_wavenumbers = wavenumbers * rmin
+
+    order = angular_wavenumbers.argsort()[::-1]
+    e_xys = e_xys[order]
+    angular_wavenumbers = angular_wavenumbers[order]
+
+    return e_xys, angular_wavenumbers
 
 
 def solve_mode(
@@ -160,8 +165,49 @@ def solve_mode(
        **kwargs: passed to `solve_modes()`
 
     Returns:
-        (e_xy, wavenumber)
+        (e_xy, angular_wavenumber)
     """
     kwargs['mode_numbers'] = [mode_number]
     e_xys, wavenumbers = solve_modes(*args, **kwargs)
-    return e_xys[0], wavenumbers[0]
+    return e_xys[0], angular_wavenumbers[0]
+
+
+def linear_wavenumbers(
+        e_xys: vcfdfield_t,
+        angular_wavenumbers: ArrayLike,
+        epsilon: vfdfield_t,
+        dxes: dx_lists_t,
+        rmin: float,
+        ) -> NDArray[numpy.complex128]:
+    """
+    Calculate linear wavenumbers (1/distance) based on angular wavenumbers (1/rad)
+      and the mode's energy distribution.
+
+    Args:
+        e_xys: Vectorized mode fields with shape [num_modes, 2 * x *y)
+        angular_wavenumbers: Angular wavenumbers corresponding to the fields in `e_xys`
+        epsilon: Vectorized dielectric constant grid with shape (3, x, y)
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
+        rmin: Radius at the left edge of the simulation domain (minimum 'x')
+
+    Returns:
+        NDArray containing the calculated linear (1/distance) wavenumbers
+    """
+    angular_wavenumbers = numpy.asarray(angular_wavenumbers)
+    mode_radii = numpy.empty_like(angular_wavenumbers, dtype=float)
+
+    wavenumbers = numpy.empty_like(angular_wavenumbers)
+    shape2d = (len(dxes[0][0]), len(dxes[0][1]))
+    epsilon2d = unvec(epsilon, shape2d)[:2]
+    grid_radii = rmin + numpy.cumsum(dxes[0][0])
+    for ii in range(angular_wavenumbers.size):
+        efield = unvec(e_xys[ii], shape2d, 2)
+        energy = numpy.real((efield * efield.conj()) * epsilon2d)
+        energy_vs_x = energy.sum(axis=(0, 2))
+        mode_radii[ii] = (grid_radii * energy_vs_x).sum() / energy_vs_x.sum()
+
+    logger.info(f'{mode_radii=}')
+    lin_wavenumbers = angular_wavenumbers / mode_radii
+    return lin_wavenumbers
+
+

meanas/fdmath/vectorization.py

@@ -28,14 +28,16 @@ def vec(f: cfdfield_t) -> vcfdfield_t:
 def vec(f: ArrayLike) -> vfdfield_t | vcfdfield_t:
     pass
 
-def vec(f: fdfield_t | cfdfield_t | ArrayLike | None) -> vfdfield_t | vcfdfield_t | None:
+def vec(
+        f: fdfield_t | cfdfield_t | ArrayLike | None,
+        ) -> 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 vector field which spans a 1-3D region.
 
     Returns `None` if called with `f=None`.
 
     Args:
-        f: A vector field, `[f_x, f_y, f_z]` where each `f_` component is a 1- to
+        f: A vector field, e.g. `[f_x, f_y, f_z]` where each `f_` component is a 1- to
            3-D ndarray (`f_*` should all be the same size). Doesn't fail with `f=None`.
 
     Returns:
@@ -47,33 +49,38 @@ def vec(f: fdfield_t | cfdfield_t | ArrayLike | None) -> vfdfield_t | vcfdfield_
 
 
 @overload
-def unvec(v: None, shape: Sequence[int]) -> None:
+def unvec(v: None, shape: Sequence[int], nvdim: int) -> None:
     pass
 
 @overload
-def unvec(v: vfdfield_t, shape: Sequence[int]) -> fdfield_t:
+def unvec(v: vfdfield_t, shape: Sequence[int], nvdim: int) -> fdfield_t:
     pass
 
 @overload
-def unvec(v: vcfdfield_t, shape: Sequence[int]) -> cfdfield_t:
+def unvec(v: vcfdfield_t, shape: Sequence[int], nvdim: int) -> cfdfield_t:
     pass
 
-def unvec(v: vfdfield_t | vcfdfield_t | None, shape: Sequence[int]) -> fdfield_t | cfdfield_t | None:
+def unvec(
+        v: vfdfield_t | vcfdfield_t | None,
+        shape: Sequence[int],
+        nvdim: int = 3,
+        ) -> fdfield_t | cfdfield_t | None:
     """
-    Perform the inverse of vec(): take a 1D ndarray and output a 3D field
-     of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional
+    Perform the inverse of vec(): take a 1D ndarray and output an `nvdim`-component field
+     of form e.g. `[f_x, f_y, f_z]` (`nvdim=3`) where each of `f_*` is a len(shape)-dimensional
      ndarray.
 
     Returns `None` if called with `v=None`.
 
     Args:
-        v: 1D ndarray representing a 3D vector field of shape shape (or None)
+        v: 1D ndarray representing a vector field of shape shape (or None)
         shape: shape of the vector field
+        nvdim: Number of components in each vector
 
     Returns:
         `[f_x, f_y, f_z]` where each `f_` is a `len(shape)` dimensional ndarray (or `None`)
     """
     if v is None:
         return None
-    return v.reshape((3, *shape), order='C')
+    return v.reshape((nvdim, *shape), order='C')