Big documentation and structure updates

- Split math into fdmath package
- Rename waveguide into _2d _3d and _cyl variants
- pdoc-based documentation

make_docs.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+cd ~/projects/meanas
+pdoc3 --html --force --template-dir pdoc_templates -o doc .

meanas/eigensolvers.py

@@ -15,10 +15,13 @@ def power_iteration(operator: sparse.spmatrix,
     """
     Use power iteration to estimate the dominant eigenvector of a matrix.
 
-    :param operator: Matrix to analyze.
+    Args:
-    :param guess_vector: Starting point for the eigenvector. Default is a randomly chosen vector.
+       operator: Matrix to analyze.
-    :param iterations: Number of iterations to perform. Default 20.
+       guess_vector: Starting point for the eigenvector. Default is a randomly chosen vector.
-    :return: (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
+       iterations: Number of iterations to perform. Default 20.
+
+    Returns:
+        (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
     """
     if numpy.any(numpy.equal(guess_vector, None)):
         v = numpy.random.rand(operator.shape[0])
@@ -91,12 +94,15 @@ def signed_eigensolve(operator: sparse.spmatrix or spalg.LinearOperator,
     Find the largest-magnitude positive-only (or negative-only) eigenvalues and
      eigenvectors of the provided matrix.
 
-    :param operator: Matrix to analyze.
+    Args:
-    :param how_many: How many eigenvalues to find.
+        operator: Matrix to analyze.
-    :param negative: Whether to find negative-only eigenvalues.
+        how_many: How many eigenvalues to find.
+        negative: Whether to find negative-only eigenvalues.
                   Default False (positive only).
-    :return: (sorted list of eigenvalues, 2D ndarray of corresponding eigenvectors)
+
-        eigenvectors[:, k] corresponds to the k-th eigenvalue
+    Returns:
+        (sorted list of eigenvalues, 2D ndarray of corresponding eigenvectors)
+        `eigenvectors[:, k]` corresponds to the k-th eigenvalue
     """
     # Use power iteration to estimate the dominant eigenvector
     lm_eigval, _ = power_iteration(operator)

meanas/fdfd/__init__.py

@@ -1,2 +1,17 @@
-from . import solvers, operators, functional, scpml, waveguide, waveguide_mode
+"""
+Tools for finite difference frequency-domain (FDFD) simulations and calculations.
+
+These mostly involve picking a single frequency, then setting up and solving a
+matrix equation (Ax=b) or eigenvalue problem.
+
+
+Submodules:
+
+- `operators`, `functional`: General FDFD problem setup.
+- `solvers`: Solver interface and reference implementation.
+- `scpml`: Stretched-coordinate perfectly matched layer (scpml) boundary conditions
+- `waveguide_2d`: Operators and mode-solver for waveguides with constant cross-section.
+- `waveguide_3d`: Functions for transforming `waveguide_2d` results into 3D.
+"""
+from . import solvers, operators, functional, scpml, waveguide_2d, waveguide_3d
 # from . import farfield, bloch TODO

meanas/fdfd/bloch.py

@@ -83,7 +83,7 @@ import scipy.optimize
 from scipy.linalg import norm
 import scipy.sparse.linalg as spalg
 
-from . import field_t
+from .. import field_t
 
 logger = logging.getLogger(__name__)
 

meanas/fdfd/farfield.py

@@ -1,7 +1,7 @@
 """
 Functions for performing near-to-farfield transformation (and the reverse).
 """
-from typing import Dict, List
+from typing import Dict, List, Any
 import numpy
 from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
 from numpy import pi
@@ -14,7 +14,7 @@ def near_to_farfield(E_near: field_t,
                      dx: float,
                      dy: float,
                      padded_size: List[int] = None
-                     ) -> Dict[str]:
+                     ) -> Dict[str, Any]:
     """
     Compute the farfield, i.e. the distribution of the fields after propagation
       through several wavelengths of uniform medium.
@@ -122,7 +122,7 @@ def far_to_nearfield(E_far: field_t,
                      dkx: float,
                      dky: float,
                      padded_size: List[int] = None
-                     ) -> Dict[str]:
+                     ) -> Dict[str, Any]:
     """
     Compute the farfield, i.e. the distribution of the fields after propagation
       through several wavelengths of uniform medium.

meanas/fdfd/functional.py

@@ -2,88 +2,41 @@
 Functional versions of many FDFD operators. These can be useful for performing
  FDFD calculations without needing to construct large matrices in memory.
 
-The functions generated here expect field inputs with shape (3, X, Y, Z),
+The functions generated here expect `field_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)
 """
-from typing import List, Callable
+from typing import List, Callable, Tuple
 import numpy
 
 from .. import dx_lists_t, field_t
+from ..fdmath.functional import curl_forward, curl_back
 
 __author__ = 'Jan Petykiewicz'
 
 
-functional_matrix = Callable[[field_t], field_t]
+field_transform_t = Callable[[field_t], field_t]
-
-
-def curl_h(dxes: dx_lists_t) -> functional_matrix:
-    """
-    Curl operator for use with the H field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :return: Function for taking the discretized curl of the H-field, F(H) -> curlH
-    """
-    dxyz_b = numpy.meshgrid(*dxes[1], indexing='ij')
-
-    def dh(f, ax):
-        return (f - numpy.roll(f, 1, axis=ax)) / dxyz_b[ax]
-
-    def ch_fun(h: field_t) -> field_t:
-        e = numpy.empty_like(h)
-        e[0] = dh(h[2], 1)
-        e[0] -= dh(h[1], 2)
-        e[1] = dh(h[0], 2)
-        e[1] -= dh(h[2], 0)
-        e[2] = dh(h[1], 0)
-        e[2] -= dh(h[0], 1)
-        return e
-
-    return ch_fun
-
-
-def curl_e(dxes: dx_lists_t) -> functional_matrix:
-    """
-    Curl operator for use with the E field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :return: Function for taking the discretized curl of the E-field, F(E) -> curlE
-    """
-    dxyz_a = numpy.meshgrid(*dxes[0], indexing='ij')
-
-    def de(f, ax):
-        return (numpy.roll(f, -1, axis=ax) - f) / dxyz_a[ax]
-
-    def ce_fun(e: field_t) -> field_t:
-        h = numpy.empty_like(e)
-        h[0] = de(e[2], 1)
-        h[0] -= de(e[1], 2)
-        h[1] = de(e[0], 2)
-        h[1] -= de(e[2], 0)
-        h[2] = de(e[1], 0)
-        h[2] -= de(e[0], 1)
-        return h
-
-    return ce_fun
 
 
 def e_full(omega: complex,
            dxes: dx_lists_t,
            epsilon: field_t,
            mu: field_t = None
-           ) -> functional_matrix:
+           ) -> field_transform_t:
     """
-    Wave operator del x (1/mu * del x) - omega**2 * epsilon, for use with E-field,
+    Wave operator for use with E-field. See `operators.e_full` for details.
-     with wave equation
-    (del x (1/mu * del x) - omega**2 * epsilon) E = -i * omega * J
 
-    :param omega: Angular frequency of the simulation
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-    :param epsilon: Dielectric constant
+        dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :param mu: Magnetic permeability (default 1 everywhere)
+        epsilon: Dielectric constant
-    :return: Function implementing the wave operator A(E) -> E
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Return:
+        Function `f` implementing the wave operator
+        `f(E)` -> `-i * omega * J`
     """
-    ch = curl_h(dxes)
+    ch = curl_back(dxes[1])
-    ce = curl_e(dxes)
+    ce = curl_forward(dxes[0])
 
     def op_1(e):
         curls = ch(ce(e))
@@ -103,18 +56,23 @@ def eh_full(omega: complex,
             dxes: dx_lists_t,
             epsilon: field_t,
             mu: field_t = None
-            ) -> functional_matrix:
+            ) -> Callable[[field_t, field_t], Tuple[field_t, field_t]]:
     """
     Wave operator for full (both E and H) field representation.
+    See `operators.eh_full`.
 
-    :param omega: Angular frequency of the simulation
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-    :param epsilon: Dielectric constant
+        dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :param mu: Magnetic permeability (default 1 everywhere)
+        epsilon: Dielectric constant
-    :return: Function implementing the wave operator A(E, H) -> (E, H)
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Function `f` implementing the wave operator
+        `f(E, H)` -> `(J, -M)`
     """
-    ch = curl_h(dxes)
+    ch = curl_back(dxes[1])
-    ce = curl_e(dxes)
+    ce = curl_forward(dxes[0])
 
     def op_1(e, h):
         return (ch(h) - 1j * omega * epsilon * e,
@@ -133,23 +91,27 @@ def eh_full(omega: complex,
 def e2h(omega: complex,
         dxes: dx_lists_t,
         mu: field_t = None,
-        ) -> functional_matrix:
+        ) -> field_transform_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`.
 
-   :param omega: Angular frequency of the simulation
+    Args:
-   :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-   :param mu: Magnetic permeability (default 1 everywhere)
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-   :return: Function for converting E to H
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Return:
+        Function `f` for converting `E` to `H`,
+        `f(E)` -> `H`
     """
-    A2 = curl_e(dxes)
+    ce = curl_forward(dxes[0])
 
     def e2h_1_1(e):
-        return A2(e) / (-1j * omega)
+        return ce(e) / (-1j * omega)
 
     def e2h_mu(e):
-        return A2(e) / (-1j * omega * mu)
+        return ce(e) / (-1j * omega * mu)
 
     if numpy.any(numpy.equal(mu, None)):
         return e2h_1_1
@@ -160,18 +122,22 @@ def e2h(omega: complex,
 def m2j(omega: complex,
         dxes: dx_lists_t,
         mu: field_t = None,
-        ) -> functional_matrix:
+        ) -> field_transform_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`.
-   For use with e.g. e_full().
+    For use with e.g. `e_full`.
 
-   :param omega: Angular frequency of the simulation
+    Args:
-   :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-   :param mu: Magnetic permeability (default 1 everywhere)
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-   :return: Function for converting M to J
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Function `f` for converting `M` to `J`,
+        `f(M)` -> `J`
     """
-    ch = curl_h(dxes)
+    ch = curl_back(dxes[1])
 
     def m2j_mu(m):
         J = ch(m / mu) / (-1j * omega)
@@ -192,10 +158,23 @@ def e_tfsf_source(TF_region: field_t,
                   dxes: dx_lists_t,
                   epsilon: field_t,
                   mu: field_t = None,
-                  ) -> functional_matrix:
+                  ) -> field_transform_t:
     """
-    Operator that turuns 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.
+
+    Args:
+        TF_region: mask which is set to 1 in the total-field region, and 0 elsewhere
+                   (i.e. in the scattered-field region).
+                   Should have the same shape as the simulation grid, e.g. `epsilon[0].shape`.
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        epsilon: Dielectric constant distribution
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Function `f` which takes an E field and returns a current distribution,
+        `f(E)` -> `J`
     """
     # TODO documentation
     A = e_full(omega, dxes, epsilon, mu)
@@ -205,7 +184,28 @@ def e_tfsf_source(TF_region: field_t,
         return neg_iwj / (-1j * omega)
 
 
-def poynting_e_cross_h(dxes: dx_lists_t):
+def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[field_t, field_t], field_t]:
+    """
+    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
+    for the Poynting vector, \\( S = E \\times H \\)
+
+    Note:
+        This function also shifts the input `E` field by one cell as required
+        for computing the Poynting cross product (see `meanas.fdfd` module docs).
+
+    Note:
+        If `E` and `H` are peak amplitudes as assumed elsewhere in this code,
+        the time-average of the poynting vector is `<S> = Re(S)/2 = Re(E x H) / 2`.
+        The factor of `1/2` can be omitted if root-mean-square quantities are used
+        instead.
+
+    Args:
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+
+    Returns:
+        Function `f` that returns E x H as required for the poynting vector.
+    """
     def exh(e: field_t, h: field_t):
         s = numpy.empty_like(e)
         ex = e[0] * dxes[0][0][:, None, None]

meanas/fdfd/operators.py

@@ -1,18 +1,19 @@
 """
 Sparse matrix operators for use with electromagnetic wave equations.
 
-These functions return sparse-matrix (scipy.sparse.spmatrix) representations of
+These functions return sparse-matrix (`scipy.sparse.spmatrix`) representations of
  a variety of operators, intended for use with E and H fields vectorized using the
- meanas.vec() and .unvec() functions (column-major/Fortran ordering).
+ `meanas.vec()` and `meanas.unvec()` functions.
 
-E- and H-field values are defined on a Yee cell; epsilon values should be calculated for
+E- and H-field values are defined on a Yee cell; `epsilon` values should be calculated for
- cells centered at each E component (mu at each H component).
+ cells centered at each E component (`mu` at each H component).
 
-Many of these functions require a 'dxes' parameter, of type meanas.dx_lists_type; see
+Many of these functions require a `dxes` parameter, of type `dx_lists_t`; see
-the meanas.types submodule for details.
+the `meanas.types` submodule for details.
 
 
 The following operators are included:
+
 - E-only wave operator
 - H-only wave operator
 - EH wave operator
@@ -20,8 +21,6 @@ The following operators are included:
 - E to H conversion
 - M to J conversion
 - Poynting cross products
-
-Also available:
 - Circular shifts
 - Discrete derivatives
 - Averaging operators
@@ -33,6 +32,7 @@ import numpy
 import scipy.sparse as sparse
 
 from .. import vec, dx_lists_t, vfield_t
+from ..fdmath.operators import shift_with_mirror, rotation, curl_forward, curl_back
 
 
 __author__ = 'Jan Petykiewicz'
@@ -46,26 +46,35 @@ def e_full(omega: complex,
            pmc: vfield_t = None,
            ) -> sparse.spmatrix:
     """
-    Wave operator del x (1/mu * del x) - omega**2 * epsilon, for use with E-field,
+    Wave operator
-     with wave equation
+     $$ \\nabla \\times (\\frac{1}{\\mu} \\nabla \\times) - \\omega^2 \\epsilon $$
+
+        del x (1/mu * del x) - omega**2 * epsilon
+
+     for use with the E-field, with wave equation
+     $$ (\\nabla \\times (\\frac{1}{\\mu} \\nabla \\times) - \\omega^2 \\epsilon) E = -\\imath \\omega J $$
+
         (del x (1/mu * del x) - omega**2 * epsilon) E = -i * omega * J
 
-    To make this matrix symmetric, use the preconditions from e_full_preconditioners().
+    To make this matrix symmetric, use the preconditioners from `e_full_preconditioners()`.
 
-    :param omega: Angular frequency of the simulation
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-    :param epsilon: Vectorized dielectric constant
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-    :param mu: Vectorized magnetic permeability (default 1 everywhere).
+        epsilon: Vectorized dielectric constant
-    :param pec: Vectorized mask specifying PEC cells. Any cells where pec != 0 are interpreted
+        mu: Vectorized magnetic permeability (default 1 everywhere).
+        pec: Vectorized mask specifying PEC cells. Any cells where `pec != 0` are interpreted
           as containing a perfect electrical conductor (PEC).
-        The PEC is applied per-field-component (ie, pec.size == epsilon.size)
+          The PEC is applied per-field-component (i.e. `pec.size == epsilon.size`)
-    :param pmc: Vectorized mask specifying PMC cells. Any cells where pmc != 0 are interpreted
+        pmc: Vectorized mask specifying PMC cells. Any cells where `pmc != 0` are interpreted
           as containing a perfect magnetic conductor (PMC).
-        The PMC is applied per-field-component (ie, pmc.size == epsilon.size)
+          The PMC is applied per-field-component (i.e. `pmc.size == epsilon.size`)
-    :return: Sparse matrix containing the wave operator
+
+    Returns:
+        Sparse matrix containing the wave operator.
     """
-    ce = curl_e(dxes)
+    ch = curl_back(dxes[1])
-    ch = curl_h(dxes)
+    ce = curl_forward(dxes[0])
 
     if numpy.any(numpy.equal(pec, None)):
         pe = sparse.eye(epsilon.size)
@@ -90,15 +99,18 @@ def e_full(omega: complex,
 def e_full_preconditioners(dxes: dx_lists_t
                            ) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
     """
-    Left and right preconditioners (Pl, Pr) for symmetrizing the e_full wave operator.
+    Left and right preconditioners `(Pl, Pr)` for symmetrizing the `e_full` wave operator.
 
-    The preconditioned matrix A_symm = (Pl @ A @ Pr) is complex-symmetric
+    The preconditioned matrix `A_symm = (Pl @ A @ Pr)` is complex-symmetric
      (non-Hermitian unless there is no loss or PMLs).
 
-    The preconditioner matrices are diagonal and complex, with Pr = 1 / Pl
+    The preconditioner matrices are diagonal and complex, with `Pr = 1 / Pl`
 
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+    Args:
-    :return: Preconditioner matrices (Pl, Pr)
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+
+    Returns:
+        Preconditioner matrices `(Pl, Pr)`.
     """
     p_squared = [dxes[0][0][:, None, None] * dxes[1][1][None, :, None] * dxes[1][2][None, None, :],
                  dxes[1][0][:, None, None] * dxes[0][1][None, :, None] * dxes[1][2][None, None, :],
@@ -118,24 +130,33 @@ def h_full(omega: complex,
            pmc: vfield_t = None,
            ) -> sparse.spmatrix:
     """
-    Wave operator del x (1/epsilon * del x) - omega**2 * mu, for use with H-field,
+    Wave operator
-     with wave equation
+     $$ \\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu $$
-    (del x (1/epsilon * del x) - omega**2 * mu) H = i * omega * M
 
-    :param omega: Angular frequency of the simulation
+        del x (1/epsilon * del x) - omega**2 * mu
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+
-    :param epsilon: Vectorized dielectric constant
+     for use with the H-field, with wave equation
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+     $$ (\\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu) E = \\imath \\omega M $$
-    :param pec: Vectorized mask specifying PEC cells. Any cells where pec != 0 are interpreted
+
+        (del x (1/epsilon * del x) - omega**2 * mu) E = i * omega * M
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        epsilon: Vectorized dielectric constant
+        mu: Vectorized magnetic permeability (default 1 everywhere)
+        pec: Vectorized mask specifying PEC cells. Any cells where `pec != 0` are interpreted
            as containing a perfect electrical conductor (PEC).
-        The PEC is applied per-field-component (ie, pec.size == epsilon.size)
+           The PEC is applied per-field-component (i.e. `pec.size == epsilon.size`)
-    :param pmc: Vectorized mask specifying PMC cells. Any cells where pmc != 0 are interpreted
+        pmc: Vectorized mask specifying PMC cells. Any cells where `pmc != 0` are interpreted
            as containing a perfect magnetic conductor (PMC).
-        The PMC is applied per-field-component (ie, pmc.size == epsilon.size)
+           The PMC is applied per-field-component (i.e. `pmc.size == epsilon.size`)
-    :return: Sparse matrix containing the wave operator
+
+    Returns:
+        Sparse matrix containing the wave operator.
     """
-    ec = curl_e(dxes)
+    ch = curl_back(dxes[1])
-    hc = curl_h(dxes)
+    ce = curl_forward(dxes[0])
 
     if numpy.any(numpy.equal(pec, None)):
         pe = sparse.eye(epsilon.size)
@@ -153,7 +174,7 @@ def h_full(omega: complex,
     else:
         m = sparse.diags(mu)
 
-    A = pm @ (ec @ pe @ e_div @ hc - omega**2 * m) @ pm
+    A = pm @ (ce @ pe @ e_div @ ch - omega**2 * m) @ pm
     return A
 
 
@@ -165,24 +186,42 @@ def eh_full(omega: complex,
             pmc: vfield_t = None
             ) -> sparse.spmatrix:
     """
-    Wave operator for [E, H] field representation. This operator implements Maxwell's
+    Wave operator for `[E, H]` field representation. This operator implements Maxwell's
      equations without cancelling out either E or H. The operator is
-     [[-i * omega * epsilon,  del x],
+    $$  \\begin{bmatrix}
+        -\\imath \\omega \\epsilon  &  \\nabla \\times      \\\\
+        \\nabla \\times             &  \\imath \\omega \\mu
+        \\end{bmatrix} $$
+
+        [[-i * omega * epsilon,  del x         ],
          [del x,                 i * omega * mu]]
 
-    for use with a field vector of the form hstack(vec(E), vec(H)).
+    for use with a field vector of the form `cat(vec(E), vec(H))`:
+    $$  \\begin{bmatrix}
+        -\\imath \\omega \\epsilon  &  \\nabla \\times      \\\\
+        \\nabla \\times             &  \\imath \\omega \\mu
+        \\end{bmatrix}
+        \\begin{bmatrix} E \\\\
+                         H
+        \\end{bmatrix}
+        = \\begin{bmatrix} J \\\\
+                          -M
+          \\end{bmatrix} $$
 
-    :param omega: Angular frequency of the simulation
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-    :param epsilon: Vectorized dielectric constant
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+        epsilon: Vectorized dielectric constant
-    :param pec: Vectorized mask specifying PEC cells. Any cells where pec != 0 are interpreted
+        mu: Vectorized magnetic permeability (default 1 everywhere)
+        pec: Vectorized mask specifying PEC cells. Any cells where `pec != 0` are interpreted
           as containing a perfect electrical conductor (PEC).
-        The PEC is applied per-field-component (i.e., pec.size == epsilon.size)
+          The PEC is applied per-field-component (i.e. `pec.size == epsilon.size`)
-    :param pmc: Vectorized mask specifying PMC cells. Any cells where pmc != 0 are interpreted
+        pmc: Vectorized mask specifying PMC cells. Any cells where `pmc != 0` are interpreted
           as containing a perfect magnetic conductor (PMC).
-        The PMC is applied per-field-component (i.e., pmc.size == epsilon.size)
+          The PMC is applied per-field-component (i.e. `pmc.size == epsilon.size`)
-    :return: Sparse matrix containing the wave operator
+
+    Returns:
+        Sparse matrix containing the wave operator.
     """
     if numpy.any(numpy.equal(pec, None)):
         pe = sparse.eye(epsilon.size)
@@ -200,34 +239,14 @@ def eh_full(omega: complex,
         iwm *= sparse.diags(mu)
     iwm = pm @ iwm @ pm
 
-    A1 = pe @ curl_h(dxes) @ pm
+    A1 = pe @ curl_back(dxes[1]) @ pm
-    A2 = pm @ curl_e(dxes) @ pe
+    A2 = pm @ curl_forward(dxes[0]) @ pe
 
     A = sparse.bmat([[-iwe, A1],
                      [A2,  iwm]])
     return A
 
 
-def curl_h(dxes: dx_lists_t) -> sparse.spmatrix:
-    """
-    Curl operator for use with the H field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :return: Sparse matrix for taking the discretized curl of the H-field
-    """
-    return cross(deriv_back(dxes[1]))
-
-
-def curl_e(dxes: dx_lists_t) -> sparse.spmatrix:
-    """
-    Curl operator for use with the E field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :return: Sparse matrix for taking the discretized curl of the E-field
-    """
-    return cross(deriv_forward(dxes[0]))
-
-
 def e2h(omega: complex,
         dxes: dx_lists_t,
         mu: vfield_t = None,
@@ -235,17 +254,20 @@ def e2h(omega: complex,
         ) -> sparse.spmatrix:
     """
     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.
 
-    :param omega: Angular frequency of the simulation
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-    :param pmc: Vectorized mask specifying PMC cells. Any cells where pmc != 0 are interpreted
+        mu: Vectorized magnetic permeability (default 1 everywhere)
+        pmc: Vectorized mask specifying PMC cells. Any cells where `pmc != 0` are interpreted
           as containing a perfect magnetic conductor (PMC).
-        The PMC is applied per-field-component (ie, pmc.size == epsilon.size)
+          The PMC is applied per-field-component (i.e. `pmc.size == epsilon.size`)
-    :return: Sparse matrix for converting E to H
+
+    Returns:
+        Sparse matrix for converting E to H.
     """
-    op = curl_e(dxes) / (-1j * omega)
+    op = curl_forward(dxes[0]) / (-1j * omega)
 
     if not numpy.any(numpy.equal(mu, None)):
         op = sparse.diags(1 / mu) @ op
@@ -261,16 +283,18 @@ def m2j(omega: complex,
         mu: vfield_t = None
         ) -> sparse.spmatrix:
     """
-    Utility operator for converting M field into J.
+    Operator for converting a magnetic current M into an electric current J.
-    Converts a magnetic current M into an electric current J.
+    For use with eg. `e_full()`.
-    For use with eg. e_full.
 
-    :param omega: Angular frequency of the simulation
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        omega: Angular frequency of the simulation
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-    :return: Sparse matrix for converting E to H
+        mu: Vectorized magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Sparse matrix for converting M to J.
     """
-    op = curl_h(dxes) / (1j * omega)
+    op = curl_back(dxes[1]) / (1j * omega)
 
     if not numpy.any(numpy.equal(mu, None)):
         op = op @ sparse.diags(1 / mu)
@@ -278,178 +302,17 @@ def m2j(omega: complex,
     return op
 
 
-def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix:
-    """
-    Utility operator for performing a circular shift along a specified axis by a
-     specified number of elements.
-
-    :param axis: Axis to shift along. x=0, y=1, z=2
-    :param shape: Shape of the grid being shifted
-    :param shift_distance: Number of cells to shift by. May be negative. Default 1.
-    :return: Sparse matrix for performing the circular shift
-    """
-    if len(shape) not in (2, 3):
-        raise Exception('Invalid shape: {}'.format(shape))
-    if axis not in range(len(shape)):
-        raise Exception('Invalid direction: {}, shape is {}'.format(axis, 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)]
-    ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
-
-    n = numpy.prod(shape)
-    i_ind = numpy.arange(n)
-    j_ind = numpy.ravel_multi_index(ijk, shape, order='C')
-
-    vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C')))
-
-    d = sparse.csr_matrix(vij, shape=(n, n))
-
-    if shift_distance < 0:
-        d = d.T
-
-    return d
-
-
-def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix:
-    """
-    Utility operator for performing an n-element shift along a specified axis, with mirror
-    boundary conditions applied to the cells beyond the receding edge.
-
-    :param axis: Axis to shift along. x=0, y=1, z=2
-    :param shape: Shape of the grid being shifted
-    :param shift_distance: Number of cells to shift by. May be negative. Default 1.
-    :return: Sparse matrix for performing the circular shift
-    """
-    if len(shape) not in (2, 3):
-        raise Exception('Invalid shape: {}'.format(shape))
-    if axis not in range(len(shape)):
-        raise Exception('Invalid direction: {}, shape is {}'.format(axis, shape))
-    if shift_distance >= shape[axis]:
-        raise Exception('Shift ({}) is too large for axis {} of size {}'.format(
-                        shift_distance, axis, shape[axis]))
-
-    def mirrored_range(n, s):
-        v = numpy.arange(n) + s
-        v = numpy.where(v >= n, 2 * n - v - 1, v)
-        v = numpy.where(v < 0, - 1 - v, v)
-        return v
-
-    shifts = [shift_distance if a == axis else 0 for a in range(3)]
-    shifted_diags = [mirrored_range(n, s) for n, s in zip(shape, shifts)]
-    ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
-
-    n = numpy.prod(shape)
-    i_ind = numpy.arange(n)
-    j_ind = numpy.ravel_multi_index(ijk, shape, order='C')
-
-    vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C')))
-
-    d = sparse.csr_matrix(vij, shape=(n, n))
-    return d
-
-
-def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]:
-    """
-    Utility operators for taking discretized derivatives (forward variant).
-
-    :param dx_e: Lists of cell sizes for all axes [[dx_0, dx_1, ...], ...].
-    :return: List of operators for taking forward derivatives along each axis.
-    """
-    shape = [s.size for s in dx_e]
-    n = numpy.prod(shape)
-
-    dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij')
-
-    def deriv(axis):
-        return rotation(axis, shape, 1) - sparse.eye(n)
-
-    Ds = [sparse.diags(+1 / dx.ravel(order='C')) @ deriv(a)
-          for a, dx in enumerate(dx_e_expanded)]
-
-    return Ds
-
-
-def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]:
-    """
-    Utility operators for taking discretized derivatives (backward variant).
-
-    :param dx_h: Lists of cell sizes for all axes [[dx_0, dx_1, ...], ...].
-    :return: List of operators for taking forward derivatives along each axis.
-    """
-    shape = [s.size for s in dx_h]
-    n = numpy.prod(shape)
-
-    dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij')
-
-    def deriv(axis):
-        return rotation(axis, shape, -1) - sparse.eye(n)
-
-    Ds = [sparse.diags(-1 / dx.ravel(order='C')) @ deriv(a)
-          for a, dx in enumerate(dx_h_expanded)]
-
-    return Ds
-
-
-def cross(B: List[sparse.spmatrix]) -> sparse.spmatrix:
-    """
-    Cross product operator
-
-    :param B: List [Bx, By, Bz] of sparse matrices corresponding to the x, y, z
-            portions of the operator on the left side of the cross product.
-    :return: Sparse matrix corresponding to (B x), where x is the cross product
-    """
-    n = B[0].shape[0]
-    zero = sparse.csr_matrix((n, n))
-    return sparse.bmat([[zero, -B[2], B[1]],
-                        [B[2], zero, -B[0]],
-                        [-B[1], B[0], zero]])
-
-
-def vec_cross(b: vfield_t) -> sparse.spmatrix:
-    """
-    Vector cross product operator
-
-    :param b: Vector on the left side of the cross product
-    :return: Sparse matrix corresponding to (b x), where x is the cross product
-    """
-    B = [sparse.diags(c) for c in numpy.split(b, 3)]
-    return cross(B)
-
-
-def avgf(axis: int, shape: List[int]) -> sparse.spmatrix:
-    """
-    Forward average operator (x4 = (x4 + x5) / 2)
-
-    :param axis: Axis to average along (x=0, y=1, z=2)
-    :param shape: Shape of the grid to average
-    :return: Sparse matrix for forward average operation
-    """
-    if len(shape) not in (2, 3):
-        raise Exception('Invalid shape: {}'.format(shape))
-
-    n = numpy.prod(shape)
-    return 0.5 * (sparse.eye(n) + rotation(axis, shape))
-
-
-def avgb(axis: int, shape: List[int]) -> sparse.spmatrix:
-    """
-    Backward average operator (x4 = (x4 + x3) / 2)
-
-    :param axis: Axis to average along (x=0, y=1, z=2)
-    :param shape: Shape of the grid to average
-    :return: Sparse matrix for backward average operation
-    """
-    return avgf(axis, shape).T
-
-
 def poynting_e_cross(e: vfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     """
-    Operator for computing the Poynting vector, containing the (E x) portion of the Poynting vector.
+    Operator for computing the Poynting vector, containing the
+    (E x) portion of the Poynting vector.
 
-    :param e: Vectorized E-field for the ExH cross product
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        e: Vectorized E-field for the ExH cross product
-    :return: Sparse matrix containing (E x) portion of Poynting cross product
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+
+    Returns:
+        Sparse matrix containing (E x) portion of Poynting cross product.
     """
     shape = [len(dx) for dx in dxes[0]]
 
@@ -472,9 +335,12 @@ def poynting_h_cross(h: vfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
     """
     Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector.
 
-    :param h: Vectorized H-field for the HxE cross product
+    Args:
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+        h: Vectorized H-field for the HxE cross product
-    :return: Sparse matrix containing (H x) portion of Poynting cross product
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+
+    Returns:
+        Sparse matrix containing (H x) portion of Poynting cross product.
     """
     shape = [len(dx) for dx in dxes[0]]
 
@@ -499,8 +365,22 @@ def e_tfsf_source(TF_region: vfield_t,
                   mu: vfield_t = None,
                   ) -> sparse.spmatrix:
     """
-    Operator that turns an E-field distribution into a total-field/scattered-field
+    Operator that turns a desired E-field distribution into a
-    (TFSF) source.
+     total-field/scattered-field (TFSF) source.
+
+    TODO: Reference Rumpf paper
+
+    Args:
+        TF_region: Mask, which is set to 1 inside the total-field region and 0 in the
+                   scattered-field region
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        epsilon: Vectorized dielectric constant
+        mu: Vectorized magnetic permeability (default 1 everywhere).
+
+    Returns:
+        Sparse matrix that turns an E-field into a current (J) distribution.
+
     """
     # TODO documentation
     A = e_full(omega, dxes, epsilon, mu)
@@ -518,7 +398,19 @@ def e_boundary_source(mask: vfield_t,
     """
     Operator that turns an E-field distrubtion into a current (J) distribution
       along the edges (external and internal) of the provided mask. This is just an
-      e_tfsf_source with an additional masking step.
+      `e_tfsf_source()` with an additional masking step.
+
+    Args:
+        mask: The current distribution is generated at the edges of the mask,
+              i.e. any points where shifting the mask by one cell in any direction
+              would change its value.
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        epsilon: Vectorized dielectric constant
+        mu: Vectorized magnetic permeability (default 1 everywhere).
+
+    Returns:
+        Sparse matrix that turns an E-field into a current (J) distribution.
     """
     full = e_tfsf_source(TF_region=mask, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu)
 

meanas/fdfd/scpml.py

@@ -1,5 +1,5 @@
 """
-Functions for creating stretched coordinate PMLs.
+Functions for creating stretched coordinate perfectly matched layer (PML) absorbers.
 """
 
 from typing import List, Callable
@@ -7,23 +7,28 @@ import numpy
 
 from .. import dx_lists_t
 
+
 __author__ = 'Jan Petykiewicz'
 
 
-s_function_type = Callable[[float], float]
+s_function_t = Callable[[float], float]
+"""Typedef for s-functions"""
 
 
 def prepare_s_function(ln_R: float = -16,
                        m: float = 4
-                       ) -> s_function_type:
+                       ) -> s_function_t:
     """
     Create an s_function to pass to the SCPML functions. This is used when you would like to
     customize the PML parameters.
 
-    :param ln_R: Natural logarithm of the desired reflectance
+    Args:
-    :param m: Polynomial order for the PML (imaginary part increases as distance ** m)
+        ln_R: Natural logarithm of the desired reflectance
-    :return: An s_function, which takes an ndarray (distances) and returns an ndarray (complex part
+        m: Polynomial order for the PML (imaginary part increases as distance ** m)
-                of the cell width; needs to be divided by sqrt(epilon_effective) * real(omega))
+
+    Returns:
+        An s_function, which takes an ndarray (distances) and returns an ndarray (complex part
+        of the cell width; needs to be divided by `sqrt(epilon_effective) * real(omega))`
         before use.
     """
     def s_factor(distance: numpy.ndarray) -> numpy.ndarray:
@@ -36,26 +41,29 @@ def uniform_grid_scpml(shape: numpy.ndarray or List[int],
                        thicknesses: numpy.ndarray or List[int],
                        omega: float,
                        epsilon_effective: float = 1.0,
-                       s_function: s_function_type = None,
+                       s_function: s_function_t = None,
                        ) -> dx_lists_t:
     """
     Create dx arrays for a uniform grid with a cell width of 1 and a pml.
 
-    If you want something more fine-grained, check out stretch_with_scpml(...).
+    If you want something more fine-grained, check out `stretch_with_scpml(...)`.
 
-    :param shape: Shape of the grid, including the PMLs (which are 2*thicknesses thick)
+    Args:
-    :param thicknesses: [th_x, th_y, th_z] Thickness of the PML in each direction.
+        shape: Shape of the grid, including the PMLs (which are 2*thicknesses thick)
+        thicknesses: `[th_x, th_y, th_z]`
+                     Thickness of the PML in each direction.
                      Both polarities are added.
                      Each th_ of pml is applied twice, once on each edge of the grid along the given axis.
-        th_* may be zero, in which case no pml is added.
+                     `th_*` may be zero, in which case no pml is added.
-    :param omega: Angular frequency for the simulation
+        omega: Angular frequency for the simulation
-    :param epsilon_effective: Effective epsilon of the PML. Match this to the material
+        epsilon_effective: Effective epsilon of the PML. Match this to the material
                             at the edge of your grid.
                             Default 1.
-    :param s_function: s_function created by prepare_s_function(...), allowing
+        s_function: created by `prepare_s_function(...)`, allowing customization of pml parameters.
-        customization of pml parameters.
+                    Default uses `prepare_s_function()` with no parameters.
-        Default uses prepare_s_function() with no parameters.
+
-    :return: Complex cell widths (dx_lists)
+    Returns:
+        Complex cell widths (dx_lists_t) as discussed in `meanas.types`.
     """
     if s_function is None:
         s_function = prepare_s_function()
@@ -88,21 +96,25 @@ def stretch_with_scpml(dxes: dx_lists_t,
                        omega: float,
                        epsilon_effective: float = 1.0,
                        thickness: int = 10,
-                       s_function: s_function_type = None,
+                       s_function: s_function_t = None,
                        ) -> dx_lists_t:
     """
         Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis.
 
-        :param dxes: dx_tuple with coordinates to stretch
+        Args:
-        :param axis: axis to stretch (0=x, 1=y, 2=z)
+            dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
-        :param polarity: direction to stretch (-1 for -ve, +1 for +ve)
+            axis: axis to stretch (0=x, 1=y, 2=z)
-        :param omega: Angular frequency for the simulation
+            polarity: direction to stretch (-1 for -ve, +1 for +ve)
-        :param epsilon_effective: Effective epsilon of the PML. Match this to the material at the
+            omega: Angular frequency for the simulation
+            epsilon_effective: Effective epsilon of the PML. Match this to the material at the
                                edge of your grid. Default 1.
-        :param thickness: number of cells to use for pml (default 10)
+            thickness: number of cells to use for pml (default 10)
-        :param s_function: s_function created by prepare_s_function(...), allowing customization
+            s_function: Created by `prepare_s_function(...)`, allowing customization
-            of pml parameters. Default uses prepare_s_function() with no parameters.
+                        of pml parameters. Default uses `prepare_s_function()` with no parameters.
-        :return: Complex cell widths
+
+        Returns:
+            Complex cell widths (dx_lists_t) as discussed in `meanas.types`.
+            Multiple calls to this function may be necessary if multiple absorpbing boundaries are needed.
     """
     if s_function is None:
         s_function = prepare_s_function()
@@ -147,25 +159,3 @@ def stretch_with_scpml(dxes: dx_lists_t,
     dxes[1][axis] = dx_bi
 
     return dxes
-
-
-def generate_periodic_dx(pos: List[numpy.ndarray]) -> dx_lists_t:
-    """
-    Given a list of 3 ndarrays cell centers, creates the cell width parameters for a periodic grid.
-
-    :param pos: List of 3 ndarrays of cell centers
-    :return: (dx_a, dx_b) cell widths (no pml)
-    """
-    if len(pos) != 3:
-        raise Exception('Must have len(pos) == 3')
-
-    dx_a = [numpy.array(numpy.inf)] * 3
-    dx_b = [numpy.array(numpy.inf)] * 3
-
-    for i, p_orig in enumerate(pos):
-        p = numpy.array(p_orig, dtype=float)
-        if p.size != 1:
-            p_shifted = numpy.hstack((p[1:], p[-1] + (p[1] - p[0])))
-            dx_a[i] = numpy.diff(p)
-            dx_b[i] = numpy.diff((p + p_shifted) / 2)
-    return dx_a, dx_b

meanas/fdfd/solvers.py

@@ -1,5 +1,5 @@
 """
-Solvers for FDFD problems.
+Solvers and solver interface for FDFD problems.
 """
 
 from typing import List, Callable, Dict, Any
@@ -22,10 +22,13 @@ def _scipy_qmr(A: scipy.sparse.csr_matrix,
     """
     Wrapper for scipy.sparse.linalg.qmr
 
-    :param A: Sparse matrix
+    Args:
-    :param b: Right-hand-side vector
+        A: Sparse matrix
-    :param kwargs: Passed as **kwargs to the wrapped function
+        b: Right-hand-side vector
-    :return: Guess for solution (returned even if didn't converge)
+        kwargs: Passed as **kwargs to the wrapped function
+
+    Returns:
+        Guess for solution (returned even if didn't converge)
     """
 
     '''
@@ -70,27 +73,31 @@ def generic(omega: complex,
     """
     Conjugate gradient FDFD solver using CSR sparse matrices.
 
-    All ndarray arguments should be 1D array, as returned by meanas.vec().
+    All ndarray arguments should be 1D arrays, as returned by `meanas.vec()`.
 
-    :param omega: Complex frequency to solve at.
+    Args:
-    :param dxes: [[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]] (complex cell sizes)
+        omega: Complex frequency to solve at.
-    :param J: Electric current distribution (at E-field locations)
+        dxes: `[[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]]` (complex cell sizes) as
-    :param epsilon: Dielectric constant distribution (at E-field locations)
+            discussed in `meanas.types`
-    :param mu: Magnetic permeability distribution (at H-field locations)
+        J: Electric current distribution (at E-field locations)
-    :param pec: Perfect electric conductor distribution
+        epsilon: Dielectric constant distribution (at E-field locations)
+        mu: Magnetic permeability distribution (at H-field locations)
+        pec: Perfect electric conductor distribution
              (at E-field locations; non-zero value indicates PEC is present)
-    :param pmc: Perfect magnetic conductor distribution
+        pmc: Perfect magnetic conductor distribution
              (at H-field locations; non-zero value indicates PMC is present)
-    :param adjoint: If true, solves the adjoint problem.
+        adjoint: If true, solves the adjoint problem.
-    :param matrix_solver: Called as matrix_solver(A, b, **matrix_solver_opts) -> x
+        matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`,
-        Where A: scipy.sparse.csr_matrix
+                where `A`: `scipy.sparse.csr_matrix`;
-              b: numpy.ndarray
+                      `b`: `numpy.ndarray`;
-              x: numpy.ndarray
+                      `x`: `numpy.ndarray`;
-        Default is a wrapped version of scipy.sparse.linalg.qmr()
+                Default is a wrapped version of `scipy.sparse.linalg.qmr()`
                  which doesn't return convergence info and logs the residual
                  every 100 iterations.
-    :param matrix_solver_opts: Passed as kwargs to matrix_solver(...)
+        matrix_solver_opts: Passed as kwargs to `matrix_solver(...)`
-    :return: E-field which solves the system.
+
+    Returns:
+        E-field which solves the system.
     """
 
     if matrix_solver_opts is None:

meanas/fdfd/waveguide.py

@@ -1,492 +0,0 @@
-"""
-Various operators and helper functions for solving for waveguide modes.
-
-Assuming a z-dependence of the from exp(-i * wavenumber * z), we can simplify Maxwell's
- equations in the absence of sources to the form
-
-A @ [H_x, H_y] = wavenumber**2 * [H_x, H_y]
-
-with A =
-omega**2 * epsilon * mu +
-epsilon * [[-Dy], [Dx]] / epsilon * [-Dy, Dx] +
-[[Dx], [Dy]] / mu * [Dx, Dy] * mu
-
-which is the form used in this file.
-
-As the z-dependence is known, all the functions in this file assume a 2D grid
- (ie. dxes = [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...]], [[dx_h_0, ...], [dy_h_0, ...]]])
- with propagation along the z axis.
-"""
-# TODO update module docs
-
-from typing import List, Tuple
-import numpy
-from numpy.linalg import norm
-import scipy.sparse as sparse
-
-from .. import vec, unvec, dx_lists_t, field_t, vfield_t
-from . import operators
-
-
-__author__ = 'Jan Petykiewicz'
-
-
-def operator_e(omega: complex,
-               dxes: dx_lists_t,
-               epsilon: vfield_t,
-               mu: vfield_t = None,
-               ) -> sparse.spmatrix:
-    if numpy.any(numpy.equal(mu, None)):
-        mu = numpy.ones_like(epsilon)
-
-    Dfx, Dfy = operators.deriv_forward(dxes[0])
-    Dbx, Dby = operators.deriv_back(dxes[1])
-
-    eps_parts = numpy.split(epsilon, 3)
-    eps_xy = sparse.diags(numpy.hstack((eps_parts[0], eps_parts[1])))
-    eps_z_inv = sparse.diags(1 / eps_parts[2])
-
-    mu_parts = numpy.split(mu, 3)
-    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 + \
-        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
-    return op
-
-
-def operator_h(omega: complex,
-               dxes: dx_lists_t,
-               epsilon: vfield_t,
-               mu: vfield_t = None,
-               ) -> sparse.spmatrix:
-    """
-    Waveguide operator of the form
-
-    omega**2 * epsilon * mu +
-    epsilon * [[-Dy], [Dx]] / epsilon * [-Dy, Dx] +
-    [[Dx], [Dy]] / mu * [Dx, Dy] * mu
-
-    for use with a field vector of the form [H_x, H_y].
-
-    This operator can be used to form an eigenvalue problem of the form
-    A @ [H_x, H_y] = wavenumber**2 * [H_x, H_y]
-
-    which can then be solved for the eigenmodes of the system (an exp(-i * wavenumber * z)
-    z-dependence is assumed for the fields).
-
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Sparse matrix representation of the operator
-    """
-    if numpy.any(numpy.equal(mu, None)):
-        mu = numpy.ones_like(epsilon)
-
-    Dfx, Dfy = operators.deriv_forward(dxes[0])
-    Dbx, Dby = operators.deriv_back(dxes[1])
-
-    eps_parts = numpy.split(epsilon, 3)
-    eps_yx = sparse.diags(numpy.hstack((eps_parts[1], eps_parts[0])))
-    eps_z_inv = sparse.diags(1 / eps_parts[2])
-
-    mu_parts = numpy.split(mu, 3)
-    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 + \
-        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
-
-    return op
-
-
-def normalized_fields_e(e_xy: numpy.ndarray,
-                        wavenumber: complex,
-                        omega: complex,
-                        dxes: dx_lists_t,
-                        epsilon: vfield_t,
-                        mu: vfield_t = None,
-                        prop_phase: float = 0,
-                        ) -> Tuple[vfield_t, vfield_t]:
-    """
-    Given a vector e_xy containing the vectorized E_x and E_y fields,
-     returns normalized, vectorized E and H fields for the system.
-
-    :param e_xy: Vector containing E_x and E_y fields
-    :param wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :param prop_phase: Phase shift (dz * corrected_wavenumber) over 1 cell in propagation direction.
-        Default 0 (continuous propagation direction, i.e. dz->0).
-    :return: Normalized, vectorized (e, h) containing all vector components.
-    """
-    e = exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon) @ e_xy
-    h = exy2h(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ e_xy
-    e_norm, h_norm = _normalized_fields(e=e, h=h, omega=omega, dxes=dxes, epsilon=epsilon,
-                                        mu=mu, prop_phase=prop_phase)
-    return e_norm, h_norm
-
-
-def normalized_fields_h(h_xy: numpy.ndarray,
-                        wavenumber: complex,
-                        omega: complex,
-                        dxes: dx_lists_t,
-                        epsilon: vfield_t,
-                        mu: vfield_t = None,
-                        prop_phase: float = 0,
-                        ) -> Tuple[vfield_t, vfield_t]:
-    """
-    Given a vector e_xy containing the vectorized E_x and E_y fields,
-     returns normalized, vectorized E and H fields for the system.
-
-    :param e_xy: Vector containing E_x and E_y fields
-    :param wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :param dxes_prop: Grid cell width in the propagation direction. Default 0 (continuous).
-    :return: Normalized, vectorized (e, h) containing all vector components.
-    """
-    e = hxy2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ h_xy
-    h = hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu) @ h_xy
-    e_norm, h_norm = _normalized_fields(e=e, h=h, omega=omega, dxes=dxes, epsilon=epsilon,
-                                        mu=mu, prop_phase=prop_phase)
-    return e_norm, h_norm
-
-
-def _normalized_fields(e: numpy.ndarray,
-                       h: numpy.ndarray,
-                       omega: complex,
-                       dxes: dx_lists_t,
-                       epsilon: vfield_t,
-                       mu: vfield_t = None,
-                       prop_phase: float = 0,
-                       ) -> Tuple[vfield_t, vfield_t]:
-    # TODO documentation
-    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
-    assert Sz_tavg > 0, 'Found a mode propagating in the wrong direction! Sz_tavg={}'.format(Sz_tavg)
-
-    energy = 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
-
-    # Try to break symmetry to assign a consistent sign [experimental TODO]
-    E_weighted = unvec(e * energy * numpy.exp(1j * norm_angle), shape)
-    sign = numpy.sign(E_weighted[:, :max(shape[0]//2, 1), :max(shape[1]//2, 1)].real.sum())
-
-    norm_factor = sign * norm_amplitude * numpy.exp(1j * norm_angle)
-
-    e *= norm_factor
-    h *= norm_factor
-
-    return e, h
-
-
-def exy2h(wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfield_t,
-          mu: vfield_t = None
-          ) -> sparse.spmatrix:
-    """
-    Operator which transforms the vector e_xy containing the vectorized E_x and E_y fields,
-     into a vectorized H containing all three H components
-
-    :param wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Sparse matrix representing the operator
-    """
-    e2hop = e2h(wavenumber=wavenumber, omega=omega, dxes=dxes, mu=mu)
-    return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)
-
-
-def hxy2e(wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfield_t,
-          mu: vfield_t = None
-          ) -> sparse.spmatrix:
-    """
-    Operator which transforms the vector h_xy containing the vectorized H_x and H_y fields,
-     into a vectorized E containing all three E components
-
-    :param wavenumber: Wavenumber satisfying `operator_h(...) @ h_xy == wavenumber**2 * h_xy`
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Sparse matrix representing the operator
-    """
-    h2eop = h2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon)
-    return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)
-
-
-def hxy2h(wavenumber: complex,
-          dxes: dx_lists_t,
-          mu: vfield_t = None
-          ) -> sparse.spmatrix:
-    """
-    Operator which transforms the vector h_xy containing the vectorized H_x and H_y fields,
-     into a vectorized H containing all three H components
-
-    :param wavenumber: Wavenumber satisfying `operator_h(...) @ h_xy == wavenumber**2 * h_xy`
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Sparse matrix representing the operator
-    """
-    Dfx, Dfy = operators.deriv_forward(dxes[0])
-    hxy2hz = sparse.hstack((Dfx, Dfy)) / (1j * wavenumber)
-
-    if not numpy.any(numpy.equal(mu, None)):
-        mu_parts = numpy.split(mu, 3)
-        mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1])))
-        mu_z_inv = sparse.diags(1 / mu_parts[2])
-
-        hxy2hz = mu_z_inv @ hxy2hz @ mu_xy
-
-    n_pts = dxes[1][0].size * dxes[1][1].size
-    op = sparse.vstack((sparse.eye(2 * n_pts),
-                        hxy2hz))
-    return op
-
-
-def exy2e(wavenumber: complex,
-          dxes: dx_lists_t,
-          epsilon: vfield_t,
-          ) -> sparse.spmatrix:
-    """
-    Operator which transforms the vector e_xy containing the vectorized E_x and E_y fields,
-     into a vectorized E containing all three E components
-
-    :param wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :return: Sparse matrix representing the operator
-    """
-    Dbx, Dby = operators.deriv_back(dxes[1])
-    exy2ez = sparse.hstack((Dbx, Dby)) / (1j * wavenumber)
-
-    if not numpy.any(numpy.equal(epsilon, None)):
-        epsilon_parts = numpy.split(epsilon, 3)
-        epsilon_xy = sparse.diags(numpy.hstack((epsilon_parts[0], epsilon_parts[1])))
-        epsilon_z_inv = sparse.diags(1 / epsilon_parts[2])
-
-        exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy
-
-    n_pts = dxes[0][0].size * dxes[0][1].size
-    op = sparse.vstack((sparse.eye(2 * n_pts),
-                        exy2ez))
-    return op
-
-
-def e2h(wavenumber: complex,
-        omega: complex,
-        dxes: dx_lists_t,
-        mu: vfield_t = None
-        ) -> sparse.spmatrix:
-    """
-    Returns an operator which, when applied to a vectorized E eigenfield, produces
-     the vectorized H eigenfield.
-
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Sparse matrix representation of the operator
-    """
-    op = curl_e(wavenumber, dxes) / (-1j * omega)
-    if not numpy.any(numpy.equal(mu, None)):
-        op = sparse.diags(1 / mu) @ op
-    return op
-
-
-def h2e(wavenumber: complex,
-        omega: complex,
-        dxes: dx_lists_t,
-        epsilon: vfield_t
-        ) -> sparse.spmatrix:
-    """
-    Returns an operator which, when applied to a vectorized H eigenfield, produces
-     the vectorized E eigenfield.
-
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :return: Sparse matrix representation of the operator
-    """
-    op = sparse.diags(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes)
-    return op
-
-
-def curl_e(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
-    """
-    Discretized curl operator for use with the waveguide E field.
-
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :return: Sparse matrix representation of the operator
-    """
-    n = 1
-    for d in dxes[0]:
-        n *= len(d)
-
-    Bz = -1j * wavenumber * sparse.eye(n)
-    Dfx, Dfy = operators.deriv_forward(dxes[0])
-    return operators.cross([Dfx, Dfy, Bz])
-
-
-def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
-    """
-    Discretized curl operator for use with the waveguide H field.
-
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :return: Sparse matrix representation of the operator
-    """
-    n = 1
-    for d in dxes[1]:
-        n *= len(d)
-
-    Bz = -1j * wavenumber * sparse.eye(n)
-    Dbx, Dby = operators.deriv_back(dxes[1])
-    return operators.cross([Dbx, Dby, Bz])
-
-
-def h_err(h: vfield_t,
-          wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfield_t,
-          mu: vfield_t = None
-          ) -> float:
-    """
-    Calculates the relative error in the H field
-
-    :param h: Vectorized H field
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Relative error norm(OP @ h) / norm(h)
-    """
-    ce = curl_e(wavenumber, dxes)
-    ch = curl_h(wavenumber, dxes)
-
-    eps_inv = sparse.diags(1 / epsilon)
-
-    if numpy.any(numpy.equal(mu, None)):
-        op = ce @ eps_inv @ ch @ h - omega ** 2 * h
-    else:
-        op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h)
-
-    return norm(op) / norm(h)
-
-
-def e_err(e: vfield_t,
-          wavenumber: complex,
-          omega: complex,
-          dxes: dx_lists_t,
-          epsilon: vfield_t,
-          mu: vfield_t = None
-          ) -> float:
-    """
-    Calculates the relative error in the E field
-
-    :param e: Vectorized E field
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Relative error norm(OP @ e) / norm(e)
-    """
-    ce = curl_e(wavenumber, dxes)
-    ch = curl_h(wavenumber, dxes)
-
-    if numpy.any(numpy.equal(mu, None)):
-        op = ch @ ce @ e - omega ** 2 * (epsilon * e)
-    else:
-        mu_inv = sparse.diags(1 / mu)
-        op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
-
-    return norm(op) / norm(e)
-
-
-def cylindrical_operator(omega: complex,
-                         dxes: dx_lists_t,
-                         epsilon: vfield_t,
-                         r0: float,
-                         ) -> sparse.spmatrix:
-    """
-    Cylindrical coordinate waveguide operator of the form
-
-    TODO
-
-    for use with a field vector of the form [E_r, E_y].
-
-    This operator can be used to form an eigenvalue problem of the form
-    A @ [E_r, E_y] = wavenumber**2 * [E_r, E_y]
-
-    which can then be solved for the eigenmodes of the system (an exp(-i * wavenumber * theta)
-     theta-dependence is assumed for the fields).
-
-    :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param r0: Radius of curvature for the simulation. This should be the minimum value of
-        r within the simulation domain.
-    :return: Sparse matrix representation of the operator
-    """
-
-    Dfx, Dfy = operators.deriv_forward(dxes[0])
-    Dbx, Dby = operators.deriv_back(dxes[1])
-
-    rx = r0 + numpy.cumsum(dxes[0][0])
-    ry = r0 + dxes[0][0]/2.0 + numpy.cumsum(dxes[1][0])
-    tx = rx/r0
-    ty = ry/r0
-
-    Tx = sparse.diags(vec(tx[:, None].repeat(dxes[0][1].size, axis=1)))
-    Ty = sparse.diags(vec(ty[:, None].repeat(dxes[1][1].size, axis=1)))
-
-    eps_parts = numpy.split(epsilon, 3)
-    eps_x = sparse.diags(eps_parts[0])
-    eps_y = sparse.diags(eps_parts[1])
-    eps_z_inv = sparse.diags(1 / eps_parts[2])
-
-    pa = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dbx, Dby))
-    pb = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dby, Dbx))
-    a0 = Ty @ eps_x + omega**-2 * Dby @ Ty @ Dfy
-    a1 = Tx @ eps_y + omega**-2 * Dbx @ Ty @ Dfx
-    b0 = Dbx @ Ty @ Dfy
-    b1 = Dby @ Ty @ Dfx
-
-    diag = sparse.block_diag
-    op = (omega**2 * diag((Tx, Ty)) + pa) @ diag((a0, a1)) + \
-        - (sparse.bmat(((None, Ty), (Tx, None))) + omega**-2 * pb) @ diag((b0, b1))
-
-    return op
-

meanas/fdfd/waveguide_2d.py

@@ -0,0 +1,607 @@
+"""
+Operators and helper functions for waveguides with unchanging cross-section.
+
+The propagation direction is chosen to be along the z axis, and all fields
+are given an implicit z-dependence of the form `exp(-1 * wavenumber * z)`.
+
+As the z-dependence is known, all the functions in this file assume a 2D grid
+ (i.e. `dxes = [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...]], [[dx_h_0, ...], [dy_h_0, ...]]]`).
+"""
+# TODO update module docs
+
+from typing import List, Tuple
+import numpy
+from numpy.linalg import norm
+import scipy.sparse as sparse
+
+from .. import vec, unvec, dx_lists_t, field_t, vfield_t
+from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
+from . import operators
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def operator_e(omega: complex,
+               dxes: dx_lists_t,
+               epsilon: vfield_t,
+               mu: vfield_t = None,
+               ) -> sparse.spmatrix:
+    """
+    Waveguide operator of the form
+
+        omega**2 * mu * epsilon +
+        mu * [[-Dy], [Dx]] / mu * [-Dy, Dx] +
+        [[Dx], [Dy]] / epsilon * [Dx, Dy] * epsilon
+
+    for use with a field vector of the form `cat([E_x, E_y])`.
+
+    More precisely, the operator is
+    $$ \\omega^2 \\mu_{yx} \\epsilon_{xy} +
+       \\mu_{yx} \\begin{bmatrix} -D_{by} \\\\
+                                   D_{bx} \\end{bmatrix} \\mu_z^{-1}
+                 \\begin{bmatrix} -D_{fy} & D_{fx} \\end{bmatrix} +
+      \\begin{bmatrix} D_{fx} \\\\
+                       D_{fy} \\end{bmatrix} \\epsilon_z^{-1}
+                 \\begin{bmatrix} D_{bx} & D_{by} \\end{bmatrix} \\epsilon_{xy} $$
+
+    where
+    \\( \\epsilon_{xy} = \\begin{bmatrix}
+                                 \\epsilon_x & 0 \\\\
+                                           0 & \\epsilon_y
+                                 \\end{bmatrix} \\),
+    \\( \\mu_{yx} = \\begin{bmatrix}
+                           \\mu_y & 0 \\\\
+                                0 & \\mu_x
+                           \\end{bmatrix} \\),
+    \\( D_{fx} \\) and \\( D_{bx} \\) are the forward and backward derivatives along x,
+    and each \\( \\epsilon_x, \\mu_y, \\) etc. is a diagonal matrix representing
+
+
+    This operator can be used to form an eigenvalue problem of the form
+    `operator_e(...) @ [E_x, E_y] = wavenumber**2 * [E_x, E_y]`
+
+    which can then be solved for the eigenmodes of the system (an `exp(-i * wavenumber * z)`
+    z-dependence is assumed for the fields).
+
+    Args:
+        omega: The angular frequency of the system.
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representation of the operator.
+    """
+    if numpy.any(numpy.equal(mu, None)):
+        mu = numpy.ones_like(epsilon)
+
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    Dbx, Dby = operators.deriv_back(dxes[1])
+
+    eps_parts = numpy.split(epsilon, 3)
+    eps_xy = sparse.diags(numpy.hstack((eps_parts[0], eps_parts[1])))
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    mu_parts = numpy.split(mu, 3)
+    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 + \
+        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
+    return op
+
+
+def operator_h(omega: complex,
+               dxes: dx_lists_t,
+               epsilon: vfield_t,
+               mu: vfield_t = None,
+               ) -> sparse.spmatrix:
+    """
+    Waveguide operator of the form
+
+        omega**2 * epsilon * mu +
+        epsilon * [[-Dy], [Dx]] / epsilon * [-Dy, Dx] +
+        [[Dx], [Dy]] / mu * [Dx, Dy] * mu
+
+    for use with a field vector of the form `cat([H_x, H_y])`.
+
+    More precisely, the operator is
+    $$ \\omega^2 \\epsilon_{yx} \\mu_{xy} +
+       \\epsilon_{yx} \\begin{bmatrix} -D_{fy} \\\\
+                                        D_{fx} \\end{bmatrix} \\epsilon_z^{-1}
+                 \\begin{bmatrix} -D_{by} & D_{bx} \\end{bmatrix} +
+      \\begin{bmatrix} D_{bx} \\\\
+                       D_{by} \\end{bmatrix} \\mu_z^{-1}
+                 \\begin{bmatrix} D_{fx} & D_{fy} \\end{bmatrix} \\mu_{xy} $$
+
+    where
+    \\( \\epsilon_{yx} = \\begin{bmatrix}
+                                 \\epsilon_y & 0 \\\\
+                                           0 & \\epsilon_x
+                                 \\end{bmatrix} \\),
+    \\( \\mu_{xy} = \\begin{bmatrix}
+                           \\mu_x & 0 \\\\
+                                0 & \\mu_y
+                           \\end{bmatrix} \\),
+    \\( D_{fx} \\) and \\( D_{bx} \\) are the forward and backward derivatives along x,
+    and each \\( \\epsilon_x, \\mu_y, \\) etc. is a diagonal matrix.
+
+
+    This operator can be used to form an eigenvalue problem of the form
+    `operator_h(...) @ [H_x, H_y] = wavenumber**2 * [H_x, H_y]`
+
+    which can then be solved for the eigenmodes of the system (an `exp(-i * wavenumber * z)`
+    z-dependence is assumed for the fields).
+
+    Args:
+        omega: The angular frequency of the system.
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representation of the operator.
+    """
+    if numpy.any(numpy.equal(mu, None)):
+        mu = numpy.ones_like(epsilon)
+
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    Dbx, Dby = operators.deriv_back(dxes[1])
+
+    eps_parts = numpy.split(epsilon, 3)
+    eps_yx = sparse.diags(numpy.hstack((eps_parts[1], eps_parts[0])))
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    mu_parts = numpy.split(mu, 3)
+    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 + \
+        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
+
+    return op
+
+
+def normalized_fields_e(e_xy: numpy.ndarray,
+                        wavenumber: complex,
+                        omega: complex,
+                        dxes: dx_lists_t,
+                        epsilon: vfield_t,
+                        mu: vfield_t = None,
+                        prop_phase: float = 0,
+                        ) -> Tuple[vfield_t, vfield_t]:
+    """
+    Given a vector `e_xy` containing the vectorized E_x and E_y fields,
+     returns normalized, vectorized E and H fields for the system.
+
+    Args:
+        e_xy: Vector containing E_x and E_y fields
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
+                    It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+        prop_phase: Phase shift `(dz * corrected_wavenumber)` over 1 cell in propagation direction.
+                    Default 0 (continuous propagation direction, i.e. dz->0).
+
+    Returns:
+        `(e, h)`, where each field is vectorized, normalized,
+        and contains all three vector components.
+    """
+    e = exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon) @ e_xy
+    h = exy2h(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ e_xy
+    e_norm, h_norm = _normalized_fields(e=e, h=h, omega=omega, dxes=dxes, epsilon=epsilon,
+                                        mu=mu, prop_phase=prop_phase)
+    return e_norm, h_norm
+
+
+def normalized_fields_h(h_xy: numpy.ndarray,
+                        wavenumber: complex,
+                        omega: complex,
+                        dxes: dx_lists_t,
+                        epsilon: vfield_t,
+                        mu: vfield_t = None,
+                        prop_phase: float = 0,
+                        ) -> Tuple[vfield_t, vfield_t]:
+    """
+    Given a vector `h_xy` containing the vectorized H_x and H_y fields,
+     returns normalized, vectorized E and H fields for the system.
+
+    Args:
+        h_xy: Vector containing H_x and H_y fields
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
+                    It should satisfy `operator_h() @ h_xy == wavenumber**2 * h_xy`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+        prop_phase: Phase shift `(dz * corrected_wavenumber)` over 1 cell in propagation direction.
+                    Default 0 (continuous propagation direction, i.e. dz->0).
+
+    Returns:
+        `(e, h)`, where each field is vectorized, normalized,
+        and contains all three vector components.
+    """
+    e = hxy2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ h_xy
+    h = hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu) @ h_xy
+    e_norm, h_norm = _normalized_fields(e=e, h=h, omega=omega, dxes=dxes, epsilon=epsilon,
+                                        mu=mu, prop_phase=prop_phase)
+    return e_norm, h_norm
+
+
+def _normalized_fields(e: numpy.ndarray,
+                       h: numpy.ndarray,
+                       omega: complex,
+                       dxes: dx_lists_t,
+                       epsilon: vfield_t,
+                       mu: vfield_t = None,
+                       prop_phase: float = 0,
+                       ) -> Tuple[vfield_t, vfield_t]:
+    # TODO documentation
+    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
+    assert Sz_tavg > 0, 'Found a mode propagating in the wrong direction! Sz_tavg={}'.format(Sz_tavg)
+
+    energy = 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
+
+    # Try to break symmetry to assign a consistent sign [experimental TODO]
+    E_weighted = unvec(e * energy * numpy.exp(1j * norm_angle), shape)
+    sign = numpy.sign(E_weighted[:, :max(shape[0]//2, 1), :max(shape[1]//2, 1)].real.sum())
+
+    norm_factor = sign * norm_amplitude * numpy.exp(1j * norm_angle)
+
+    e *= norm_factor
+    h *= norm_factor
+
+    return e, h
+
+
+def exy2h(wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> sparse.spmatrix:
+    """
+    Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
+     into a vectorized H containing all three H components
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
+                    It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representing the operator.
+    """
+    e2hop = e2h(wavenumber=wavenumber, omega=omega, dxes=dxes, mu=mu)
+    return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)
+
+
+def hxy2e(wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> sparse.spmatrix:
+    """
+    Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
+     into a vectorized E containing all three E components
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
+                    It should satisfy `operator_h() @ h_xy == wavenumber**2 * h_xy`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representing the operator.
+    """
+    h2eop = h2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon)
+    return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)
+
+
+def hxy2h(wavenumber: complex,
+          dxes: dx_lists_t,
+          mu: vfield_t = None
+          ) -> sparse.spmatrix:
+    """
+    Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
+     into a vectorized H containing all three H components
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
+                    It should satisfy `operator_h() @ h_xy == wavenumber**2 * h_xy`
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representing the operator.
+    """
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    hxy2hz = sparse.hstack((Dfx, Dfy)) / (1j * wavenumber)
+
+    if not numpy.any(numpy.equal(mu, None)):
+        mu_parts = numpy.split(mu, 3)
+        mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1])))
+        mu_z_inv = sparse.diags(1 / mu_parts[2])
+
+        hxy2hz = mu_z_inv @ hxy2hz @ mu_xy
+
+    n_pts = dxes[1][0].size * dxes[1][1].size
+    op = sparse.vstack((sparse.eye(2 * n_pts),
+                        hxy2hz))
+    return op
+
+
+def exy2e(wavenumber: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          ) -> sparse.spmatrix:
+    """
+    Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
+     into a vectorized E containing all three E components
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+                    It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+
+    Returns:
+        Sparse matrix representing the operator.
+    """
+    Dbx, Dby = operators.deriv_back(dxes[1])
+    exy2ez = sparse.hstack((Dbx, Dby)) / (1j * wavenumber)
+
+    if not numpy.any(numpy.equal(epsilon, None)):
+        epsilon_parts = numpy.split(epsilon, 3)
+        epsilon_xy = sparse.diags(numpy.hstack((epsilon_parts[0], epsilon_parts[1])))
+        epsilon_z_inv = sparse.diags(1 / epsilon_parts[2])
+
+        exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy
+
+    n_pts = dxes[0][0].size * dxes[0][1].size
+    op = sparse.vstack((sparse.eye(2 * n_pts),
+                        exy2ez))
+    return op
+
+
+def e2h(wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        mu: vfield_t = None
+        ) -> sparse.spmatrix:
+    """
+    Returns an operator which, when applied to a vectorized E eigenfield, produces
+     the vectorized H eigenfield.
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representation of the operator.
+    """
+    op = curl_e(wavenumber, dxes) / (-1j * omega)
+    if not numpy.any(numpy.equal(mu, None)):
+        op = sparse.diags(1 / mu) @ op
+    return op
+
+
+def h2e(wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfield_t
+        ) -> sparse.spmatrix:
+    """
+    Returns an operator which, when applied to a vectorized H eigenfield, produces
+     the vectorized E eigenfield.
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+
+    Returns:
+        Sparse matrix representation of the operator.
+    """
+    op = sparse.diags(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes)
+    return op
+
+
+def curl_e(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Discretized curl operator for use with the waveguide E field.
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+
+    Return:
+        Sparse matrix representation of the operator.
+    """
+    n = 1
+    for d in dxes[0]:
+        n *= len(d)
+
+    Bz = -1j * wavenumber * sparse.eye(n)
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    return operators.cross([Dfx, Dfy, Bz])
+
+
+def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Discretized curl operator for use with the waveguide H field.
+
+    Args:
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+
+    Return:
+        Sparse matrix representation of the operator.
+    """
+    n = 1
+    for d in dxes[1]:
+        n *= len(d)
+
+    Bz = -1j * wavenumber * sparse.eye(n)
+    Dbx, Dby = operators.deriv_back(dxes[1])
+    return operators.cross([Dbx, Dby, Bz])
+
+
+def h_err(h: vfield_t,
+          wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> float:
+    """
+    Calculates the relative error in the H field
+
+    Args:
+        h: Vectorized H field
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Relative error `norm(A_h @ h) / norm(h)`.
+    """
+    ce = curl_e(wavenumber, dxes)
+    ch = curl_h(wavenumber, dxes)
+
+    eps_inv = sparse.diags(1 / epsilon)
+
+    if numpy.any(numpy.equal(mu, None)):
+        op = ce @ eps_inv @ ch @ h - omega ** 2 * h
+    else:
+        op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h)
+
+    return norm(op) / norm(h)
+
+
+def e_err(e: vfield_t,
+          wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> float:
+    """
+    Calculates the relative error in the E field
+
+    Args:
+        e: Vectorized E field
+        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Relative error `norm(A_e @ e) / norm(e)`.
+    """
+    ce = curl_e(wavenumber, dxes)
+    ch = curl_h(wavenumber, dxes)
+
+    if numpy.any(numpy.equal(mu, None)):
+        op = ch @ ce @ e - omega ** 2 * (epsilon * e)
+    else:
+        mu_inv = sparse.diags(1 / mu)
+        op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
+
+    return norm(op) / norm(e)
+
+
+def solve_modes(mode_numbers: List[int],
+                omega: complex,
+                dxes: dx_lists_t,
+                epsilon: vfield_t,
+                mu: vfield_t = None,
+                mode_margin: int = 2,
+                ) -> Tuple[List[vfield_t], List[complex]]:
+    """
+    Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
+
+    Args:
+       mode_numbers: List of 0-indexed mode numbers to solve for
+       omega: Angular frequency of the simulation
+       dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+       epsilon: Dielectric constant
+       mu: Magnetic permeability (default 1 everywhere)
+       mode_margin: The eigensolver will actually solve for `(max(mode_number) + mode_margin)`
+            modes, but only return the target mode. Increasing this value can improve the solver's
+            ability to find the correct mode. Default 2.
+
+    Returns:
+        (e_xys, wavenumbers)
+    """
+
+    '''
+    Solve for the largest-magnitude eigenvalue of the real operator
+    '''
+    dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
+    A_r = waveguide.operator_e(numpy.real(omega), dxes_real, numpy.real(epsilon), numpy.real(mu))
+
+    eigvals, eigvecs = signed_eigensolve(A_r, max(mode_number) + mode_margin)
+    e_xys = eigvecs[:, -(numpy.array(mode_number) + 1)]
+
+    '''
+    Now solve for the eigenvector of the full operator, using the real operator's
+     eigenvector as an initial guess for Rayleigh quotient iteration.
+    '''
+    A = waveguide.operator_e(omega, dxes, epsilon, mu)
+    eigvals, e_xys = rayleigh_quotient_iteration(A, e_xys)
+
+    # Calculate the wave-vector (force the real part to be positive)
+    wavenumbers = numpy.sqrt(eigvals)
+    wavenumbers *= numpy.sign(numpy.real(wavenumbers))
+
+    return e_xys, wavenumbers
+
+
+def solve_mode(mode_number: int,
+               *args,
+               **kwargs
+               ) -> Tuple[vfield_t, complex]:
+    """
+    Wrapper around `solve_modes()` that solves for a single mode.
+
+    Args:
+       mode_number: 0-indexed mode number to solve for
+       *args: passed to `solve_modes()`
+       **kwargs: passed to `solve_modes()`
+
+    Returns:
+        (e_xy, wavenumber)
+    """
+    return solve_modes(mode_numbers=[mode_number], *args, **kwargs)

meanas/fdfd/waveguide_3d.py

@@ -0,0 +1,236 @@
+"""
+Tools for working with waveguide modes in 3D domains.
+
+This module relies heavily on `waveguide_2d` and mostly just transforms
+its parameters into 2D equivalents and expands the results back into 3D.
+"""
+from typing import Dict, List, Tuple
+import numpy
+import scipy.sparse as sparse
+
+from .. import vec, unvec, dx_lists_t, vfield_t, field_t
+from . import operators, waveguide_2d, functional
+
+
+def solve_mode(mode_number: int,
+               omega: complex,
+               dxes: dx_lists_t,
+               axis: int,
+               polarity: int,
+               slices: List[slice],
+               epsilon: field_t,
+               mu: field_t = None,
+               ) -> Dict[str, complex or numpy.ndarray]:
+    """
+    Given a 3D grid, selects a slice from the grid and attempts to
+     solve for an eigenmode propagating through that slice.
+
+    Args:
+        mode_number: Number of the mode, 0-indexed
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        axis: Propagation axis (0=x, 1=y, 2=z)
+        polarity: Propagation direction (+1 for +ve, -1 for -ve)
+        slices: `epsilon[tuple(slices)]` is used to select the portion of the grid to use
+                as the waveguide cross-section. `slices[axis]` should select only one item.
+        epsilon: Dielectric constant
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        `{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}`
+    """
+    if mu is None:
+        mu = numpy.ones_like(epsilon)
+
+    slices = tuple(slices)
+
+    '''
+    Solve the 2D problem in the specified plane
+    '''
+    # Define rotation to set z as propagation direction
+    order = numpy.roll(range(3), 2 - axis)
+    reverse_order = numpy.roll(range(3), axis - 2)
+
+    # Find dx in propagation direction
+    dxab_forward = numpy.array([dx[order[2]][slices[order[2]]] for dx in dxes])
+    dx_prop = 0.5 * sum(dxab_forward)[0]
+
+    # Reduce to 2D and solve the 2D problem
+    args_2d = {
+        'omega': omega,
+        'dxes': [[dx[i][slices[i]] for i in order[:2]] for dx in dxes],
+        'epsilon': vec([epsilon[i][slices].transpose(order) for i in order]),
+        'mu': vec([mu[i][slices].transpose(order) for i in order]),
+    }
+    e_xy, wavenumber_2d = waveguide_2d.solve_mode(mode_number, **args_2d)
+
+    '''
+    Apply corrections and expand to 3D
+    '''
+    # Correct wavenumber to account for numerical dispersion.
+    wavenumber = 2/dx_prop * numpy.arcsin(wavenumber_2d * dx_prop/2)
+
+    shape = [d.size for d in args_2d['dxes'][0]]
+    ve, vh = waveguide.normalized_fields_e(e_xy, wavenumber=wavenumber_2d, **args_2d, prop_phase=dx_prop * wavenumber)
+    e = unvec(ve, shape)
+    h = unvec(vh, shape)
+
+    # Adjust for propagation direction
+    h *= polarity
+
+    # Apply phase shift to H-field
+    h[:2] *= numpy.exp(-1j * polarity * 0.5 * wavenumber * dx_prop)
+    e[2]  *= numpy.exp(-1j * polarity * 0.5 * wavenumber * dx_prop)
+
+    # Expand E, H to full epsilon space we were given
+    E = numpy.zeros_like(epsilon, dtype=complex)
+    H = numpy.zeros_like(epsilon, dtype=complex)
+    for a, o in enumerate(reverse_order):
+        E[(a, *slices)] = e[o][:, :, None].transpose(reverse_order)
+        H[(a, *slices)] = h[o][:, :, None].transpose(reverse_order)
+
+    results = {
+        'wavenumber': wavenumber,
+        'wavenumber_2d': wavenumber_2d,
+        'H': H,
+        'E': E,
+    }
+    return results
+
+
+def compute_source(E: field_t,
+                   wavenumber: complex,
+                   omega: complex,
+                   dxes: dx_lists_t,
+                   axis: int,
+                   polarity: int,
+                   slices: List[slice],
+                   epsilon: field_t,
+                   mu: field_t = None,
+                   ) -> field_t:
+    """
+    Given an eigenmode obtained by `solve_mode`, returns the current source distribution
+    necessary to position a unidirectional source at the slice location.
+
+    Args:
+        E: E-field of the mode
+        wavenumber: Wavenumber of the mode
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        axis: Propagation axis (0=x, 1=y, 2=z)
+        polarity: Propagation direction (+1 for +ve, -1 for -ve)
+        slices: `epsilon[tuple(slices)]` is used to select the portion of the grid to use
+                as the waveguide cross-section. `slices[axis]` should select only one item.
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        J distribution for the unidirectional source
+    """
+    E_expanded = expand_e(E=E, dxes=dxes, wavenumber=wavenumber, axis=axis,
+                          polarity=polarity, slices=slices)
+
+    smask = [slice(None)] * 4
+    if polarity > 0:
+        smask[axis + 1] = slice(slices[axis].start, None)
+    else:
+        smask[axis + 1] = slice(None, slices[axis].stop)
+
+    mask = numpy.zeros_like(E_expanded, dtype=int)
+    mask[tuple(smask)] = 1
+
+    masked_e2j = operators.e_boundary_source(mask=vec(mask), omega=omega, dxes=dxes, epsilon=vec(epsilon), mu=vec(mu))
+    J = unvec(masked_e2j @ vec(E_expanded), E.shape[1:])
+    return J
+
+
+def compute_overlap_e(E: field_t,
+                      wavenumber: complex,
+                      dxes: dx_lists_t,
+                      axis: int,
+                      polarity: int,
+                      slices: List[slice],
+                      ) -> field_t:                 # TODO DOCS
+    """
+    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)
+    [assumes reflection symmetry].
+
+    Args:
+        E: E-field of the mode
+        H: H-field of the mode (advanced by half of a Yee cell from E)
+        wavenumber: Wavenumber of the mode
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        axis: Propagation axis (0=x, 1=y, 2=z)
+        polarity: Propagation direction (+1 for +ve, -1 for -ve)
+        slices: `epsilon[tuple(slices)]` is used to select the portion of the grid to use
+                as the waveguide cross-section. slices[axis] should select only one item.
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        overlap_e such that `numpy.sum(overlap_e * other_e)` computes the overlap integral
+    """
+    slices = tuple(slices)
+
+    Ee = expand_e(E=E, wavenumber=wavenumber, dxes=dxes,
+                  axis=axis, polarity=polarity, slices=slices)
+
+    start, stop = sorted((slices[axis].start, slices[axis].start - 2 * polarity))
+
+    slices2 = list(slices)
+    slices2[axis] = slice(start, stop)
+    slices2 = (slice(None), *slices2)
+
+    Etgt = numpy.zeros_like(Ee)
+    Etgt[slices2] = Ee[slices2]
+
+    Etgt /= (Etgt.conj() * Etgt).sum()
+    return Etgt
+
+
+def expand_e(E: field_t,
+             wavenumber: complex,
+             dxes: dx_lists_t,
+             axis: int,
+             polarity: int,
+             slices: List[slice],
+             ) -> field_t:
+    """
+    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
+    axis). This assumes the epsilon cross-section remains constant throughout the
+    entire domain; it is up to the caller to truncate the expansion to any regions
+    where it is valid.
+
+    Args:
+        E: E-field of the mode
+        wavenumber: Wavenumber of the mode
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types`
+        axis: Propagation axis (0=x, 1=y, 2=z)
+        polarity: Propagation direction (+1 for +ve, -1 for -ve)
+        slices: `epsilon[tuple(slices)]` is used to select the portion of the grid to use
+                as the waveguide cross-section. slices[axis] should select only one item.
+
+    Returns:
+        `E`, with the original field expanded along the specified `axis`.
+    """
+    slices = tuple(slices)
+
+    # Determine phase factors for parallel slices
+    a_shape = numpy.roll([1, -1, 1, 1], axis)
+    a_E = numpy.real(dxes[0][axis]).cumsum()
+    r_E = a_E - a_E[slices[axis]]
+    iphi = polarity * -1j * wavenumber
+    phase_E = numpy.exp(iphi * r_E).reshape(a_shape)
+
+    # Expand our slice to the entire grid using the phase factors
+    E_expanded = numpy.zeros_like(E)
+
+    slices_exp = list(slices)
+    slices_exp[axis] = slice(E.shape[axis + 1])
+    slices_exp = (slice(None), *slices_exp)
+
+    slices_in = (slice(None), *slices)
+
+    E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in]
+    return E_expanded

meanas/fdfd/waveguide_cyl.py

@@ -0,0 +1,138 @@
+"""
+Operators and helper functions for cylindrical waveguides with unchanging cross-section.
+
+WORK IN PROGRESS, CURRENTLY BROKEN
+
+As the z-dependence is known, all the functions in this file assume a 2D grid
+ (i.e. `dxes = [[[dr_e_0, dx_e_1, ...], [dy_e_0, ...]], [[dr_h_0, ...], [dy_h_0, ...]]]`).
+"""
+# TODO update module docs
+
+from typing import List, Tuple, Dict
+import numpy
+from numpy.linalg import norm
+import scipy.sparse as sparse
+
+from .. import vec, unvec, dx_lists_t, field_t, vfield_t
+from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
+from . import operators
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def cylindrical_operator(omega: complex,
+                         dxes: dx_lists_t,
+                         epsilon: vfield_t,
+                         r0: float,
+                         ) -> sparse.spmatrix:
+    """
+    Cylindrical coordinate waveguide operator of the form
+
+    TODO
+
+    for use with a field vector of the form `[E_r, E_y]`.
+
+    This operator can be used to form an eigenvalue problem of the form
+        A @ [E_r, E_y] = wavenumber**2 * [E_r, E_y]
+
+    which can then be solved for the eigenmodes of the system
+    (an `exp(-i * wavenumber * theta)` theta-dependence is assumed for the fields).
+
+    Args:
+        omega: The angular frequency of the system
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+        r0: Radius of curvature for the simulation. This should be the minimum value of
+            r within the simulation domain.
+
+    Returns:
+        Sparse matrix representation of the operator
+    """
+
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    Dbx, Dby = operators.deriv_back(dxes[1])
+
+    rx = r0 + numpy.cumsum(dxes[0][0])
+    ry = r0 + dxes[0][0]/2.0 + numpy.cumsum(dxes[1][0])
+    tx = rx/r0
+    ty = ry/r0
+
+    Tx = sparse.diags(vec(tx[:, None].repeat(dxes[0][1].size, axis=1)))
+    Ty = sparse.diags(vec(ty[:, None].repeat(dxes[1][1].size, axis=1)))
+
+    eps_parts = numpy.split(epsilon, 3)
+    eps_x = sparse.diags(eps_parts[0])
+    eps_y = sparse.diags(eps_parts[1])
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    pa = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dbx, Dby))
+    pb = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dby, Dbx))
+    a0 = Ty @ eps_x + omega**-2 * Dby @ Ty @ Dfy
+    a1 = Tx @ eps_y + omega**-2 * Dbx @ Ty @ Dfx
+    b0 = Dbx @ Ty @ Dfy
+    b1 = Dby @ Ty @ Dfx
+
+    diag = sparse.block_diag
+    op = (omega**2 * diag((Tx, Ty)) + pa) @ diag((a0, a1)) + \
+        - (sparse.bmat(((None, Ty), (Tx, None))) + omega**-2 * pb) @ diag((b0, b1))
+
+    return op
+
+
+def solve_mode(mode_number: int,
+               omega: complex,
+               dxes: dx_lists_t,
+               epsilon: vfield_t,
+               r0: float,
+               ) -> Dict[str, complex or field_t]:
+    """
+    TODO: fixup
+    Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
+     of the bent waveguide with the specified mode number.
+
+    Args:
+        mode_number: Number of the mode, 0-indexed
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters [dx_e, dx_h] as described in meanas.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.
+
+    Returns:
+        `{'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}`
+    """
+
+    '''
+    Solve for the largest-magnitude eigenvalue of the real operator
+    '''
+    dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
+
+    A_r = waveguide.cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0)
+    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + 3)
+    e_xy = eigvecs[:, -(mode_number+1)]
+
+    '''
+    Now solve for the eigenvector of the full operator, using the real operator's
+     eigenvector as an initial guess for Rayleigh quotient iteration.
+    '''
+    A = waveguide.cylindrical_operator(omega, dxes, epsilon, r0)
+    eigval, e_xy = rayleigh_quotient_iteration(A, e_xy)
+
+    # Calculate the wave-vector (force the real part to be positive)
+    wavenumber = numpy.sqrt(eigval)
+    wavenumber *= numpy.sign(numpy.real(wavenumber))
+
+    # TODO: Perform correction on wavenumber to account for numerical dispersion.
+
+    shape = [d.size for d in dxes[0]]
+    e_xy = numpy.hstack((e_xy, numpy.zeros(shape[0] * shape[1])))
+    fields = {
+        'wavenumber': wavenumber,
+        'E': unvec(e_xy, shape),
+#        'E': unvec(e, shape),
+#        'H': unvec(h, shape),
+    }
+
+    return fields

meanas/fdfd/waveguide_mode.py

@@ -1,307 +0,0 @@
-from typing import Dict, List, Tuple
-import numpy
-import scipy.sparse as sparse
-
-from .. import vec, unvec, dx_lists_t, vfield_t, field_t
-from . import operators, waveguide, functional
-from ..eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
-
-
-def vsolve_waveguide_mode_2d(mode_number: int,
-                             omega: complex,
-                             dxes: dx_lists_t,
-                             epsilon: vfield_t,
-                             mu: vfield_t = None,
-                             mode_margin: int = 2,
-                             ) -> Tuple[vfield_t, complex]:
-    """
-    Given a 2d region, attempts to solve for the eigenmode with the specified mode number.
-
-    :param mode_number: Number of the mode, 0-indexed.
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :param epsilon: Dielectric constant
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :param mode_margin: The eigensolver will actually solve for (mode_number + mode_margin)
-        modes, but only return the target mode. Increasing this value can improve the solver's
-        ability to find the correct mode. Default 2.
-    :return: (e_xy, wavenumber)
-    """
-
-    '''
-    Solve for the largest-magnitude eigenvalue of the real operator
-    '''
-    dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
-    A_r = waveguide.operator_e(numpy.real(omega), dxes_real, numpy.real(epsilon), numpy.real(mu))
-
-    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + mode_margin)
-    e_xy = eigvecs[:, -(mode_number + 1)]
-
-    '''
-    Now solve for the eigenvector of the full operator, using the real operator's
-     eigenvector as an initial guess for Rayleigh quotient iteration.
-    '''
-    A = waveguide.operator_e(omega, dxes, epsilon, mu)
-    eigval, e_xy = rayleigh_quotient_iteration(A, e_xy)
-
-    # Calculate the wave-vector (force the real part to be positive)
-    wavenumber = numpy.sqrt(eigval)
-    wavenumber *= numpy.sign(numpy.real(wavenumber))
-
-    return e_xy, wavenumber
-
-
-
-def solve_waveguide_mode(mode_number: int,
-                         omega: complex,
-                         dxes: dx_lists_t,
-                         axis: int,
-                         polarity: int,
-                         slices: List[slice],
-                         epsilon: field_t,
-                         mu: field_t = None,
-                         ) -> Dict[str, complex or numpy.ndarray]:
-    """
-    Given a 3D grid, selects a slice from the grid and attempts to
-     solve for an eigenmode propagating through that slice.
-
-    :param mode_number: Number of the mode, 0-indexed
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :param axis: Propagation axis (0=x, 1=y, 2=z)
-    :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
-    :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
-        as the waveguide cross-section. slices[axis] should select only one
-    :param epsilon: Dielectric constant
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :return: {'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}
-    """
-    if mu is None:
-        mu = numpy.ones_like(epsilon)
-
-    slices = tuple(slices)
-
-    '''
-    Solve the 2D problem in the specified plane
-    '''
-    # Define rotation to set z as propagation direction
-    order = numpy.roll(range(3), 2 - axis)
-    reverse_order = numpy.roll(range(3), axis - 2)
-
-    # Find dx in propagation direction
-    dxab_forward = numpy.array([dx[order[2]][slices[order[2]]] for dx in dxes])
-    dx_prop = 0.5 * sum(dxab_forward)[0]
-
-    # Reduce to 2D and solve the 2D problem
-    args_2d = {
-        'omega': omega,
-        'dxes': [[dx[i][slices[i]] for i in order[:2]] for dx in dxes],
-        'epsilon': vec([epsilon[i][slices].transpose(order) for i in order]),
-        'mu': vec([mu[i][slices].transpose(order) for i in order]),
-    }
-    e_xy, wavenumber_2d = vsolve_waveguide_mode_2d(mode_number, **args_2d)
-
-    '''
-    Apply corrections and expand to 3D
-    '''
-    # Correct wavenumber to account for numerical dispersion.
-    wavenumber = 2/dx_prop * numpy.arcsin(wavenumber_2d * dx_prop/2)
-    print(wavenumber_2d / wavenumber)
-
-    shape = [d.size for d in args_2d['dxes'][0]]
-    ve, vh = waveguide.normalized_fields_e(e_xy, wavenumber=wavenumber_2d, **args_2d, prop_phase=dx_prop * wavenumber)
-    e = unvec(ve, shape)
-    h = unvec(vh, shape)
-
-    # Adjust for propagation direction
-    h *= polarity
-
-    # Apply phase shift to H-field
-    h[:2] *= numpy.exp(-1j * polarity * 0.5 * wavenumber * dx_prop)
-    e[2]  *= numpy.exp(-1j * polarity * 0.5 * wavenumber * dx_prop)
-
-    # Expand E, H to full epsilon space we were given
-    E = numpy.zeros_like(epsilon, dtype=complex)
-    H = numpy.zeros_like(epsilon, dtype=complex)
-    for a, o in enumerate(reverse_order):
-        E[(a, *slices)] = e[o][:, :, None].transpose(reverse_order)
-        H[(a, *slices)] = h[o][:, :, None].transpose(reverse_order)
-
-    results = {
-        'wavenumber': wavenumber,
-        'wavenumber_2d': wavenumber_2d,
-        'H': H,
-        'E': E,
-    }
-    return results
-
-
-def compute_source(E: field_t,
-                   wavenumber: complex,
-                   omega: complex,
-                   dxes: dx_lists_t,
-                   axis: int,
-                   polarity: int,
-                   slices: List[slice],
-                   epsilon: field_t,
-                   mu: field_t = None,
-                   ) -> field_t:
-    """
-    Given an eigenmode obtained by solve_waveguide_mode, returns the current source distribution
-    necessary to position a unidirectional source at the slice location.
-
-    :param E: E-field of the mode
-    :param wavenumber: Wavenumber of the mode
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :param axis: Propagation axis (0=x, 1=y, 2=z)
-    :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
-    :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
-        as the waveguide cross-section. slices[axis] should select only one
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :return: J distribution for the unidirectional source
-    """
-    E_expanded = expand_wgmode_e(E=E, dxes=dxes, wavenumber=wavenumber, axis=axis,
-                                 polarity=polarity, slices=slices)
-
-    smask = [slice(None)] * 4
-    if polarity > 0:
-        smask[axis + 1] = slice(slices[axis].start, None)
-    else:
-        smask[axis + 1] = slice(None, slices[axis].stop)
-
-    mask = numpy.zeros_like(E_expanded, dtype=int)
-    mask[tuple(smask)] = 1
-
-    masked_e2j = operators.e_boundary_source(mask=vec(mask), omega=omega, dxes=dxes, epsilon=vec(epsilon), mu=vec(mu))
-    J = unvec(masked_e2j @ vec(E_expanded), E.shape[1:])
-    return J
-
-
-def compute_overlap_e(E: field_t,
-                      wavenumber: complex,
-                      dxes: dx_lists_t,
-                      axis: int,
-                      polarity: int,
-                      slices: List[slice],
-                      ) -> field_t:                 # TODO DOCS
-    """
-    Given an eigenmode obtained by solve_waveguide_mode, calculates overlap_e for the
-    mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn)
-    [assumes reflection symmetry].i
-
-    overlap_e makes use of the e2h operator to collapse the above expression into
-     (vec(E) @ vec(overlap_e)), allowing for simple calculation of the mode overlap.
-
-    :param E: E-field of the mode
-    :param H: H-field of the mode (advanced by half of a Yee cell from E)
-    :param wavenumber: Wavenumber of the mode
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :param axis: Propagation axis (0=x, 1=y, 2=z)
-    :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
-    :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
-        as the waveguide cross-section. slices[axis] should select only one
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :return: overlap_e for calculating the mode overlap
-    """
-    slices = tuple(slices)
-
-    Ee = expand_wgmode_e(E=E, wavenumber=wavenumber, dxes=dxes,
-                         axis=axis, polarity=polarity, slices=slices)
-
-    start, stop = sorted((slices[axis].start, slices[axis].start - 2 * polarity))
-
-    slices2 = list(slices)
-    slices2[axis] = slice(start, stop)
-    slices2 = (slice(None), *slices2)
-
-    Etgt = numpy.zeros_like(Ee)
-    Etgt[slices2] = Ee[slices2]
-
-    Etgt /= (Etgt.conj() * Etgt).sum()
-    return Etgt
-
-
-def solve_waveguide_mode_cylindrical(mode_number: int,
-                                     omega: complex,
-                                     dxes: dx_lists_t,
-                                     epsilon: vfield_t,
-                                     r0: float,
-                                     ) -> Dict[str, complex or field_t]:
-    """
-    TODO: fixup
-    Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
-     of the bent waveguide with the specified mode number.
-
-    :param mode_number: Number of the mode, 0-indexed
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types.
-        The first coordinate is assumed to be r, the second is y.
-    :param epsilon: Dielectric constant
-    :param r0: Radius of curvature for the simulation. This should be the minimum value of
-        r within the simulation domain.
-    :return: {'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}
-    """
-
-    '''
-    Solve for the largest-magnitude eigenvalue of the real operator
-    '''
-    dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
-
-    A_r = waveguide.cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0)
-    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + 3)
-    e_xy = eigvecs[:, -(mode_number+1)]
-
-    '''
-    Now solve for the eigenvector of the full operator, using the real operator's
-     eigenvector as an initial guess for Rayleigh quotient iteration.
-    '''
-    A = waveguide.cylindrical_operator(omega, dxes, epsilon, r0)
-    eigval, e_xy = rayleigh_quotient_iteration(A, e_xy)
-
-    # Calculate the wave-vector (force the real part to be positive)
-    wavenumber = numpy.sqrt(eigval)
-    wavenumber *= numpy.sign(numpy.real(wavenumber))
-
-    # TODO: Perform correction on wavenumber to account for numerical dispersion.
-
-    shape = [d.size for d in dxes[0]]
-    e_xy = numpy.hstack((e_xy, numpy.zeros(shape[0] * shape[1])))
-    fields = {
-        'wavenumber': wavenumber,
-        'E': unvec(e_xy, shape),
-#        'E': unvec(e, shape),
-#        'H': unvec(h, shape),
-    }
-
-    return fields
-
-
-def expand_wgmode_e(E: field_t,
-                    wavenumber: complex,
-                    dxes: dx_lists_t,
-                    axis: int,
-                    polarity: int,
-                    slices: List[slice],
-                    ) -> field_t:
-    slices = tuple(slices)
-
-    # Determine phase factors for parallel slices
-    a_shape = numpy.roll([1, -1, 1, 1], axis)
-    a_E = numpy.real(dxes[0][axis]).cumsum()
-    r_E = a_E - a_E[slices[axis]]
-    iphi = polarity * -1j * wavenumber
-    phase_E = numpy.exp(iphi * r_E).reshape(a_shape)
-
-    # Expand our slice to the entire grid using the phase factors
-    E_expanded = numpy.zeros_like(E)
-
-    slices_exp = list(slices)
-    slices_exp[axis] = slice(E.shape[axis + 1])
-    slices_exp = (slice(None), *slices_exp)
-
-    slices_in = (slice(None), *slices)
-
-    E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in]
-    return E_expanded

meanas/fdmath/functional.py

@@ -0,0 +1,109 @@
+"""
+Math functions for finite difference simulations
+
+Basic discrete calculus etc.
+"""
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from .. import field_t, field_updater
+
+
+def deriv_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
+    """
+    Utility operators for taking discretized derivatives (backward variant).
+
+    Args:
+        dx_e: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        List of functions for taking forward derivatives along each axis.
+    """
+    if dx_e:
+        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=2) - f) / dx_e[2][None, None, :]]
+    else:
+        derivs = [lambda f: numpy.roll(f, -1, axis=0) - f,
+                  lambda f: numpy.roll(f, -1, axis=1) - f,
+                  lambda f: numpy.roll(f, -1, axis=2) - f]
+    return derivs
+
+
+def deriv_back(dx_h: List[numpy.ndarray] = None) -> field_updater:
+    """
+    Utility operators for taking discretized derivatives (forward variant).
+
+    Args:
+        dx_h: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        List of functions for taking forward derivatives along each axis.
+    """
+    if dx_h:
+        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=2)) / dx_h[2][None, None, :]]
+    else:
+        derivs = [lambda f: f - numpy.roll(f, 1, axis=0),
+                  lambda f: f - numpy.roll(f, 1, axis=1),
+                  lambda f: f - numpy.roll(f, 1, axis=2)]
+    return derivs
+
+
+def curl_forward(dx_e: List[numpy.ndarray] = None) -> field_updater:
+    """
+    Curl operator for use with the E field.
+
+    Args:
+        dx_e: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        Function `f` for taking the discrete forward curl of a field,
+        `f(E)` -> curlE \\( = \\nabla_f \\times E \\)
+    """
+    Dx, Dy, Dz = deriv_forward(dx_e)
+
+    def ce_fun(e: field_t) -> field_t:
+        output = numpy.empty_like(e)
+        output[0] = Dy(e[2])
+        output[1] = Dz(e[0])
+        output[2] = Dx(e[1])
+        output[0] -= Dz(e[1])
+        output[1] -= Dx(e[2])
+        output[2] -= Dy(e[0])
+        return output
+
+    return ce_fun
+
+
+def curl_back(dx_h: List[numpy.ndarray] = None) -> field_updater:
+    """
+    Create a function which takes the backward curl of a field.
+
+    Args:
+        dx_h: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        Function `f` for taking the discrete backward curl of a field,
+        `f(H)` -> curlH \\( = \\nabla_b \\times H \\)
+    """
+    Dx, Dy, Dz = deriv_back(dx_h)
+
+    def ch_fun(h: field_t) -> field_t:
+        output = numpy.empty_like(h)
+        output[0] = Dy(h[2])
+        output[1] = Dz(h[0])
+        output[2] = Dx(h[1])
+        output[0] -= Dz(h[1])
+        output[1] -= Dx(h[2])
+        output[2] -= Dy(h[0])
+        return output
+
+    return ch_fun
+
+

meanas/fdmath/operators.py

@@ -0,0 +1,231 @@
+"""
+Matrix operators for finite difference simulations
+
+Basic discrete calculus etc.
+"""
+from typing import List, Callable, Tuple, Dict
+import numpy
+import scipy.sparse as sparse
+
+from .. import field_t, vfield_t
+
+
+def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix:
+    """
+    Utility operator for performing a circular shift along a specified axis by a
+     specified number of elements.
+
+    Args:
+        axis: Axis to shift along. x=0, y=1, z=2
+        shape: Shape of the grid being shifted
+        shift_distance: Number of cells to shift by. May be negative. Default 1.
+
+    Returns:
+        Sparse matrix for performing the circular shift.
+    """
+    if len(shape) not in (2, 3):
+        raise Exception('Invalid shape: {}'.format(shape))
+    if axis not in range(len(shape)):
+        raise Exception('Invalid direction: {}, shape is {}'.format(axis, 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)]
+    ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
+
+    n = numpy.prod(shape)
+    i_ind = numpy.arange(n)
+    j_ind = numpy.ravel_multi_index(ijk, shape, order='C')
+
+    vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C')))
+
+    d = sparse.csr_matrix(vij, shape=(n, n))
+
+    if shift_distance < 0:
+        d = d.T
+
+    return d
+
+
+def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix:
+    """
+    Utility operator for performing an n-element shift along a specified axis, with mirror
+    boundary conditions applied to the cells beyond the receding edge.
+
+    Args:
+        axis: Axis to shift along. x=0, y=1, z=2
+        shape: Shape of the grid being shifted
+        shift_distance: Number of cells to shift by. May be negative. Default 1.
+
+    Returns:
+        Sparse matrix for performing the shift-with-mirror.
+    """
+    if len(shape) not in (2, 3):
+        raise Exception('Invalid shape: {}'.format(shape))
+    if axis not in range(len(shape)):
+        raise Exception('Invalid direction: {}, shape is {}'.format(axis, shape))
+    if shift_distance >= shape[axis]:
+        raise Exception('Shift ({}) is too large for axis {} of size {}'.format(
+                        shift_distance, axis, shape[axis]))
+
+    def mirrored_range(n, s):
+        v = numpy.arange(n) + s
+        v = numpy.where(v >= n, 2 * n - v - 1, v)
+        v = numpy.where(v < 0, - 1 - v, v)
+        return v
+
+    shifts = [shift_distance if a == axis else 0 for a in range(3)]
+    shifted_diags = [mirrored_range(n, s) for n, s in zip(shape, shifts)]
+    ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
+
+    n = numpy.prod(shape)
+    i_ind = numpy.arange(n)
+    j_ind = numpy.ravel_multi_index(ijk, shape, order='C')
+
+    vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C')))
+
+    d = sparse.csr_matrix(vij, shape=(n, n))
+    return d
+
+
+def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]:
+    """
+    Utility operators for taking discretized derivatives (forward variant).
+
+    Args:
+        dx_e: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        List of operators for taking forward derivatives along each axis.
+    """
+    shape = [s.size for s in dx_e]
+    n = numpy.prod(shape)
+
+    dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij')
+
+    def deriv(axis):
+        return rotation(axis, shape, 1) - sparse.eye(n)
+
+    Ds = [sparse.diags(+1 / dx.ravel(order='C')) @ deriv(a)
+          for a, dx in enumerate(dx_e_expanded)]
+
+    return Ds
+
+
+def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]:
+    """
+    Utility operators for taking discretized derivatives (backward variant).
+
+    Args:
+        dx_h: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        List of operators for taking forward derivatives along each axis.
+    """
+    shape = [s.size for s in dx_h]
+    n = numpy.prod(shape)
+
+    dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij')
+
+    def deriv(axis):
+        return rotation(axis, shape, -1) - sparse.eye(n)
+
+    Ds = [sparse.diags(-1 / dx.ravel(order='C')) @ deriv(a)
+          for a, dx in enumerate(dx_h_expanded)]
+
+    return Ds
+
+
+def cross(B: List[sparse.spmatrix]) -> sparse.spmatrix:
+    """
+    Cross product operator
+
+    Args:
+        B: List `[Bx, By, Bz]` of sparse matrices corresponding to the x, y, z
+           portions of the operator on the left side of the cross product.
+
+    Returns:
+        Sparse matrix corresponding to (B x), where x is the cross product.
+    """
+    n = B[0].shape[0]
+    zero = sparse.csr_matrix((n, n))
+    return sparse.bmat([[zero, -B[2], B[1]],
+                        [B[2], zero, -B[0]],
+                        [-B[1], B[0], zero]])
+
+
+def vec_cross(b: vfield_t) -> sparse.spmatrix:
+    """
+    Vector cross product operator
+
+    Args:
+        b: Vector on the left side of the cross product.
+
+    Returns:
+
+        Sparse matrix corresponding to (b x), where x is the cross product.
+
+    """
+    B = [sparse.diags(c) for c in numpy.split(b, 3)]
+    return cross(B)
+
+
+def avg_forward(axis: int, shape: List[int]) -> sparse.spmatrix:
+    """
+    Forward average operator `(x4 = (x4 + x5) / 2)`
+
+    Args:
+        axis: Axis to average along (x=0, y=1, z=2)
+        shape: Shape of the grid to average
+
+    Returns:
+        Sparse matrix for forward average operation.
+    """
+    if len(shape) not in (2, 3):
+        raise Exception('Invalid shape: {}'.format(shape))
+
+    n = numpy.prod(shape)
+    return 0.5 * (sparse.eye(n) + rotation(axis, shape))
+
+
+def avg_back(axis: int, shape: List[int]) -> sparse.spmatrix:
+    """
+    Backward average operator `(x4 = (x4 + x3) / 2)`
+
+    Args:
+        axis: Axis to average along (x=0, y=1, z=2)
+        shape: Shape of the grid to average
+
+    Returns:
+        Sparse matrix for backward average operation.
+    """
+    return avg_forward(axis, shape).T
+
+
+def curl_forward(dx_e: List[numpy.ndarray]) -> sparse.spmatrix:
+    """
+    Curl operator for use with the E field.
+
+    Args:
+        dx_e: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        Sparse matrix for taking the discretized curl of the E-field
+    """
+    return cross(deriv_forward(dx_e))
+
+
+def curl_back(dx_h: List[numpy.ndarray]) -> sparse.spmatrix:
+    """
+    Curl operator for use with the H field.
+
+    Args:
+        dx_h: Lists of cell sizes for all axes
+              `[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
+
+    Returns:
+        Sparse matrix for taking the discretized curl of the H-field
+    """
+    return cross(deriv_back(dx_h))

meanas/fdtd/__init__.py

@@ -1,5 +1,5 @@
 """
-Basic FDTD functionality
+Utilities for running finite-difference time-domain (FDTD) simulations
 """
 
 from .base import maxwell_e, maxwell_h

meanas/fdtd/base.py

@@ -5,70 +5,14 @@ from typing import List, Callable, Tuple, Dict
 import numpy
 
 from .. import dx_lists_t, field_t, field_updater
+from ..fdmath.functional import curl_forward, curl_back
+
 
 __author__ = 'Jan Petykiewicz'
 
 
-def curl_h(dxes: dx_lists_t = None) -> field_updater:
-    """
-    Curl operator for use with the H field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :return: Function for taking the discretized curl of the H-field, F(H) -> curlH
-    """
-    if dxes:
-        dxyz_b = numpy.meshgrid(*dxes[1], indexing='ij')
-
-        def dh(f, ax):
-            return (f - numpy.roll(f, 1, axis=ax)) / dxyz_b[ax]
-    else:
-        def dh(f, ax):
-            return f - numpy.roll(f, 1, axis=ax)
-
-    def ch_fun(h: field_t) -> field_t:
-        output = numpy.empty_like(h)
-        output[0] = dh(h[2], 1)
-        output[1] = dh(h[0], 2)
-        output[2] = dh(h[1], 0)
-        output[0] -= dh(h[1], 2)
-        output[1] -= dh(h[2], 0)
-        output[2] -= dh(h[0], 1)
-        return output
-
-    return ch_fun
-
-
-def curl_e(dxes: dx_lists_t = None) -> field_updater:
-    """
-    Curl operator for use with the E field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
-    :return: Function for taking the discretized curl of the E-field, F(E) -> curlE
-    """
-    if dxes is not None:
-        dxyz_a = numpy.meshgrid(*dxes[0], indexing='ij')
-
-        def de(f, ax):
-            return (numpy.roll(f, -1, axis=ax) - f) / dxyz_a[ax]
-    else:
-        def de(f, ax):
-            return numpy.roll(f, -1, axis=ax) - f
-
-    def ce_fun(e: field_t) -> field_t:
-        output = numpy.empty_like(e)
-        output[0] = de(e[2], 1)
-        output[1] = de(e[0], 2)
-        output[2] = de(e[1], 0)
-        output[0] -= de(e[1], 2)
-        output[1] -= de(e[2], 0)
-        output[2] -= de(e[0], 1)
-        return output
-
-    return ce_fun
-
-
 def maxwell_e(dt: float, dxes: dx_lists_t = None) -> field_updater:
-    curl_h_fun = curl_h(dxes)
+    curl_h_fun = curl_back(dxes[1])
 
     def me_fun(e: field_t, h: field_t, epsilon: field_t):
         e += dt * curl_h_fun(h) / epsilon
@@ -78,7 +22,7 @@ def maxwell_e(dt: float, dxes: dx_lists_t = None) -> field_updater:
 
 
 def maxwell_h(dt: float, dxes: dx_lists_t = None) -> field_updater:
-    curl_e_fun = curl_e(dxes)
+    curl_e_fun = curl_forward(dxes[0])
 
     def mh_fun(e: field_t, h: field_t):
         h -= dt * curl_e_fun(e)

meanas/fdtd/energy.py

@@ -35,6 +35,7 @@ def poynting_divergence(s: field_t = None,
     if s is None:
         s = poynting(e, h, dxes=dxes)
 
+    #TODO use deriv operators
     ds = ((s[0] - numpy.roll(s[0], 1, axis=0)) +
           (s[1] - numpy.roll(s[1], 1, axis=1)) +
           (s[2] - numpy.roll(s[2], 1, axis=2)))

meanas/test/__init__.py

@@ -0,0 +1,3 @@
+"""
+Tests (run with `python3 -m pytest -rxPXs | tee results.txt`)
+"""

meanas/test/test_fdfd_pml.py

@@ -101,7 +101,7 @@ def j_distribution(request, shape, epsilon, dxes, omega, src_polarity):
     slices[dim] = slice(shape[dim + 1] // 2,
                         shape[dim + 1] // 2 + 1)
 
-    j = fdfd.waveguide_mode.compute_source(E=e, wavenumber=wavenumber_corrected, omega=omega, dxes=dxes,
+    j = fdfd.waveguide_3d.compute_source(E=e, wavenumber=wavenumber_corrected, omega=omega, dxes=dxes,
                                          axis=dim, polarity=src_polarity, slices=slices, epsilon=epsilon)
     yield j
 
@@ -145,4 +145,3 @@ def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
         )
 
     return sim
-

meanas/types.py

@@ -6,17 +6,20 @@ from typing import List, Callable
 
 
 # Field types
-field_t = numpy.ndarray        # vector field with shape (3, X, Y, Z) (e.g. [E_x, E_y, E_z])
+field_t = numpy.ndarray
-vfield_t = numpy.ndarray       # linearized vector field (vector of length 3*X*Y*Z)
+"""vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)"""
 
+vfield_t = numpy.ndarray
+"""Linearized vector field (vector of length 3*X*Y*Z)"""
+
+dx_lists_t = List[List[numpy.ndarray]]
 '''
  'dxes' datastructure which contains grid cell width information in the following format:
- [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...], [dz_e_0, ...]],
+ `[[[dx_e_0, dx_e_1, ...], [dy_e_0, ...], [dz_e_0, ...]],
-  [[dx_h_0, dx_h_1, ...], [dy_h_0, ...], [dz_h_0, ...]]]
+   [[dx_h_0, dx_h_1, ...], [dy_h_0, ...], [dz_h_0, ...]]]`
- where dx_e_0 is the x-width of the x=0 cells, as used when calculating dE/dx,
+ where `dx_e_0` is the x-width of the `x=0` cells, as used when calculating dE/dx,
- and dy_h_0 is  the y-width of the y=0 cells, as used when calculating dH/dy, etc.
+ and `dy_h_0` is  the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
 '''
-dx_lists_t = List[List[numpy.ndarray]]
 
 
 field_updater = Callable[[field_t], field_t]

pdoc_templates/config.mako

@@ -0,0 +1,46 @@
+<%!
+    # Template configuration. Copy over in your template directory
+    # (used with --template-dir) and adapt as required.
+    html_lang = 'en'
+    show_inherited_members = False
+    extract_module_toc_into_sidebar = True
+    list_class_variables_in_index = True
+    sort_identifiers = True
+    show_type_annotations = True
+
+    # Show collapsed source code block next to each item.
+    # Disabling this can improve rendering speed of large modules.
+    show_source_code = True
+
+    # If set, format links to objects in online source code repository
+    # according to this template. Supported keywords for interpolation
+    # are: commit, path, start_line, end_line.
+    #git_link_template = 'https://github.com/USER/PROJECT/blob/{commit}/{path}#L{start_line}-L{end_line}'
+    #git_link_template = 'https://gitlab.com/USER/PROJECT/blob/{commit}/{path}#L{start_line}-L{end_line}'
+    #git_link_template = 'https://bitbucket.org/USER/PROJECT/src/{commit}/{path}#lines-{start_line}:{end_line}'
+    #git_link_template = 'https://CGIT_HOSTNAME/PROJECT/tree/{path}?id={commit}#n{start-line}'
+    git_link_template = None
+
+    # A prefix to use for every HTML hyperlink in the generated documentation.
+    # No prefix results in all links being relative.
+    link_prefix = ''
+
+    # Enable syntax highlighting for code/source blocks by including Highlight.js
+    syntax_highlighting = True
+
+    # Set the style keyword such as 'atom-one-light' or 'github-gist'
+    #     Options: https://github.com/highlightjs/highlight.js/tree/master/src/styles
+    #     Demo: https://highlightjs.org/static/demo/
+    hljs_style = 'github'
+
+    # If set, insert Google Analytics tracking code. Value is GA
+    # tracking id (UA-XXXXXX-Y).
+    google_analytics = ''
+
+    # If set, render LaTeX math syntax within \(...\) (inline equations),
+    # or within \[...\] or $$...$$ or `.. math::` (block equations)
+    # as nicely-formatted math formulas using MathJax.
+    # Note: in Python docstrings, either all backslashes need to be escaped (\\)
+    # or you need to use raw r-strings.
+    latex_math = True
+%>

pdoc_templates/css.mako

@@ -0,0 +1,389 @@
+<%!
+    from pdoc.html_helpers import minify_css
+%>
+
+<%def name="mobile()" filter="minify_css">
+  .flex {
+    display: flex !important;
+  }
+
+  body {
+    line-height: 1.5em;
+    background: black;
+    color: #DDD;
+  }
+
+  #content {
+    padding: 20px;
+  }
+
+  #sidebar {
+    padding: 30px;
+    overflow: hidden;
+  }
+
+  .http-server-breadcrumbs {
+    font-size: 130%;
+    margin: 0 0 15px 0;
+  }
+
+  #footer {
+    font-size: .75em;
+    padding: 5px 30px;
+    border-top: 1px solid #ddd;
+    text-align: right;
+  }
+    #footer p {
+      margin: 0 0 0 1em;
+      display: inline-block;
+    }
+    #footer p:last-child {
+      margin-right: 30px;
+    }
+
+  h1, h2, h3, h4, h5 {
+    font-weight: 300;
+  }
+  h1 {
+    font-size: 2.5em;
+    line-height: 1.1em;
+  }
+  h2 {
+    font-size: 1.75em;
+    margin: 1em 0 .50em 0;
+  }
+  h3 {
+    font-size: 1.4em;
+    margin: 25px 0 10px 0;
+  }
+  h4 {
+    margin: 0;
+    font-size: 105%;
+  }
+
+  a {
+    color: #999;
+    text-decoration: none;
+    transition: color .3s ease-in-out;
+  }
+  a:hover {
+    color: #18d;
+  }
+
+  .title code {
+    font-weight: bold;
+  }
+  h2[id^="header-"] {
+    margin-top: 2em;
+  }
+  .ident {
+    color: #7ff;
+  }
+
+  pre code {
+    background: transparent;
+    font-size: .8em;
+    line-height: 1.4em;
+  }
+  code {
+    background: #0d0d0e;
+    padding: 1px 4px;
+    overflow-wrap: break-word;
+  }
+  h1 code { background: transparent }
+
+  pre {
+    background: #111;
+    border: 0;
+    border-top: 1px solid #ccc;
+    border-bottom: 1px solid #ccc;
+    margin: 1em 0;
+    padding: 1ex;
+  }
+
+  #http-server-module-list {
+    display: flex;
+    flex-flow: column;
+  }
+    #http-server-module-list div {
+      display: flex;
+    }
+    #http-server-module-list dt {
+      min-width: 10%;
+    }
+    #http-server-module-list p {
+      margin-top: 0;
+    }
+
+  .toc ul,
+  #index {
+    list-style-type: none;
+    margin: 0;
+    padding: 0;
+  }
+    #index code {
+      background: transparent;
+    }
+    #index h3 {
+      border-bottom: 1px solid #ddd;
+    }
+    #index ul {
+      padding: 0;
+    }
+    #index h4 {
+      font-weight: bold;
+    }
+    #index h4 + ul {
+      margin-bottom:.6em;
+    }
+    /* Make TOC lists have 2+ columns when viewport is wide enough.
+       Assuming ~20-character identifiers and ~30% wide sidebar. */
+    @media (min-width: 200ex) { #index .two-column { column-count: 2 } }
+    @media (min-width: 300ex) { #index .two-column { column-count: 3 } }
+
+  dl {
+    margin-bottom: 2em;
+  }
+    dl dl:last-child {
+      margin-bottom: 4em;
+    }
+  dd {
+    margin: 0 0 1em 3em;
+  }
+    #header-classes + dl > dd {
+      margin-bottom: 3em;
+    }
+    dd dd {
+      margin-left: 2em;
+    }
+    dd p {
+      margin: 10px 0;
+    }
+    .name {
+      background: #111;
+      font-weight: bold;
+      font-size: .85em;
+      padding: 5px 10px;
+      display: inline-block;
+      min-width: 40%;
+    }
+      .name:hover {
+        background: #101010;
+      }
+      .name > span:first-child {
+        white-space: nowrap;
+      }
+      .name.class > span:nth-child(2) {
+        margin-left: .4em;
+      }
+    .inherited {
+      color: #777;
+      border-left: 5px solid #eee;
+      padding-left: 1em;
+    }
+    .inheritance em {
+      font-style: normal;
+      font-weight: bold;
+    }
+
+    /* Docstrings titles, e.g. in numpydoc format */
+    .desc h2 {
+      font-weight: 400;
+      font-size: 1.25em;
+    }
+    .desc h3 {
+      font-size: 1em;
+    }
+    .desc dt code {
+      background: inherit;  /* Don't grey-back parameters */
+    }
+
+    .source summary,
+    .git-link-div {
+      color: #aaa;
+      text-align: right;
+      font-weight: 400;
+      font-size: .8em;
+      text-transform: uppercase;
+    }
+      .source summary > * {
+        white-space: nowrap;
+        cursor: pointer;
+      }
+      .git-link {
+        color: inherit;
+        margin-left: 1em;
+      }
+    .source pre {
+      max-height: 500px;
+      overflow: auto;
+      margin: 0;
+    }
+    .source pre code {
+      font-size: 12px;
+      overflow: visible;
+    }
+  .hlist {
+    list-style: none;
+  }
+    .hlist li {
+      display: inline;
+    }
+    .hlist li:after {
+      content: ',\2002';
+    }
+    .hlist li:last-child:after {
+      content: none;
+    }
+    .hlist .hlist {
+      display: inline;
+      padding-left: 1em;
+    }
+
+  img {
+    max-width: 100%;
+  }
+
+  .admonition {
+    padding: .1em .5em;
+    margin-bottom: 1em;
+  }
+    .admonition-title {
+      font-weight: bold;
+    }
+    .admonition.note,
+    .admonition.info,
+    .admonition.important {
+      background: #610;
+    }
+    .admonition.todo,
+    .admonition.versionadded,
+    .admonition.tip,
+    .admonition.hint {
+      background: #202;
+    }
+    .admonition.warning,
+    .admonition.versionchanged,
+    .admonition.deprecated {
+      background: #02b;
+    }
+    .admonition.error,
+    .admonition.danger,
+    .admonition.caution {
+      background: darkpink;
+    }
+</%def>
+
+<%def name="desktop()" filter="minify_css">
+  @media screen and (min-width: 700px) {
+    #sidebar {
+      width: 30%;
+    }
+    #content {
+      width: 70%;
+      max-width: 100ch;
+      padding: 3em 4em;
+      border-left: 1px solid #ddd;
+    }
+    pre code {
+      font-size: 1em;
+    }
+    .item .name {
+      font-size: 1em;
+    }
+    main {
+      display: flex;
+      flex-direction: row-reverse;
+      justify-content: flex-end;
+    }
+    .toc ul ul,
+    #index ul {
+      padding-left: 1.5em;
+    }
+    .toc > ul > li {
+      margin-top: .5em;
+    }
+  }
+</%def>
+
+<%def name="print()" filter="minify_css">
+@media print {
+  #sidebar h1 {
+    page-break-before: always;
+  }
+  .source {
+    display: none;
+  }
+}
+@media print {
+    * {
+        background: transparent !important;
+        color: #000 !important; /* Black prints faster: h5bp.com/s */
+        box-shadow: none !important;
+        text-shadow: none !important;
+    }
+
+    a[href]:after {
+        content: " (" attr(href) ")";
+        font-size: 90%;
+    }
+    /* Internal, documentation links, recognized by having a title,
+       don't need the URL explicity stated. */
+    a[href][title]:after {
+        content: none;
+    }
+
+    abbr[title]:after {
+        content: " (" attr(title) ")";
+    }
+
+    /*
+     * Don't show links for images, or javascript/internal links
+     */
+
+    .ir a:after,
+    a[href^="javascript:"]:after,
+    a[href^="#"]:after {
+        content: "";
+    }
+
+    pre,
+    blockquote {
+        border: 1px solid #999;
+        page-break-inside: avoid;
+    }
+
+    thead {
+        display: table-header-group; /* h5bp.com/t */
+    }
+
+    tr,
+    img {
+        page-break-inside: avoid;
+    }
+
+    img {
+        max-width: 100% !important;
+    }
+
+    @page {
+        margin: 0.5cm;
+    }
+
+    p,
+    h2,
+    h3 {
+        orphans: 3;
+        widows: 3;
+    }
+
+    h1,
+    h2,
+    h3,
+    h4,
+    h5,
+    h6 {
+        page-break-after: avoid;
+    }
+}
+</%def>

pdoc_templates/html.mako

@@ -0,0 +1,421 @@
+<%
+  import os
+
+  import pdoc
+  from pdoc.html_helpers import extract_toc, glimpse, to_html as _to_html, format_git_link
+
+
+  def link(d, name=None, fmt='{}'):
+    name = fmt.format(name or d.qualname + ('()' if isinstance(d, pdoc.Function) else ''))
+    if not isinstance(d, pdoc.Doc) or isinstance(d, pdoc.External) and not external_links:
+        return name
+    url = d.url(relative_to=module, link_prefix=link_prefix,
+                top_ancestor=not show_inherited_members)
+    return '<a title="{}" href="{}">{}</a>'.format(d.refname, url, name)
+
+
+  def to_html(text):
+    return _to_html(text, module=module, link=link, latex_math=latex_math)
+%>
+
+<%def name="ident(name)"><span class="ident">${name}</span></%def>
+
+<%def name="show_source(d)">
+  % if (show_source_code or git_link_template) and d.source and d.obj is not getattr(d.inherits, 'obj', None):
+    <% git_link = format_git_link(git_link_template, d) %>
+    % if show_source_code:
+      <details class="source">
+        <summary>
+            <span>Expand source code</span>
+            % if git_link:
+              <a href="${git_link}" class="git-link">Browse git</a>
+            %endif
+        </summary>
+        <pre><code class="python">${d.source | h}</code></pre>
+      </details>
+    % elif git_link:
+      <div class="git-link-div"><a href="${git_link}" class="git-link">Browse git</a></div>
+    %endif
+  %endif
+</%def>
+
+<%def name="show_desc(d, short=False)">
+  <%
+  inherits = ' inherited' if d.inherits else ''
+  docstring = glimpse(d.docstring) if short or inherits else d.docstring
+  %>
+  % if d.inherits:
+      <p class="inheritance">
+          <em>Inherited from:</em>
+          % if hasattr(d.inherits, 'cls'):
+              <code>${link(d.inherits.cls)}</code>.<code>${link(d.inherits, d.name)}</code>
+          % else:
+              <code>${link(d.inherits)}</code>
+          % endif
+      </p>
+  % endif
+  <section class="desc${inherits}">${docstring | to_html}</section>
+  % if not isinstance(d, pdoc.Module):
+  ${show_source(d)}
+  % endif
+</%def>
+
+<%def name="show_module_list(modules)">
+<h1>Python module list</h1>
+
+% if not modules:
+  <p>No modules found.</p>
+% else:
+  <dl id="http-server-module-list">
+  % for name, desc in modules:
+      <div class="flex">
+      <dt><a href="${link_prefix}${name}">${name}</a></dt>
+      <dd>${desc | glimpse, to_html}</dd>
+      </div>
+  % endfor
+  </dl>
+% endif
+</%def>
+
+<%def name="show_column_list(items)">
+  <%
+      two_column = len(items) >= 6 and all(len(i.name) < 20 for i in items)
+  %>
+  <ul class="${'two-column' if two_column else ''}">
+  % for item in items:
+    <li><code>${link(item, item.name)}</code></li>
+  % endfor
+  </ul>
+</%def>
+
+<%def name="show_module(module)">
+  <%
+  variables = module.variables(sort=sort_identifiers)
+  classes = module.classes(sort=sort_identifiers)
+  functions = module.functions(sort=sort_identifiers)
+  submodules = module.submodules()
+  %>
+
+  <%def name="show_func(f)">
+    <dt id="${f.refname}"><code class="name flex">
+        <%
+            params = ', '.join(f.params(annotate=show_type_annotations, link=link))
+            returns = show_type_annotations and f.return_annotation(link=link) or ''
+            if returns:
+                returns = ' ->\N{NBSP}' + returns
+        %>
+        <span>${f.funcdef()} ${ident(f.name)}</span>(<span>${params})${returns}</span>
+    </code></dt>
+    <dd>${show_desc(f)}</dd>
+  </%def>
+
+  <header>
+  % if http_server:
+    <nav class="http-server-breadcrumbs">
+      <a href="/">All packages</a>
+      <% parts = module.name.split('.')[:-1] %>
+      % for i, m in enumerate(parts):
+        <% parent = '.'.join(parts[:i+1]) %>
+        :: <a href="/${parent.replace('.', '/')}/">${parent}</a>
+      % endfor
+    </nav>
+  % endif
+  <h1 class="title">${'Namespace' if module.is_namespace else 'Module'} <code>${module.name}</code></h1>
+  </header>
+
+  <section id="section-intro">
+  ${module.docstring | to_html}
+  ${show_source(module)}
+  </section>
+
+  <section>
+    % if submodules:
+    <h2 class="section-title" id="header-submodules">Sub-modules</h2>
+    <dl>
+    % for m in submodules:
+      <dt><code class="name">${link(m)}</code></dt>
+      <dd>${show_desc(m, short=True)}</dd>
+    % endfor
+    </dl>
+    % endif
+  </section>
+
+  <section>
+    % if variables:
+    <h2 class="section-title" id="header-variables">Global variables</h2>
+    <dl>
+    % for v in variables:
+      <dt id="${v.refname}"><code class="name">var ${ident(v.name)}</code></dt>
+      <dd>${show_desc(v)}</dd>
+    % endfor
+    </dl>
+    % endif
+  </section>
+
+  <section>
+    % if functions:
+    <h2 class="section-title" id="header-functions">Functions</h2>
+    <dl>
+    % for f in functions:
+      ${show_func(f)}
+    % endfor
+    </dl>
+    % endif
+  </section>
+
+  <section>
+    % if classes:
+    <h2 class="section-title" id="header-classes">Classes</h2>
+    <dl>
+    % for c in classes:
+      <%
+      class_vars = c.class_variables(show_inherited_members, sort=sort_identifiers)
+      smethods = c.functions(show_inherited_members, sort=sort_identifiers)
+      inst_vars = c.instance_variables(show_inherited_members, sort=sort_identifiers)
+      methods = c.methods(show_inherited_members, sort=sort_identifiers)
+      mro = c.mro()
+      subclasses = c.subclasses()
+      params = ', '.join(c.params(annotate=show_type_annotations, link=link))
+      %>
+      <dt id="${c.refname}"><code class="flex name class">
+          <span>class ${ident(c.name)}</span>
+          % if params:
+              <span>(</span><span>${params})</span>
+          % endif
+      </code></dt>
+
+      <dd>${show_desc(c)}
+
+      % if mro:
+          <h3>Ancestors</h3>
+          <ul class="hlist">
+          % for cls in mro:
+              <li>${link(cls)}</li>
+          % endfor
+          </ul>
+      %endif
+
+      % if subclasses:
+          <h3>Subclasses</h3>
+          <ul class="hlist">
+          % for sub in subclasses:
+              <li>${link(sub)}</li>
+          % endfor
+          </ul>
+      % endif
+      % if class_vars:
+          <h3>Class variables</h3>
+          <dl>
+          % for v in class_vars:
+              <dt id="${v.refname}"><code class="name">var ${ident(v.name)}</code></dt>
+              <dd>${show_desc(v)}</dd>
+          % endfor
+          </dl>
+      % endif
+      % if smethods:
+          <h3>Static methods</h3>
+          <dl>
+          % for f in smethods:
+              ${show_func(f)}
+          % endfor
+          </dl>
+      % endif
+      % if inst_vars:
+          <h3>Instance variables</h3>
+          <dl>
+          % for v in inst_vars:
+              <dt id="${v.refname}"><code class="name">var ${ident(v.name)}</code></dt>
+              <dd>${show_desc(v)}</dd>
+          % endfor
+          </dl>
+      % endif
+      % if methods:
+          <h3>Methods</h3>
+          <dl>
+          % for f in methods:
+              ${show_func(f)}
+          % endfor
+          </dl>
+      % endif
+
+      % if not show_inherited_members:
+          <%
+              members = c.inherited_members()
+          %>
+          % if members:
+              <h3>Inherited members</h3>
+              <ul class="hlist">
+              % for cls, mems in members:
+                  <li><code><b>${link(cls)}</b></code>:
+                      <ul class="hlist">
+                          % for m in mems:
+                              <li><code>${link(m, name=m.name)}</code></li>
+                          % endfor
+                      </ul>
+
+                  </li>
+              % endfor
+              </ul>
+          % endif
+      % endif
+
+      </dd>
+    % endfor
+    </dl>
+    % endif
+  </section>
+</%def>
+
+<%def name="module_index(module)">
+  <%
+  variables = module.variables(sort=sort_identifiers)
+  classes = module.classes(sort=sort_identifiers)
+  functions = module.functions(sort=sort_identifiers)
+  submodules = module.submodules()
+  supermodule = module.supermodule
+  %>
+  <nav id="sidebar">
+
+    <%include file="logo.mako"/>
+
+    <h1>Index</h1>
+    ${extract_toc(module.docstring) if extract_module_toc_into_sidebar else ''}
+    <ul id="index">
+    % if supermodule:
+    <li><h3>Super-module</h3>
+      <ul>
+        <li><code>${link(supermodule)}</code></li>
+      </ul>
+    </li>
+    % endif
+
+    % if submodules:
+    <li><h3><a href="#header-submodules">Sub-modules</a></h3>
+      <ul>
+      % for m in submodules:
+        <li><code>${link(m)}</code></li>
+      % endfor
+      </ul>
+    </li>
+    % endif
+
+    % if variables:
+    <li><h3><a href="#header-variables">Global variables</a></h3>
+      ${show_column_list(variables)}
+    </li>
+    % endif
+
+    % if functions:
+    <li><h3><a href="#header-functions">Functions</a></h3>
+      ${show_column_list(functions)}
+    </li>
+    % endif
+
+    % if classes:
+    <li><h3><a href="#header-classes">Classes</a></h3>
+      <ul>
+      % for c in classes:
+        <li>
+        <h4><code>${link(c)}</code></h4>
+        <%
+            members = c.functions(sort=sort_identifiers) + c.methods(sort=sort_identifiers)
+            if list_class_variables_in_index:
+                members += (c.instance_variables(sort=sort_identifiers) +
+                            c.class_variables(sort=sort_identifiers))
+            if not show_inherited_members:
+                members = [i for i in members if not i.inherits]
+            if sort_identifiers:
+              members = sorted(members)
+        %>
+        % if members:
+          ${show_column_list(members)}
+        % endif
+        </li>
+      % endfor
+      </ul>
+    </li>
+    % endif
+
+    </ul>
+  </nav>
+</%def>
+
+<!doctype html>
+<html lang="${html_lang}">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1" />
+  <meta name="generator" content="pdoc ${pdoc.__version__}" />
+
+<%
+    module_list = 'modules' in context.keys()  # Whether we're showing module list in server mode
+%>
+
+  % if module_list:
+    <title>Python module list</title>
+    <meta name="description" content="A list of documented Python modules." />
+  % else:
+    <title>${module.name} API documentation</title>
+    <meta name="description" content="${module.docstring | glimpse, trim, h}" />
+  % endif
+
+    <link href='https://mpxd.net/scripts/normalize.css/normalize.css' rel='stylesheet'>
+    <link href='https://mpxd.net/scripts/sanitize.css/sanitize.css' rel='stylesheet'>
+  % if syntax_highlighting:
+    <link href="https://mpxd.net/scripts/highlightjs/styles/${hljs_style}.min.css" rel="stylesheet">
+  %endif
+
+  <%namespace name="css" file="css.mako" />
+  <style>${css.mobile()}</style>
+  <style media="screen and (min-width: 700px)">${css.desktop()}</style>
+  <style media="print">${css.print()}</style>
+
+  % if google_analytics:
+    <script>
+    window.ga=window.ga||function(){(ga.q=ga.q||[]).push(arguments)};ga.l=+new Date;
+    ga('create', '${google_analytics}', 'auto'); ga('send', 'pageview');
+    </script><script async src='https://www.google-analytics.com/analytics.js'></script>
+  % endif
+
+  % if latex_math:
+    <script async src='https://mpxd.net/scripts/MathJax/MathJax.js?config=TeX-AMS_CHTML'></script>
+  % endif
+
+  <%include file="head.mako"/>
+</head>
+<body>
+<main>
+  % if module_list:
+    <article id="content">
+      ${show_module_list(modules)}
+    </article>
+  % else:
+    <article id="content">
+      ${show_module(module)}
+    </article>
+    ${module_index(module)}
+  % endif
+</main>
+
+<footer id="footer">
+    <%include file="credits.mako"/>
+    <p>Generated by <a href="https://pdoc3.github.io/pdoc"><cite>pdoc</cite> ${pdoc.__version__}</a>.</p>
+</footer>
+
+% if syntax_highlighting:
+    <script src="https://mpxd.net/scripts/highlightjs/highlight.pack.js"></script>
+    <script>hljs.initHighlightingOnLoad()</script>
+% endif
+
+% if http_server and module:  ## Auto-reload on file change in dev mode
+    <script>
+    setInterval(() =>
+        fetch(window.location.href, {
+            method: "HEAD",
+            cache: "no-store",
+            headers: {"If-None-Match": "${os.stat(module.obj.__file__).st_mtime}"},
+        }).then(response => response.ok && window.location.reload()), 700);
+    </script>
+% endif
+</body>
+</html>