better axis label

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

.gitignore

@@ -52,11 +52,17 @@ coverage.xml
 # Django stuff:
 *.log
 
-# Sphinx documentation
-docs/_build/
+# documentation
+doc/
 
 # PyBuilder
 target/
 
 
 .idea/
+
+
+.*.sw[op]
+
+*.svg
+*.html

MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include LICENSE.md
+include meanas/VERSION

README.md

@@ -1,46 +1,111 @@
-# fdfd_tools
+# meanas
 
-**fdfd_tools** is a python package containing utilities for
-creating and analyzing 2D and 3D finite-difference frequency-domain (FDFD)
-electromagnetic simulations.
+**meanas** is a python package for electromagnetic simulations
+
+** UNSTABLE / WORK IN PROGRESS **
+
+Formerly known as [fdfd_tools](https://mpxd.net/code/jan/fdfd_tools).
+
+This package is intended for building simulation inputs, analyzing
+simulation outputs, and running short simulations on unspecialized hardware.
+It is designed to provide tooling and a baseline for other, high-performance
+purpose- and hardware-specific solvers.
 
 
 **Contents**
-* Library of sparse matrices for representing the electromagnetic wave
- equation in 3D, as well as auxiliary matrices for conversion between fields
-* Waveguide mode solver and waveguide mode operators
-* Stretched-coordinate PML boundaries (SCPML)
-* Functional versions of most operators
-* Anisotropic media (eps_xx, eps_yy, eps_zz, mu_xx, ...)
-* Arbitrary distributions of perfect electric and magnetic conductors (PEC / PMC)
+
+- Finite difference frequency domain (FDFD)
+    * Library of sparse matrices for representing the electromagnetic wave
+    equation in 3D, as well as auxiliary matrices for conversion between fields
+    * Waveguide mode operators
+    * Waveguide mode eigensolver
+    * Stretched-coordinate PML boundaries (SCPML)
+    * Functional versions of most operators
+    * Anisotropic media (limited to diagonal elements eps_xx, eps_yy, eps_zz, mu_xx, ...)
+    * Arbitrary distributions of perfect electric and magnetic conductors (PEC / PMC)
+- Finite difference time domain (FDTD)
+    * Basic Maxwell time-steps
+    * Poynting vector and energy calculation
+    * Convolutional PMLs
 
 This package does *not* provide a fast matrix solver, though by default
-```fdfd_tools.solvers.generic(...)``` will call
-```scipy.sparse.linalg.qmr(...)``` to perform a solve.
-For 2D problems this should be fine; likewise,  the waveguide mode
+`meanas.fdfd.solvers.generic(...)` will call
+`scipy.sparse.linalg.qmr(...)` to perform a solve.
+For 2D FDFD problems this should be fine; likewise, the waveguide mode
 solver uses scipy's eigenvalue solver, with reasonable results.
 
-For solving large (or 3D) problems, I recommend a GPU-based iterative
-solver, such as [opencl_fdfd](https://mpxd.net/gogs/jan/opencl_fdfd) or
-those included in [MAGMA](http://icl.cs.utk.edu/magma/index.html)). Your
+For solving large (or 3D) FDFD problems, I recommend a GPU-based iterative
+solver, such as [opencl_fdfd](https://mpxd.net/code/jan/opencl_fdfd) or
+those included in [MAGMA](http://icl.cs.utk.edu/magma/index.html). Your
 solver will need the ability to solve complex symmetric (non-Hermitian)
 linear systems, ideally with double precision.
 
+- [WIP Source repository](https://mpxd.net/code/jan/fdfd_tools/src/branch/wip)
+- *TODO* [Source repository](https://mpxd.net/code/jan/meanas)
+- PyPI *TBD*
+
+
 ## Installation
 
 **Requirements:**
-* python 3 (written and tested with 3.5)
+
+* python 3 (tests require 3.7)
 * numpy
 * scipy
 
 
-Install with pip, via git:
+Install from PyPI with pip:
 ```bash
-pip install git+https://mpxd.net/gogs/jan/fdfd_tools.git@release
+pip3 install 'meanas[test,examples]'
 ```
 
+### Development install
+Install python3.7, virtualenv, and git:
+```bash
+# This is for Debian/Ubuntu/other-apt-based systems; you may need an alternative command
+sudo apt install python3.7 virtualenv build-essential python3.7-dev git
+```
+
+If python 3.7 is not your default python3 version, create a virtualenv:
+```bash
+# Check python3 version:
+python3 --version
+# output on my system: Python 3.7.5rc1
+# If this indicates a version >= 3.7, you can skip all
+#  the steps involving virtualenv or referencing the venv/ directory
+
+# Create a virtual environment using python3.7 and place it in the directory `venv/`
+virtualenv -p python3.7 venv
+```
+
+In-place development install:
+```bash
+# Download using git
+git clone --branch wip https://mpxd.net/code/jan/fdfd_tools.git meanas/
+
+# NOTE: In the future this will become
+#git clone https://mpxd.net/code/jan/meanas.git
+
+# If you are using a virtualenv, activate it
+source venv/bin/activate
+
+# Install in-place (-e, editable) from ./meanas, including testing and example dependencies ([test, examples])
+pip3 install --user -e './meanas[test,examples]'
+
+# Run tests
+cd meanas
+python3 -m pytest -rsxX | tee test_results.txt
+```
+
+#### See also:
+- [git book](https://git-scm.com/book/en/v2)
+- [virtualenv documentation](https://virtualenv.pypa.io/en/stable/userguide/)
+- [python language reference](https://docs.python.org/3/reference/index.html)
+- [python standard library](https://docs.python.org/3/library/index.html)
+
+
 ## Use
 
-See examples/test.py for some simple examples; you may need additional
-packages such as [gridlock](https://mpxd.net/gogs/jan/gridlock)
+See `examples/` for some simple examples; you may need additional
+packages such as [gridlock](https://mpxd.net/code/jan/gridlock)
 to run the examples.

examples/bloch.py

@@ -1,12 +1,44 @@
-import numpy, scipy, gridlock, fdfd_tools
-from fdfd_tools import bloch
+import numpy, scipy, gridlock, meanas
+from meanas.fdfd import bloch
 from numpy.linalg import norm
 import logging
+from pathlib import Path
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
+WISDOM_FILEPATH = pathlib.Path.home() / '.local' / 'share' / 'pyfftw' / 'wisdom.pickle'
 
+
+def pyfftw_save_wisdom(path):
+    path = pathlib.Path(path)
+    try:
+        import pyfftw
+        import pickle
+    except ImportError as e:
+        pass
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, 'wb') as f:
+        pickle.dump(wisdom, f)
+
+
+def pyfftw_load_wisdom(path):
+    path = pathlib.Path(path)
+    try:
+        import pyfftw
+        import pickle
+    except ImportError as e:
+        pass
+
+    try:
+        with open(path, 'rb') as f:
+            wisdom = pickle.load(f)
+        pyfftw.import_wisdom(wisdom)
+    except FileNotFoundError as e:
+        pass
+
+logger.info('Drawing grid...')
 dx = 40
 x_period = 400
 y_period = z_period = 2000
@@ -30,11 +62,13 @@ g2.shifts = numpy.zeros((6,3))
 g2.grids = [numpy.zeros(g.shape) for _ in range(6)]
 
 epsilon = [g.grids[0],] * 3
-reciprocal_lattice = numpy.diag(1e6/numpy.array([x_period, y_period, z_period])) #cols are vectors
+reciprocal_lattice = numpy.diag(1000/numpy.array([x_period, y_period, z_period])) #cols are vectors
+
+pyfftw_load_wisdom(WISDOM_FILEPATH)
 
 #print('Finding k at 1550nm')
-#k, f = bloch.find_k(frequency=1/1550,
-#                    tolerance=(1/1550 - 1/1551),
+#k, f = bloch.find_k(frequency=1000/1550,
+#                    tolerance=(1000 * (1/1550 - 1/1551)),
 #                    direction=[1, 0, 0],
 #                    G_matrix=reciprocal_lattice,
 #                    epsilon=epsilon,
@@ -42,15 +76,15 @@ reciprocal_lattice = numpy.diag(1e6/numpy.array([x_period, y_period, z_period]))
 #
 #print("k={}, f={}, 1/f={}, k/f={}".format(k, f, 1/f, norm(reciprocal_lattice @ k) / f ))
 
-print('Finding f at [0.25, 0, 0]')
+logger.info('Finding f at [0.25, 0, 0]')
 for k0x in [.25]:
     k0 = numpy.array([k0x, 0, 0])
 
     kmag = norm(reciprocal_lattice @ k0)
-    tolerance = (1e6/1550) * 1e-4/1.5  # df = f * dn_eff / n
+    tolerance = (1000/1550) * 1e-4/1.5  # df = f * dn_eff / n
     logger.info('tolerance {}'.format(tolerance))
 
-    n, v = bloch.eigsolve(4, k0, G_matrix=reciprocal_lattice, epsilon=epsilon, tolerance=tolerance)
+    n, v = bloch.eigsolve(4, k0, G_matrix=reciprocal_lattice, epsilon=epsilon, tolerance=tolerance**2)
     v2e = bloch.hmn_2_exyz(k0, G_matrix=reciprocal_lattice, epsilon=epsilon)
     v2h = bloch.hmn_2_hxyz(k0, G_matrix=reciprocal_lattice, epsilon=epsilon)
     ki = bloch.generate_kmn(k0, reciprocal_lattice, g.shape)
@@ -66,3 +100,4 @@ for k0x in [.25]:
     n_eff = norm(reciprocal_lattice @ k0) / f
     print('kmag/f = n_eff =  {} \n wl = {}\n'.format(n_eff, 1/f ))
 
+pyfftw_save_wisdom(WISDOM_FILEPATH)

examples/fdfd.py

@@ -2,16 +2,19 @@ import importlib
 import numpy
 from numpy.linalg import norm
 
-from fdfd_tools import vec, unvec, waveguide_mode
-import fdfd_tools
-import fdfd_tools.functional
-import fdfd_tools.grid
-from fdfd_tools.solvers import generic as generic_solver
+import meanas
+from meanas import fdtd
+from meanas.fdmath import vec, unvec
+from meanas.fdfd import waveguide_3d, functional, scpml, operators
+from meanas.fdfd.solvers import generic as generic_solver
 
 import gridlock
 
 from matplotlib import pyplot
 
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
 
 __author__ = 'Jan Petykiewicz'
 
@@ -57,18 +60,24 @@ def test0(solver=generic_solver):
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (0, 1, 2):
         for p in (-1, 1):
-            dxes = fdfd_tools.grid.stretch_with_scpml(dxes, axis=a, polarity=p, omega=omega,
-                                                      thickness=pml_thickness)
+            dxes = meanas.fdfd.scpml.stretch_with_scpml(dxes, axis=a, polarity=p, omega=omega,
+                                                        thickness=pml_thickness)
 
     J = [numpy.zeros_like(grid.grids[0], dtype=complex) for _ in range(3)]
-    J[1][15, grid.shape[1]//2, grid.shape[2]//2] = 1e5
+    J[1][15, grid.shape[1]//2, grid.shape[2]//2] = 1
+
 
     '''
     Solve!
     '''
+    sim_args = {
+        'omega': omega,
+        'dxes': dxes,
+        'epsilon': vec(grid.grids),
+    }
     x = solver(J=vec(J), **sim_args)
 
-    A = fdfd_tools.functional.e_full(omega, dxes, vec(grid.grids)).tocsr()
+    A = operators.e_full(omega, dxes, vec(grid.grids)).tocsr()
     b = -1j * omega * vec(J)
     print('Norm of the residual is ', norm(A @ x - b))
 
@@ -113,25 +122,26 @@ def test1(solver=generic_solver):
     dxes = [grid.dxyz, grid.autoshifted_dxyz()]
     for a in (0, 1, 2):
         for p in (-1, 1):
-            dxes = fdfd_tools.grid.stretch_with_scpml(dxes,omega=omega, axis=a, polarity=p,
-                                                      thickness=pml_thickness)
+            dxes = scpml.stretch_with_scpml(dxes,omega=omega, axis=a, polarity=p,
+                                            thickness=pml_thickness)
 
     half_dims = numpy.array([10, 20, 15]) * dx
     dims = [-half_dims, half_dims]
     dims[1][0] = dims[0][0]
     ind_dims = (grid.pos2ind(dims[0], which_shifts=None).astype(int),
                 grid.pos2ind(dims[1], which_shifts=None).astype(int))
+    src_axis = 0
     wg_args = {
-        'omega': omega,
         'slices': [slice(i, f+1) for i, f in zip(*ind_dims)],
         'dxes': dxes,
-        'axis': 0,
+        'axis': src_axis,
         'polarity': +1,
     }
 
-    wg_results = waveguide_mode.solve_waveguide_mode(mode_number=0, **wg_args, epsilon=grid.grids)
-    J = waveguide_mode.compute_source(**wg_args, **wg_results)
-    H_overlap = waveguide_mode.compute_overlap_e(**wg_args, **wg_results)
+    wg_results = waveguide_3d.solve_mode(mode_number=0, omega=omega, epsilon=grid.grids, **wg_args)
+    J = waveguide_3d.compute_source(E=wg_results['E'], wavenumber=wg_results['wavenumber'],
+                                    omega=omega, epsilon=grid.grids, **wg_args)
+    e_overlap = waveguide_3d.compute_overlap_e(E=wg_results['E'], wavenumber=wg_results['wavenumber'], **wg_args)
 
     pecg = gridlock.Grid(edge_coords, initial=0.0, num_grids=3)
     # pecg.draw_cuboid(center=[700, 0, 0], dimensions=[80, 1e8, 1e8], eps=1)
@@ -141,6 +151,19 @@ def test1(solver=generic_solver):
     # pmcg.draw_cuboid(center=[700, 0, 0], dimensions=[80, 1e8, 1e8], eps=1)
     # pmcg.visualize_isosurface()
 
+    def pcolor(v):
+        vmax = numpy.max(numpy.abs(v))
+        pyplot.pcolor(v, cmap='seismic', vmin=-vmax, vmax=vmax)
+        pyplot.axis('equal')
+        pyplot.colorbar()
+
+    ss = (1, slice(None), J.shape[2]//2+6, slice(None))
+#    pyplot.figure()
+#    pcolor(J3[ss].T.imag)
+#    pyplot.figure()
+#    pcolor((numpy.abs(J3).sum(axis=2).sum(axis=0) > 0).astype(float).T)
+    pyplot.show(block=True)
+
     '''
     Solve!
     '''
@@ -155,7 +178,7 @@ def test1(solver=generic_solver):
     x = solver(J=vec(J), **sim_args)
 
     b = -1j * omega * vec(J)
-    A = fdfd_tools.operators.e_full(**sim_args).tocsr()
+    A = operators.e_full(**sim_args).tocsr()
     print('Norm of the residual is ', norm(A @ x - b))
 
     E = unvec(x, grid.shape)
@@ -163,42 +186,45 @@ def test1(solver=generic_solver):
     '''
     Plot results
     '''
-    def pcolor(v):
-        vmax = numpy.max(numpy.abs(v))
-        pyplot.pcolor(v, cmap='seismic', vmin=-vmax, vmax=vmax)
-        pyplot.axis('equal')
-        pyplot.colorbar()
-
     center = grid.pos2ind([0, 0, 0], None).astype(int)
     pyplot.figure()
     pyplot.subplot(2, 2, 1)
-    pcolor(numpy.real(E[1][center[0], :, :]))
+    pcolor(numpy.real(E[1][center[0], :, :]).T)
     pyplot.subplot(2, 2, 2)
     pyplot.plot(numpy.log10(numpy.abs(E[1][:, center[1], center[2]]) + 1e-10))
+    pyplot.grid(alpha=0.6)
+    pyplot.ylabel('log10 of Ey @ centerline')
     pyplot.subplot(2, 2, 3)
-    pcolor(numpy.real(E[1][:, :, center[2]]))
+    pcolor(numpy.real(E[1][:, :, center[2]]).T)
     pyplot.subplot(2, 2, 4)
 
     def poyntings(E):
-        e = vec(E)
-        h = fdfd_tools.operators.e2h(omega, dxes) @ e
-        cross1 = fdfd_tools.operators.poynting_e_cross(e, dxes) @ h.conj()
-        cross2 = fdfd_tools.operators.poynting_h_cross(h.conj(), dxes) @ e
-        s1 = unvec(0.5 * numpy.real(cross1), grid.shape)
-        s2 = unvec(0.5 * numpy.real(-cross2), grid.shape)
-        return s1, s2
+        H = functional.e2h(omega, dxes)(E)
+        poynting = fdtd.poynting(e=E, h=H.conj(), dxes=dxes)
+        cross1 = operators.poynting_e_cross(vec(E), dxes) @ vec(H).conj()
+        cross2 = operators.poynting_h_cross(vec(H), dxes) @ vec(E).conj() * -1
+        s1 = 0.5 * unvec(numpy.real(cross1), grid.shape)
+        s2 = 0.5 * unvec(numpy.real(cross2), grid.shape)
+        s0 = 0.5 * poynting.real
+#        s2 = poynting.imag
+        return s0, s1, s2
 
-    s1x, s2x = poyntings(E)
-    pyplot.plot(s1x[0].sum(axis=2).sum(axis=1))
-    pyplot.plot(s2x[0].sum(axis=2).sum(axis=1))
+    s0x, s1x, s2x = poyntings(E)
+    pyplot.plot(s0x[0].sum(axis=2).sum(axis=1), label='s0', marker='.')
+    pyplot.plot(s1x[0].sum(axis=2).sum(axis=1), label='s1', marker='.')
+    pyplot.plot(s2x[0].sum(axis=2).sum(axis=1), label='s2', marker='.')
+    pyplot.plot(E[1][:, center[1], center[2]].real.T, label='Ey', marker='x')
+    pyplot.grid(alpha=0.6)
+    pyplot.legend()
     pyplot.show()
 
     q = []
     for i in range(-5, 30):
-        H_rolled = [numpy.roll(h, i, axis=0) for h in H_overlap]
-        q += [numpy.abs(vec(E) @ vec(H_rolled))]
+        e_ovl_rolled = numpy.roll(e_overlap, i, axis=1)
+        q += [numpy.abs(vec(E) @ vec(e_ovl_rolled).conj())]
     pyplot.figure()
-    pyplot.plot(q)
+    pyplot.plot(q, marker='.')
+    pyplot.grid(alpha=0.6)
     pyplot.title('Overlap with mode')
     pyplot.show()
     print('Average overlap with mode:', sum(q)/len(q))
@@ -209,7 +235,8 @@ def module_available(name):
 
 
 if __name__ == '__main__':
-    # test0()
+    #test0()
+#    test1()
 
     if module_available('opencl_fdfd'):
         from opencl_fdfd import cg_solver as opencl_solver

examples/fdtd.py

@@ -10,7 +10,7 @@ import time
 import numpy
 import h5py
 
-from fdfd_tools import fdtd
+from meanas import fdtd
 from masque import Pattern, shapes
 import gridlock
 import pcgen
@@ -20,9 +20,10 @@ def perturbed_l3(a: float, radius: float, **kwargs) -> Pattern:
     """
     Generate a masque.Pattern object containing a perturbed L3 cavity.
 
-    :param a: Lattice constant.
-    :param radius: Hole radius, in units of a (lattice constant).
-    :param kwargs: Keyword arguments:
+    Args:
+        a: Lattice constant.
+        radius: Hole radius, in units of a (lattice constant).
+        **kwargs: Keyword arguments:
         hole_dose, trench_dose, hole_layer, trench_layer: Shape properties for Pattern.
                 Defaults *_dose=1, hole_layer=0, trench_layer=1.
         shifts_a, shifts_r: passed to pcgen.l3_shift; specifies lattice constant (1 -
@@ -30,11 +31,13 @@ def perturbed_l3(a: float, radius: float, **kwargs) -> Pattern:
                 holes adjacent to the defect (same row). Defaults are 0.15 shift for
                 first hole, 0.075 shift for third hole, and no radius change.
         xy_size: [x, y] number of mirror periods in each direction; total size is
-                2 * n + 1 holes in each direction. Default [10, 10].
+                `2 * n + 1` holes in each direction. Default `[10, 10]`.
         perturbed_radius: radius of holes perturbed to form an upwards-driected beam
                 (multiplicative factor). Default 1.1.
         trench width: Width of the undercut trenches. Default 1.2e3.
-    :return: masque.Pattern object containing the L3 design
+
+    Return:
+        `masque.Pattern` object containing the L3 design
     """
 
     default_args = {'hole_dose':    1,

examples/tcyl.py

@@ -0,0 +1,90 @@
+import importlib
+import numpy
+from numpy.linalg import norm
+
+from meanas import vec, unvec
+from meanas.fdfd import waveguide_mode, functional, scpml
+from meanas.fdfd.solvers import generic as generic_solver
+
+import gridlock
+
+from matplotlib import pyplot
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def test1(solver=generic_solver):
+    dx = 20                 # discretization (nm/cell)
+    pml_thickness = 10      # (number of cells)
+
+    wl = 1550               # Excitation wavelength
+    omega = 2 * numpy.pi / wl
+
+    # Device design parameters
+    w = 800
+    th = 220
+    center = [0, 0, 0]
+    r0 = 8e3
+
+    # refractive indices
+    n_wg = numpy.sqrt(12.6)  # ~Si
+    n_air = 1.0              # air
+
+    # Half-dimensions of the simulation grid
+    y_max = 1200
+    z_max = 900
+    xyz_max = numpy.array([800, y_max, z_max]) + (pml_thickness + 2) * dx
+
+    # Coordinates of the edges of the cells.
+    half_edge_coords = [numpy.arange(dx/2, m + dx/2, step=dx) for m in xyz_max]
+    edge_coords = [numpy.hstack((-h[::-1], h)) for h in half_edge_coords]
+    edge_coords[0] = numpy.array([-dx, dx])
+
+    # #### Create the grid and draw the device ####
+    grid = gridlock.Grid(edge_coords, initial=n_air**2, num_grids=3)
+    grid.draw_cuboid(center=center, dimensions=[8e3, w, th], eps=n_wg**2)
+
+    dxes = [grid.dxyz, grid.autoshifted_dxyz()]
+    for a in (1, 2):
+        for p in (-1, 1):
+            dxes = scmpl.stretch_with_scpml(dxes, omega=omega, axis=a, polarity=p,
+                                            thickness=pml_thickness)
+
+    wg_args = {
+        'omega': omega,
+        'dxes': [(d[1], d[2]) for d in dxes],
+        'epsilon': vec(g.transpose([1, 2, 0]) for g in grid.grids),
+        'r0': r0,
+    }
+
+    wg_results = waveguide_mode.solve_waveguide_mode_cylindrical(mode_number=0, **wg_args)
+
+    E = wg_results['E']
+
+    n_eff = wl / (2 * numpy.pi / wg_results['wavenumber'])
+    print('n =', n_eff)
+    print('alpha (um^-1) =', -4 * numpy.pi * numpy.imag(n_eff) / (wl * 1e-3))
+
+    '''
+    Plot results
+    '''
+    def pcolor(v):
+        vmax = numpy.max(numpy.abs(v))
+        pyplot.pcolor(v.T, cmap='seismic', vmin=-vmax, vmax=vmax)
+        pyplot.axis('equal')
+        pyplot.colorbar()
+
+    pyplot.figure()
+    pyplot.subplot(2, 2, 1)
+    pcolor(numpy.real(E[0][:, :]))
+    pyplot.subplot(2, 2, 2)
+    pcolor(numpy.real(E[1][:, :]))
+    pyplot.subplot(2, 2, 3)
+    pcolor(numpy.real(E[2][:, :]))
+    pyplot.subplot(2, 2, 4)
+    pyplot.show()
+
+
+if __name__ == '__main__':
+    test1()

fdfd_tools/__init__.py

@@ -1,25 +0,0 @@
-"""
-Electromagnetic FDFD simulation tools
-
-Tools for 3D and 2D Electromagnetic Finite Difference Frequency Domain (FDFD)
-simulations. These tools handle conversion of fields to/from vector form,
-creation of the wave operator matrix, stretched-coordinate PMLs, PECs and PMCs,
-field conversion operators, waveguide mode operator, and waveguide mode
-solver.
-
-This package only contains a solver for the waveguide mode eigenproblem;
-if you want to solve 3D problems you can use your favorite iterative sparse
-matrix solver (so long as it can handle complex symmetric [non-Hermitian]
-matrices, ideally with double precision).
-
-
-Dependencies:
-- numpy
-- scipy
-
-"""
-
-from .vectorization import vec, unvec, field_t, vfield_t
-from .grid import dx_lists_t
-
-__author__ = 'Jan Petykiewicz'

fdfd_tools/bloch.py

@@ -1,513 +0,0 @@
-'''
-Bloch eigenmode solver/operators
-
-This module contains functions for generating and solving the
- 3D Bloch eigenproblem. The approach is to transform the problem
- into the (spatial) fourier domain, transforming the equation
-   1/mu * curl(1/eps * curl(H)) = (w/c)^2 H
- into
-   conv(1/mu_k, ik x conv(1/eps_k, ik x H_k)) = (w/c)^2 H_k
- where:
-  - the _k subscript denotes a 3D fourier transformed field
-  - each component of H_k corresponds to a plane wave with wavevector k
-  - x is the cross product
-  - conv denotes convolution
-
- Since k and H are orthogonal for each plane wave, we can use each
- k to create an orthogonal basis (k, m, n), with k x m = n, and
- |m| = |n| = 1. The cross products are then simplified with
-
- k @ h = kx hx + ky hy + kz hz = 0 = hk
- h = hk + hm + hn = hm + hn
- k = kk + km + kn = kk = |k|
-
- k x h = (ky hz - kz hy,
-          kz hx - kx hz,
-          kx hy - ky hx)
-       = ((k x h) @ k, (k x h) @ m, (k x h) @ n)_kmn
-       = (0, (m x k) @ h, (n x k) @ h)_kmn         # triple product ordering
-       = (0, kk (-n @ h), kk (m @ h))_kmn          # (m x k) = -|k| n, etc.
-       = |k| (0, -h @ n, h @ m)_kmn
-
- k x h = (km hn - kn hm,
-          kn hk - kk hn,
-          kk hm - km hk)_kmn
-       = (0, -kk hn, kk hm)_kmn
-       = (-kk hn)(mx, my, mz) + (kk hm)(nx, ny, nz)
-       = |k| (hm * (nx, ny, nz) - hn * (mx, my, mz))
-
- where h is shorthand for H_k, (...)_kmn deontes the (k, m, n) basis,
- and e.g. hm is the component of h in the m direction.
-
- We can also simplify conv(X_k, Y_k) as fftn(X * ifftn(Y_k)).
-
- Using these results and storing H_k as h = (hm, hn), we have
-  e_xyz = fftn(1/eps * ifftn(|k| (hm * n - hn * m)))
-  b_mn = |k| (-e_xyz @ n, e_xyz @ m)
-  h_mn = fftn(1/mu * ifftn(b_m * m + b_n * n))
- which forms the operator from the left side of the equation.
-
- We can then use a preconditioned block Rayleigh iteration algorithm, as in
-  SG Johnson and JD Joannopoulos, Block-iterative frequency-domain methods
-  for Maxwell's equations in a planewave basis, Optics Express 8, 3, 173-190 (2001)
- (similar to that used in MPB) to find the eigenvectors for this operator.
-
- ===
-
- Typically you will want to do something like
-
-    recip_lattice = numpy.diag(1/numpy.array(epsilon[0].shape * dx))
-    n, v = bloch.eigsolve(5, k0, recip_lattice, epsilon)
-    f = numpy.sqrt(-numpy.real(n[0]))
-    n_eff = norm(recip_lattice @ k0) / f
-
-    v2e = bloch.hmn_2_exyz(k0, recip_lattice, epsilon)
-    e_field = v2e(v[0])
-
-    k, f = find_k(frequency=1/1550,
-                  tolerance=(1/1550 - 1/1551),
-                  direction=[1, 0, 0],
-                  G_matrix=recip_lattice,
-                  epsilon=epsilon,
-                  band=0)
-
-'''
-
-from typing import List, Tuple, Callable, Dict
-import logging
-import numpy
-from numpy.fft import fftn, ifftn, fftfreq
-import scipy
-import scipy.optimize
-from scipy.linalg import norm
-import scipy.sparse.linalg as spalg
-
-from .eigensolvers import rayleigh_quotient_iteration
-from . import field_t
-
-logger = logging.getLogger(__name__)
-
-
-def generate_kmn(k0: numpy.ndarray,
-                 G_matrix: numpy.ndarray,
-                 shape: numpy.ndarray
-                 ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
-    """
-    Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.
-
-    :param k0: [k0x, k0y, k0z], Bloch wavevector, in G basis.
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param shape: [nx, ny, nz] shape of the simulation grid.
-    :return: (|k|, m, n) where |k| has shape tuple(shape) + (1,)
-            and m, n have shape tuple(shape) + (3,).
-            All are given in the xyz basis (e.g. |k|[0,0,0] = norm(G_matrix @ k0)).
-    """
-    k0 = numpy.array(k0)
-
-    Gi_grids = numpy.meshgrid(*(fftfreq(n, 1/n) for n in shape[:3]), indexing='ij')
-    Gi = numpy.stack(Gi_grids, axis=3)
-
-    k_G = k0[None, None, None, :] - Gi
-    k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2)
-
-    m = numpy.broadcast_to([0, 1, 0], tuple(shape[:3]) + (3,)).astype(float)
-    n = numpy.broadcast_to([0, 0, 1], tuple(shape[:3]) + (3,)).astype(float)
-
-    xy_non0 = numpy.any(k_xyz[:, :, :, 0:1] != 0, axis=3)
-    if numpy.any(xy_non0):
-        u = numpy.cross(k_xyz[xy_non0], [0, 0, 1])
-        m[xy_non0, :] = u / norm(u, axis=1)[:, None]
-
-    z_non0 = numpy.any(k_xyz != 0, axis=3)
-    if numpy.any(z_non0):
-        v = numpy.cross(k_xyz[z_non0], m[z_non0])
-        n[z_non0, :] = v / norm(v, axis=1)[:, None]
-
-    k_mag = norm(k_xyz, axis=3)[:, :, :, None]
-    return k_mag, m, n
-
-
-def maxwell_operator(k0: numpy.ndarray,
-                     G_matrix: numpy.ndarray,
-                     epsilon: field_t,
-                     mu: field_t = None
-                     ) -> Callable[[numpy.ndarray], numpy.ndarray]:
-    """
-    Generate the Maxwell operator
-        conv(1/mu_k, ik x conv(1/eps_k, ik x ___))
-    which is the spatial-frequency-space representation of
-        1/mu * curl(1/eps * curl(___))
-
-    The operator is a function that acts on a vector h_mn of size (2 * epsilon[0].size)
-
-    See the module-level docstring for more information.
-
-    :param k0: Bloch wavevector, [k0x, k0y, k0z].
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param epsilon: Dielectric constant distribution for the simulation.
-        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
-    :param mu: Magnetic permability distribution for the simulation.
-        Default None (1 everywhere).
-    :return: Function which applies the maxwell operator to h_mn.
-    """
-
-    shape = epsilon[0].shape + (1,)
-    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
-
-    epsilon = numpy.stack(epsilon, 3)
-    if mu is not None:
-        mu = numpy.stack(mu, 3)
-
-    def operator(h: numpy.ndarray):
-        """
-        Maxwell operator for Bloch eigenmode simulation.
-
-        h is complex 2-field in (m, n) basis, vectorized
-
-        :param h: Raveled h_mn; size (2 * epsilon[0].size).
-        :return: Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)).
-        """
-        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
-
-        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
-
-        # cross product and transform into xyz basis
-        d_xyz = (n * hin_m -
-                 m * hin_n) * k_mag
-
-        # divide by epsilon
-        e_xyz = fftn(ifftn(d_xyz, axes=range(3)) / epsilon, axes=range(3))
-
-        # cross product and transform into mn basis
-        b_m = numpy.sum(e_xyz * n, axis=3)[:, :, :, None] * -k_mag
-        b_n = numpy.sum(e_xyz * m, axis=3)[:, :, :, None] * +k_mag
-
-        if mu is None:
-            h_m, h_n = b_m, b_n
-        else:
-            # transform from mn to xyz
-            b_xyz = (m * b_m[:, :, :, None] +
-                     n * b_n[:, :, :, None])
-
-            # divide by mu
-            h_xyz = fftn(ifftn(b_xyz, axes=range(3)) / mu, axes=range(3))
-
-            # transform back to mn
-            h_m = numpy.sum(h_xyz * m, axis=3)
-            h_n = numpy.sum(h_xyz * n, axis=3)
-        return numpy.hstack((h_m.ravel(), h_n.ravel()))
-
-    return operator
-
-
-def hmn_2_exyz(k0: numpy.ndarray,
-               G_matrix: numpy.ndarray,
-               epsilon: field_t,
-               ) -> Callable[[numpy.ndarray], field_t]:
-    """
-    Generate an operator which converts a vectorized spatial-frequency-space
-     h_mn into an E-field distribution, i.e.
-        ifft(conv(1/eps_k, ik x h_mn))
-
-    The operator is a function that acts on a vector h_mn of size (2 * epsilon[0].size)
-
-    See the module-level docstring for more information.
-
-    :param k0: Bloch wavevector, [k0x, k0y, k0z].
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param epsilon: Dielectric constant distribution for the simulation.
-        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
-    :return: Function for converting h_mn into E_xyz
-    """
-    shape = epsilon[0].shape + (1,)
-    epsilon = numpy.stack(epsilon, 3)
-
-    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
-
-    def operator(h: numpy.ndarray) -> field_t:
-        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
-        d_xyz = (n * hin_m -
-                 m * hin_n) * k_mag
-
-        # divide by epsilon
-        return [ei for ei in numpy.rollaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3)]
-
-    return operator
-
-
-def hmn_2_hxyz(k0: numpy.ndarray,
-               G_matrix: numpy.ndarray,
-               epsilon: field_t
-               ) -> Callable[[numpy.ndarray], field_t]:
-    """
-    Generate an operator which converts a vectorized spatial-frequency-space
-     h_mn into an H-field distribution, i.e.
-        ifft(h_mn)
-
-    The operator is a function that acts on a vector h_mn of size (2 * epsilon[0].size)
-
-    See the module-level docstring for more information.
-
-    :param k0: Bloch wavevector, [k0x, k0y, k0z].
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param epsilon: Dielectric constant distribution for the simulation.
-        Only epsilon[0].shape is used.
-    :return: Function for converting h_mn into H_xyz
-    """
-    shape = epsilon[0].shape + (1,)
-    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
-
-    def operator(h: numpy.ndarray):
-        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
-        h_xyz = (m * hin_m +
-                 n * hin_n)
-        return [ifftn(hi) for hi in numpy.rollaxis(h_xyz, 3)]
-
-    return operator
-
-
-def inverse_maxwell_operator_approx(k0: numpy.ndarray,
-                                    G_matrix: numpy.ndarray,
-                                    epsilon: field_t,
-                                    mu: field_t = None
-                                    ) -> Callable[[numpy.ndarray], numpy.ndarray]:
-    """
-    Generate an approximate inverse of the Maxwell operator,
-        ik x conv(eps_k, ik x conv(mu_k, ___))
-     which can be used to improve the speed of ARPACK in shift-invert mode.
-
-    See the module-level docstring for more information.
-
-    :param k0: Bloch wavevector, [k0x, k0y, k0z].
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param epsilon: Dielectric constant distribution for the simulation.
-        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
-    :param mu: Magnetic permability distribution for the simulation.
-        Default None (1 everywhere).
-    :return: Function which applies the approximate inverse of the maxwell operator to h_mn.
-    """
-    shape = epsilon[0].shape + (1,)
-    epsilon = numpy.stack(epsilon, 3)
-
-    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
-
-    if mu is not None:
-        mu = numpy.stack(mu, 3)
-
-    def operator(h: numpy.ndarray):
-        """
-        Approximate inverse Maxwell operator for Bloch eigenmode simulation.
-
-        h is complex 2-field in (m, n) basis, vectorized
-
-        :param h: Raveled h_mn; size (2 * epsilon[0].size).
-        :return: Raveled ik x conv(eps_k, ik x conv(mu_k, h_mn))
-        """
-        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
-
-        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
-
-        if mu is None:
-            b_m, b_n = hin_m, hin_n
-        else:
-            # transform from mn to xyz
-            h_xyz = (m * hin_m[:, :, :, None] +
-                     n * hin_n[:, :, :, None])
-
-            # multiply by mu
-            b_xyz = fftn(ifftn(h_xyz, axes=range(3)) * mu, axes=range(3))
-
-            # transform back to mn
-            b_m = numpy.sum(b_xyz * m, axis=3)
-            b_n = numpy.sum(b_xyz * n, axis=3)
-
-        # cross product and transform into xyz basis
-        e_xyz = (n * b_m -
-                 m * b_n) / k_mag
-
-        # multiply by epsilon
-        d_xyz = fftn(ifftn(e_xyz, axes=range(3)) * epsilon, axes=range(3))
-
-        # cross product and transform into mn basis   crossinv_t2c
-        h_m = numpy.sum(e_xyz * n, axis=3)[:, :, :, None] / +k_mag
-        h_n = numpy.sum(e_xyz * m, axis=3)[:, :, :, None] / -k_mag
-
-        return numpy.hstack((h_m.ravel(), h_n.ravel()))
-
-    return operator
-
-
-def eigsolve(num_modes: int,
-             k0: numpy.ndarray,
-             G_matrix: numpy.ndarray,
-             epsilon: field_t,
-             mu: field_t = None,
-             tolerance = 1e-8,
-             ) -> Tuple[numpy.ndarray, numpy.ndarray]:
-    """
-    Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
-     k0 of the specified structure.
-
-    :param k0: Bloch wavevector, [k0x, k0y, k0z].
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param epsilon: Dielectric constant distribution for the simulation.
-        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
-    :param mu: Magnetic permability distribution for the simulation.
-        Default None (1 everywhere).
-    :return: (eigenvalues, eigenvectors) where eigenvalues[i] corresponds to the
-        vector eigenvectors[i, :]
-    """
-    h_size = 2 * epsilon[0].size
-
-    kmag = norm(G_matrix @ k0)
-
-    '''
-    Generate the operators
-    '''
-    mop = maxwell_operator(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
-    imop = inverse_maxwell_operator_approx(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
-
-    scipy_op = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=mop)
-    scipy_iop = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=imop)
-
-    y_shape = (h_size, num_modes)
-
-    def rayleigh_quotient(Z: numpy.ndarray, approx_grad: bool = True):
-        """
-        Absolute value of the block Rayleigh quotient, and the associated gradient.
-
-        See Johnson and Joannopoulos, Opt. Expr. 8, 3 (2001) for details (full
-         citation in module docstring).
-
-        ===
-
-        Notes on my understanding of the procedure:
-
-        Minimize f(Y) = |trace((Y.H @ A @ Y)|, making use of Y = Z @ inv(Z.H @ Z)^(1/2)
-         (a polar orthogonalization of Y). This gives f(Z) = |trace(Z.H @ A @ Z @ U)|,
-         where U = inv(Z.H @ Z). We minimize the absolute value to find the eigenvalues
-         with smallest magnitude.
-
-        The gradient is P @ (A @ Z @ U), where P = (1 - Z @ U @ Z.H) is a projection
-         onto the space orthonormal to Z. If approx_grad is True, the approximate
-         inverse of the maxwell operator is used to precondition the gradient.
-        """
-        z = Z.view(dtype=complex).reshape(y_shape)
-        U = numpy.linalg.inv(z.conj().T @ z)
-        zU = z @ U
-        AzU = scipy_op @ zU
-        zTAzU = z.conj().T @ AzU
-        f = numpy.real(numpy.trace(zTAzU))
-        if approx_grad:
-            df_dy = scipy_iop @ (AzU - zU @ zTAzU)
-        else:
-            df_dy = (AzU - zU @ zTAzU)
-        
-        df_dy_flat = df_dy.view(dtype=float).ravel()
-        return numpy.abs(f), numpy.sign(f) * df_dy_flat
-
-    '''
-    Use the conjugate gradient method and the approximate gradient calculation to
-     quickly find approximate eigenvectors.
-    '''
-    result = scipy.optimize.minimize(rayleigh_quotient,
-                                     numpy.random.rand(*y_shape, 2),
-                                     jac=True,
-                                     method='L-BFGS-B',
-                                     tol=1e-20,
-                                     options={'maxiter': 2000, 'gtol':0, 'ftol':1e-20 , 'disp':True})#, 'maxls':80, 'm':30})
-
-
-    result = scipy.optimize.minimize(lambda y: rayleigh_quotient(y, True),
-                                     result.x,
-                                     jac=True,
-                                     method='L-BFGS-B',
-                                     tol=1e-20,
-                                     options={'maxiter': 2000, 'gtol':0, 'disp':True})
-
-    result = scipy.optimize.minimize(lambda y: rayleigh_quotient(y, False),
-                                     result.x,
-                                     jac=True,
-                                     method='L-BFGS-B',
-                                     tol=1e-20,
-                                     options={'maxiter': 2000, 'gtol':0, 'disp':True})
-
-    for i in range(20):
-        result = scipy.optimize.minimize(lambda y: rayleigh_quotient(y, False),
-                                     result.x,
-                                     jac=True,
-                                     method='L-BFGS-B',
-                                     tol=1e-20,
-                                     options={'maxiter': 70, 'gtol':0, 'disp':True})
-        if result.nit == 0:
-            # We took 0 steps, so re-running won't help
-            break
-
-
-    z = result.x.view(dtype=complex).reshape(y_shape)
-
-    '''
-    Recover eigenvectors from Z
-    '''
-    U = numpy.linalg.inv(z.conj().T @ z)
-    y = z @ scipy.linalg.sqrtm(U)
-    w = y.conj().T @ (scipy_op @ y)
-
-    eigvals, w_eigvecs = numpy.linalg.eig(w)
-    eigvecs = y @ w_eigvecs
-
-    for i in range(len(eigvals)):
-        v = eigvecs[:, i]
-        n = eigvals[i]
-        v /= norm(v)
-        eigness = norm(scipy_op @ v - (v.conj() @ (scipy_op @ v)) * v )
-        f = numpy.sqrt(-numpy.real(n))
-        df = numpy.sqrt(-numpy.real(n + eigness))
-        neff_err = kmag * (1/df - 1/f)
-        logger.info('eigness {}: {}\n neff_err: {}'.format(i, eigness, neff_err))
-
-    order = numpy.argsort(numpy.abs(eigvals))
-    return eigvals[order], eigvecs.T[order]
-
-
-def find_k(frequency: float,
-           tolerance: float,
-           direction: numpy.ndarray,
-           G_matrix: numpy.ndarray,
-           epsilon: field_t,
-           mu: field_t = None,
-           band: int = 0,
-           k_min: float = 0,
-           k_max: float = 0.5,
-           ) -> Tuple[numpy.ndarray, float]:
-    """
-    Search for a bloch vector that has a given frequency.
-
-    :param frequency: Target frequency.
-    :param tolerance: Target frequency tolerance.
-    :param direction: k-vector direction to search along.
-    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
-    :param epsilon: Dielectric constant distribution for the simulation.
-        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
-    :param mu: Magnetic permability distribution for the simulation.
-        Default None (1 everywhere).
-    :param band: Which band to search in. Default 0 (lowest frequency).
-    return: (k, actual_frequency) The found k-vector and its frequency
-    """
-
-    direction = numpy.array(direction) / norm(direction)
-
-    def get_f(k0_mag: float, band: int = 0):
-        k0 = direction * k0_mag
-        n, _v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon)
-        f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
-        return f
-
-    res = scipy.optimize.minimize_scalar(lambda x: abs(get_f(x, band) - frequency),
-                                         (k_min + k_max) / 2,
-                                         method='Bounded',
-                                         bounds=(k_min, k_max),
-                                         options={'xatol': abs(tolerance)})
-    return res.x * direction, res.fun + frequency
-
-

fdfd_tools/fdtd.py

@@ -1,239 +0,0 @@
-from typing import List, Callable, Tuple, Dict
-import numpy
-
-from . import dx_lists_t, field_t
-
-__author__ = 'Jan Petykiewicz'
-
-
-functional_matrix = Callable[[field_t], field_t]
-
-
-def curl_h(dxes: dx_lists_t = None) -> functional_matrix:
-    """
-    Curl operator for use with the H field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :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:
-        e = [dh(h[2], 1) - dh(h[1], 2),
-             dh(h[0], 2) - dh(h[2], 0),
-             dh(h[1], 0) - dh(h[0], 1)]
-        return e
-
-    return ch_fun
-
-
-def curl_e(dxes: dx_lists_t = None) -> functional_matrix:
-    """
-    Curl operator for use with the E field.
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :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:
-        h = [de(e[2], 1) - de(e[1], 2),
-             de(e[0], 2) - de(e[2], 0),
-             de(e[1], 0) - de(e[0], 1)]
-        return h
-
-    return ce_fun
-
-
-def maxwell_e(dt: float, dxes: dx_lists_t = None) -> functional_matrix:
-    curl_h_fun = curl_h(dxes)
-
-    def me_fun(e: field_t, h: field_t, epsilon: field_t):
-        ch = curl_h_fun(h)
-        for ei, ci, epsi in zip(e, ch, epsilon):
-            ei += dt * ci / epsi
-        return e
-
-    return me_fun
-
-
-def maxwell_h(dt: float, dxes: dx_lists_t = None) -> functional_matrix:
-    curl_e_fun = curl_e(dxes)
-
-    def mh_fun(e: field_t, h: field_t):
-        ce = curl_e_fun(e)
-        for hi, ci in zip(h, ce):
-            hi -= dt * ci
-        return h
-
-    return mh_fun
-
-
-def conducting_boundary(direction: int,
-                        polarity: int
-                        ) -> Tuple[functional_matrix, functional_matrix]:
-    dirs = [0, 1, 2]
-    if direction not in dirs:
-        raise Exception('Invalid direction: {}'.format(direction))
-    dirs.remove(direction)
-    u, v = dirs
-
-    if polarity < 0:
-        boundary_slice = [slice(None)] * 3
-        shifted1_slice = [slice(None)] * 3
-        boundary_slice[direction] = 0
-        shifted1_slice[direction] = 1
-
-        def en(e: field_t):
-            e[direction][boundary_slice] = 0
-            e[u][boundary_slice] = e[u][shifted1_slice]
-            e[v][boundary_slice] = e[v][shifted1_slice]
-            return e
-
-        def hn(h: field_t):
-            h[direction][boundary_slice] = h[direction][shifted1_slice]
-            h[u][boundary_slice] = 0
-            h[v][boundary_slice] = 0
-            return h
-
-        return en, hn
-
-    elif polarity > 0:
-        boundary_slice = [slice(None)] * 3
-        shifted1_slice = [slice(None)] * 3
-        shifted2_slice = [slice(None)] * 3
-        boundary_slice[direction] = -1
-        shifted1_slice[direction] = -2
-        shifted2_slice[direction] = -3
-
-        def ep(e: field_t):
-            e[direction][boundary_slice] = -e[direction][shifted2_slice]
-            e[direction][shifted1_slice] = 0
-            e[u][boundary_slice] = e[u][shifted1_slice]
-            e[v][boundary_slice] = e[v][shifted1_slice]
-            return e
-
-        def hp(h: field_t):
-            h[direction][boundary_slice] = h[direction][shifted1_slice]
-            h[u][boundary_slice] = -h[u][shifted2_slice]
-            h[u][shifted1_slice] = 0
-            h[v][boundary_slice] = -h[v][shifted2_slice]
-            h[v][shifted1_slice] = 0
-            return h
-
-        return ep, hp
-
-    else:
-        raise Exception('Bad polarity: {}'.format(polarity))
-
-
-def cpml(direction:int,
-         polarity: int,
-         dt: float,
-         epsilon: field_t,
-         thickness: int = 8,
-         epsilon_eff: float = 1,
-         dtype: numpy.dtype = numpy.float32,
-         ) -> Tuple[Callable, Callable, Dict[str, field_t]]:
-
-    if direction not in range(3):
-        raise Exception('Invalid direction: {}'.format(direction))
-
-    if polarity not in (-1, 1):
-        raise Exception('Invalid polarity: {}'.format(polarity))
-
-    if thickness <= 2:
-        raise Exception('It would be wise to have a pml with 4+ cells of thickness')
-
-    if epsilon_eff <= 0:
-        raise Exception('epsilon_eff must be positive')
-
-    m = (3.5, 1)
-    sigma_max = 0.8 * (m[0] + 1) / numpy.sqrt(epsilon_eff)
-    alpha_max = 0  # TODO: Decide what to do about non-zero alpha
-    transverse = numpy.delete(range(3), direction)
-    u, v = transverse
-
-    xe = numpy.arange(1, thickness+1, dtype=float)
-    xh = numpy.arange(1, thickness+1, dtype=float)
-    if polarity > 0:
-        xe -= 0.5
-    elif polarity < 0:
-        xh -= 0.5
-        xe = xe[::-1]
-        xh = xh[::-1]
-    else:
-        raise Exception('Bad polarity!')
-
-    expand_slice = [None] * 3
-    expand_slice[direction] = slice(None)
-
-    def par(x):
-        sigma = ((x / thickness) ** m[0]) * sigma_max
-        alpha = ((1 - x / thickness) ** m[1]) * alpha_max
-        p0 = numpy.exp(-(sigma + alpha) * dt)
-        p1 = sigma / (sigma + alpha) * (p0 - 1)
-        return p0[expand_slice], p1[expand_slice]
-
-    p0e, p1e = par(xe)
-    p0h, p1h = par(xh)
-
-    region = [slice(None)] * 3
-    if polarity < 0:
-        region[direction] = slice(None, thickness)
-    elif polarity > 0:
-        region[direction] = slice(-thickness, None)
-    else:
-        raise Exception('Bad polarity!')
-
-    if direction == 1:
-        se = 1
-    else:
-        se = -1
-
-    # TODO check if epsilon is uniform?
-    shape = list(epsilon[0].shape)
-    shape[direction] = thickness
-    psi_e = [numpy.zeros(shape, dtype=dtype), numpy.zeros(shape, dtype=dtype)]
-    psi_h = [numpy.zeros(shape, dtype=dtype), numpy.zeros(shape, dtype=dtype)]
-
-    fields = {
-        'psi_e_u': psi_e[0],
-        'psi_e_v': psi_e[1],
-        'psi_h_u': psi_h[0],
-        'psi_h_v': psi_h[1],
-    }
-
-    def pml_e(e: field_t, h: field_t, epsilon: field_t) -> Tuple[field_t, field_t]:
-        psi_e[0] *= p0e
-        psi_e[0] += p1e * (h[v][region] - numpy.roll(h[v], 1, axis=direction)[region])
-        psi_e[1] *= p0e
-        psi_e[1] += p1e * (h[u][region] - numpy.roll(h[u], 1, axis=direction)[region])
-        e[u][region] += se * dt * psi_e[0] / epsilon[u][region]
-        e[v][region] -= se * dt * psi_e[1] / epsilon[v][region]
-        return e, h
-
-    def pml_h(e: field_t, h: field_t) -> Tuple[field_t, field_t]:
-        psi_h[0] *= p0h
-        psi_h[0] += p1h * (numpy.roll(e[v], -1, axis=direction)[region] - e[v][region])
-        psi_h[1] *= p0h
-        psi_h[1] += p1h * (numpy.roll(e[u], -1, axis=direction)[region] - e[u][region])
-        h[u][region] -= se * dt * psi_h[0]
-        h[v][region] += se * dt * psi_h[1]
-        return e, h
-
-    return pml_e, pml_h, fields

fdfd_tools/functional.py

@@ -1,149 +0,0 @@
-"""
-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 inputs in the form E = [E_x, E_y, E_z], where each
- component E_* is an ndarray of equal shape.
-"""
-from typing import List, Callable
-import numpy
-
-from . import dx_lists_t, field_t
-
-__author__ = 'Jan Petykiewicz'
-
-
-functional_matrix = 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 fdfd_tools.operators header
-    :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 = [dh(h[2], 1) - dh(h[1], 2),
-             dh(h[0], 2) - dh(h[2], 0),
-             dh(h[1], 0) - 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 fdfd_tools.operators header
-    :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 = [de(e[2], 1) - de(e[1], 2),
-             de(e[0], 2) - de(e[2], 0),
-             de(e[1], 0) - 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:
-    """
-    Wave operator del x (1/mu * del x) - omega**2 * epsilon, for use with E-field,
-     with wave equation
-    (del x (1/mu * del x) - omega**2 * epsilon) E = -i * omega * J
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param epsilon: Dielectric constant
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :return: Function implementing the wave operator A(E) -> E
-    """
-    ch = curl_h(dxes)
-    ce = curl_e(dxes)
-
-    def op_1(e):
-        curls = ch(ce(e))
-        return [c - omega ** 2 * e * x for c, e, x in zip(curls, epsilon, e)]
-
-    def op_mu(e):
-        curls = ch([m * y for m, y in zip(mu, ce(e))])
-        return [c - omega ** 2 * p * x for c, p, x in zip(curls, epsilon, e)]
-
-    if numpy.any(numpy.equal(mu, None)):
-        return op_1
-    else:
-        return op_mu
-
-
-def eh_full(omega: complex,
-            dxes: dx_lists_t,
-            epsilon: field_t,
-            mu: field_t = None
-            ) -> functional_matrix:
-    """
-    Wave operator for full (both E and H) field representation.
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param epsilon: Dielectric constant
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :return: Function implementing the wave operator A(E, H) -> (E, H)
-    """
-    ch = curl_h(dxes)
-    ce = curl_e(dxes)
-
-    def op_1(e, h):
-        return ([c - 1j * omega * p * x for c, p, x in zip(ch(h), epsilon, e)],
-                [c + 1j * omega * y for c, y in zip(ce(e), h)])
-
-    def op_mu(e, h):
-        return ([c - 1j * omega * p * x for c, p, x in zip(ch(h), epsilon, e)],
-                [c + 1j * omega * m * y for c, m, y in zip(ce(e), mu, h)])
-
-    if numpy.any(numpy.equal(mu, None)):
-        return op_1
-    else:
-        return op_mu
-
-
-def e2h(omega: complex,
-        dxes: dx_lists_t,
-        mu: field_t = None,
-        ) -> functional_matrix:
-    """
-   Utility operator for converting the E field into the H field.
-   For use with e_full -- assumes that there is no magnetic current M.
-
-   :param omega: Angular frequency of the simulation
-   :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-   :param mu: Magnetic permeability (default 1 everywhere)
-   :return: Function for converting E to H
-   """
-    A2 = curl_e(dxes)
-
-    def e2h_1_1(e):
-        return [y / (-1j * omega) for y in A2(e)]
-
-    def e2h_mu(e):
-        return [y / (-1j * omega * m) for y, m in zip(A2(e), mu)]
-
-    if numpy.any(numpy.equal(mu, None)):
-        return e2h_1_1
-    else:
-        return e2h_mu

fdfd_tools/operators.py

@@ -1,507 +0,0 @@
-"""
-Sparse matrix operators for use with electromagnetic wave equations.
-
-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
- fdfd_tools.vec() and .unvec() functions (column-major/Fortran ordering).
-
-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).
-
-Many of these functions require a 'dxes' parameter, of type fdfd_tools.dx_lists_type,
- which contains grid cell width information in the following format:
- [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...], [dz_e_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,
- and dy_h_0 is  the y-width of the y=0 cells, as used when calculating dH/dy, etc.
-
-
-The following operators are included:
-- E-only wave operator
-- H-only wave operator
-- EH wave operator
-- Curl for use with E, H fields
-- E to H conversion
-- M to J conversion
-- Poynting cross products
-
-Also available:
-- Circular shifts
-- Discrete derivatives
-- Averaging operators
-- Cross product matrices
-"""
-
-from typing import List, Tuple
-import numpy
-import scipy.sparse as sparse
-
-from . import vec, dx_lists_t, vfield_t
-
-
-__author__ = 'Jan Petykiewicz'
-
-
-def e_full(omega: complex,
-           dxes: dx_lists_t,
-           epsilon: vfield_t,
-           mu: vfield_t = None,
-           pec: vfield_t = None,
-           pmc: vfield_t = None,
-           ) -> sparse.spmatrix:
-    """
-    Wave operator del x (1/mu * del x) - omega**2 * epsilon, for use with E-field,
-     with wave equation
-    (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().
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param epsilon: Vectorized dielectric constant
-    :param mu: Vectorized magnetic permeability (default 1 everywhere).
-    :param 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)
-    :param 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)
-    :return: Sparse matrix containing the wave operator
-    """
-    ce = curl_e(dxes)
-    ch = curl_h(dxes)
-
-    if numpy.any(numpy.equal(pec, None)):
-        pe = sparse.eye(epsilon.size)
-    else:
-        pe = sparse.diags(numpy.where(pec, 0, 1))     # Set pe to (not PEC)
-
-    if numpy.any(numpy.equal(pmc, None)):
-        pm = sparse.eye(epsilon.size)
-    else:
-        pm = sparse.diags(numpy.where(pmc, 0, 1))     # set pm to (not PMC)
-
-    e = sparse.diags(epsilon)
-    if numpy.any(numpy.equal(mu, None)):
-        m_div = sparse.eye(epsilon.size)
-    else:
-        m_div = sparse.diags(1 / mu)
-
-    op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe
-    return op
-
-
-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.
-
-    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
-
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :return: 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, :],
-                 dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[0][2][None, None, :]]
-
-    p_vector = numpy.sqrt(vec(p_squared))
-    P_left = sparse.diags(p_vector)
-    P_right = sparse.diags(1 / p_vector)
-    return P_left, P_right
-
-
-def h_full(omega: complex,
-           dxes: dx_lists_t,
-           epsilon: vfield_t,
-           mu: vfield_t = None,
-           pec: vfield_t = None,
-           pmc: vfield_t = None,
-           ) -> sparse.spmatrix:
-    """
-    Wave operator del x (1/epsilon * del x) - omega**2 * mu, for use with H-field,
-     with wave equation
-    (del x (1/epsilon * del x) - omega**2 * mu) H = i * omega * M
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param epsilon: Vectorized dielectric constant
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
-    :param 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)
-    :param 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)
-    :return: Sparse matrix containing the wave operator
-    """
-    ec = curl_e(dxes)
-    hc = curl_h(dxes)
-
-    if numpy.any(numpy.equal(pec, None)):
-        pe = sparse.eye(epsilon.size)
-    else:
-        pe = sparse.diags(numpy.where(pec, 0, 1))    # set pe to (not PEC)
-
-    if numpy.any(numpy.equal(pmc, None)):
-        pm = sparse.eye(epsilon.size)
-    else:
-        pm = sparse.diags(numpy.where(pmc, 0, 1))    # Set pe to (not PMC)
-
-    e_div = sparse.diags(1 / epsilon)
-    if mu is None:
-        m = sparse.eye(epsilon.size)
-    else:
-        m = sparse.diags(mu)
-
-    A = pm @ (ec @ pe @ e_div @ hc - omega**2 * m) @ pm
-    return A
-
-
-def eh_full(omega: complex,
-            dxes: dx_lists_t,
-            epsilon: vfield_t,
-            mu: vfield_t = None,
-            pec: vfield_t = None,
-            pmc: vfield_t = None
-            ) -> sparse.spmatrix:
-    """
-    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],
-      [del x, i * omega * mu]]
-
-    for use with a field vector of the form hstack(vec(E), vec(H)).
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param epsilon: Vectorized dielectric constant
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
-    :param 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)
-    :param 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)
-    :return: Sparse matrix containing the wave operator
-    """
-    if numpy.any(numpy.equal(pec, None)):
-        pe = sparse.eye(epsilon.size)
-    else:
-        pe = sparse.diags(numpy.where(pec, 0, 1))    # set pe to (not PEC)
-
-    if numpy.any(numpy.equal(pmc, None)):
-        pm = sparse.eye(epsilon.size)
-    else:
-        pm = sparse.diags(numpy.where(pmc, 0, 1))    # set pm to (not PMC)
-
-    iwe = pe @ (1j * omega * sparse.diags(epsilon)) @ pe
-    iwm = 1j * omega
-    if not numpy.any(numpy.equal(mu, None)):
-        iwm *= sparse.diags(mu)
-    iwm = pm @ iwm @ pm
-
-    A1 = pe @ curl_h(dxes) @ pm
-    A2 = pm @ curl_e(dxes) @ 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 fdfd_tools.operators header
-    :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 fdfd_tools.operators header
-    :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,
-        pmc: vfield_t = None,
-        ) -> 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.
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
-    :param 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)
-    :return: Sparse matrix for converting E to H
-    """
-    op = curl_e(dxes) / (-1j * omega)
-
-    if not numpy.any(numpy.equal(mu, None)):
-        op = sparse.diags(1 / mu) @ op
-
-    if not numpy.any(numpy.equal(pmc, None)):
-        op = sparse.diags(numpy.where(pmc, 0, 1)) @ op
-
-    return op
-
-
-def m2j(omega: complex,
-        dxes: dx_lists_t,
-        mu: vfield_t = None
-        ) -> sparse.spmatrix:
-    """
-    Utility operator for converting M field into J.
-    Converts a magnetic current M into an electric current J.
-    For use with eg. e_full.
-
-    :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :param mu: Vectorized magnetic permeability (default 1 everywhere)
-    :return: Sparse matrix for converting E to H
-    """
-    op = curl_h(dxes) / (1j * omega)
-
-    if not numpy.any(numpy.equal(mu, None)):
-        op = op @ sparse.diags(1 / mu)
-
-    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 = ijk[0] + ijk[1] * shape[0]
-    if len(shape) == 3:
-        j_ind += ijk[2] * shape[0] * shape[1]
-
-    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, contining the (E x) portion of the Poynting vector.
-
-    :param e: Vectorized E-field for the ExH cross product
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :return: Sparse matrix containing (E x) portion of Poynting cross product
-    """
-    shape = [len(dx) for dx in dxes[0]]
-
-    fx, fy, fz = [avgf(i, shape) for i in range(3)]
-    bx, by, bz = [avgb(i, shape) for i in range(3)]
-
-    dxag = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
-    dbgx, dbgy, dbgz = [sparse.diags(dx.ravel(order='C'))
-                        for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
-
-    Ex, Ey, Ez = [sparse.diags(ei * da) for ei, da in zip(numpy.split(e, 3), dxag)]
-
-    n = numpy.prod(shape)
-    zero = sparse.csr_matrix((n, n))
-
-    P = sparse.bmat(
-        [[ zero,                -fx @ Ez @ bz @ dbgy,  fx @ Ey @ by @ dbgz],
-         [ fy @ Ez @ bz @ dbgx,  zero,                -fy @ Ex @ bx @ dbgz],
-         [-fz @ Ey @ by @ dbgx,  fz @ Ex @ bx @ dbgy,  zero]])
-    return P
-
-
-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
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
-    :return: Sparse matrix containing (H x) portion of Poynting cross product
-    """
-    shape = [len(dx) for dx in dxes[0]]
-
-    fx, fy, fz = [avgf(i, shape) for i in range(3)]
-    bx, by, bz = [avgb(i, shape) for i in range(3)]
-
-    dxbg = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
-    dagx, dagy, dagz = [sparse.diags(dx.ravel(order='C'))
-                        for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
-
-    Hx, Hy, Hz = [sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg)]
-
-    n = numpy.prod(shape)
-    zero = sparse.csr_matrix((n, n))
-
-    P = sparse.bmat(
-        [[ zero,                -by @ Hz @ fx @ dagy,  bz @ Hy @ fx @ dagz],
-         [ bx @ Hz @ fy @ dagx,  zero,                -bz @ Hx @ fy @ dagz],
-         [-bx @ Hy @ fz @ dagx,  by @ Hx @ fz @ dagy,  zero]])
-    return P

fdfd_tools/solvers.py

@@ -1,118 +0,0 @@
-"""
-Solvers for FDFD problems.
-"""
-
-from typing import List, Callable, Dict, Any
-import logging
-
-import numpy
-from numpy.linalg import norm
-import scipy.sparse.linalg
-
-from . import operators
-
-
-logger = logging.getLogger(__name__)
-
-
-def _scipy_qmr(A: scipy.sparse.csr_matrix,
-               b: numpy.ndarray,
-               **kwargs
-               ) -> numpy.ndarray:
-    """
-    Wrapper for scipy.sparse.linalg.qmr
-
-    :param A: Sparse matrix
-    :param b: Right-hand-side vector
-    :param kwargs: Passed as **kwargs to the wrapped function
-    :return: Guess for solution (returned even if didn't converge)
-    """
-
-    '''
-    Report on our progress
-    '''
-    iter = 0
-
-    def log_residual(xk):
-        nonlocal iter
-        iter += 1
-        if iter % 100 == 0:
-            logger.info('Solver residual at iteration {} : {}'.format(iter, norm(A @ xk - b)))
-
-    if 'callback' in kwargs:
-        def augmented_callback(xk):
-            log_residual(xk)
-            kwargs['callback'](xk)
-
-        kwargs['callback'] = augmented_callback
-    else:
-        kwargs['callback'] = log_residual
-
-    '''
-    Run the actual solve
-    '''
-
-    x, _ = scipy.sparse.linalg.qmr(A, b, **kwargs)
-    return x
-
-
-def generic(omega: complex,
-            dxes: List[List[numpy.ndarray]],
-            J: numpy.ndarray,
-            epsilon: numpy.ndarray,
-            mu: numpy.ndarray = None,
-            pec: numpy.ndarray = None,
-            pmc: numpy.ndarray = None,
-            adjoint: bool = False,
-            matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr,
-            matrix_solver_opts: Dict[str, Any] = None,
-            ) -> numpy.ndarray:
-    """
-    Conjugate gradient FDFD solver using CSR sparse matrices.
-
-    All ndarray arguments should be 1D array, as returned by fdfd_tools.vec().
-
-    :param omega: Complex frequency to solve at.
-    :param dxes: [[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]] (complex cell sizes)
-    :param J: Electric current distribution (at E-field locations)
-    :param epsilon: Dielectric constant distribution (at E-field locations)
-    :param mu: Magnetic permeability distribution (at H-field locations)
-    :param pec: Perfect electric conductor distribution
-        (at E-field locations; non-zero value indicates PEC is present)
-    :param pmc: Perfect magnetic conductor distribution
-        (at H-field locations; non-zero value indicates PMC is present)
-    :param adjoint: If true, solves the adjoint problem.
-    :param matrix_solver: Called as matrix_solver(A, b, **matrix_solver_opts) -> x
-        Where A: scipy.sparse.csr_matrix
-              b: numpy.ndarray
-              x: numpy.ndarray
-        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(...)
-    :return: E-field which solves the system.
-    """
-
-    if matrix_solver_opts is None:
-        matrix_solver_opts = dict()
-
-    b0 = -1j * omega * J
-    A0 = operators.e_full(omega, dxes, epsilon=epsilon, mu=mu, pec=pec, pmc=pmc)
-
-    Pl, Pr = operators.e_full_preconditioners(dxes)
-
-    if adjoint:
-        A = (Pl @ A0 @ Pr).H
-        b = Pr.H @ b0
-    else:
-        A = Pl @ A0 @ Pr
-        b = Pl @ b0
-
-    x = matrix_solver(A.tocsr(), b, **matrix_solver_opts)
-
-    if adjoint:
-        x0 = Pl.H @ x
-    else:
-        x0 = Pr @ x
-
-    return x0

fdfd_tools/vectorization.py

@@ -1,49 +0,0 @@
-"""
-Functions for moving between a vector field (list of 3 ndarrays, [f_x, f_y, f_z])
-and a 1D array representation of that field [f_x0, f_x1, f_x2,... f_y0,... f_z0,...].
-Vectorized versions of the field use row-major (ie., C-style) ordering.
-"""
-
-
-from typing import List
-import numpy
-
-__author__ = 'Jan Petykiewicz'
-
-# Types
-field_t = List[numpy.ndarray]  # vector field (eg. [E_x, E_y, E_z]
-vfield_t = numpy.ndarray       # linearized vector field
-
-
-def vec(f: field_t) -> vfield_t:
-    """
-    Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
-
-    Returns None if called with f=None.
-
-    :param f: A vector field, [f_x, f_y, f_z] where each f_ component is a 1 to
-        3D ndarray (f_* should all be the same size). Doesn't fail with f=None.
-    :return: A 1D ndarray containing the linearized field (or None)
-    """
-    if numpy.any(numpy.equal(f, None)):
-        return None
-    return numpy.hstack(tuple((fi.ravel(order='C') for fi in f)))
-
-
-def unvec(v: vfield_t, shape: numpy.ndarray) -> field_t:
-    """
-    Perform the inverse of vec(): take a 1D ndarray and output a 3D field
-     of form [f_x, f_y, f_z] where each of f_* is a len(shape)-dimensional
-     ndarray.
-
-    Returns None if called with v=None.
-
-    :param v: 1D ndarray representing a 3D vector field of shape shape (or None)
-    :param shape: shape of the vector field
-    :return: [f_x, f_y, f_z] where each f_ is a len(shape) dimensional ndarray
-     (or None)
-    """
-    if numpy.any(numpy.equal(v, None)):
-        return None
-    return [vi.reshape(shape, order='C') for vi in numpy.split(v, 3)]
-

fdfd_tools/waveguide.py

@@ -1,368 +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.
-"""
-
-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(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 fdfd_tools.operators header (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 ** 2 * 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(v: numpy.ndarray,
-                      wavenumber: complex,
-                      omega: complex,
-                      dxes: dx_lists_t,
-                      epsilon: vfield_t,
-                      mu: vfield_t = None
-                      ) -> Tuple[vfield_t, vfield_t]:
-    """
-    Given a vector v containing the vectorized H_x and H_y fields,
-     returns normalized, vectorized E and H fields for the system.
-
-    :param v: Vector containing H_x and H_y fields
-    :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 fdfd_tools.operators header (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Normalized, vectorized (e, h) containing all vector components.
-    """
-    e = v2e(v, wavenumber, omega, dxes, epsilon, mu=mu)
-    h = v2h(v, wavenumber, dxes, mu=mu)
-
-    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)
-
-    S1 = E[0] * numpy.roll(numpy.conj(H[1]), 1, axis=0) * dxes_real[0][1] * dxes_real[1][0]
-    S2 = E[1] * numpy.roll(numpy.conj(H[0]), 1, axis=1) * dxes_real[0][0] * dxes_real[1][1]
-    S = 0.25 * ((S1 + numpy.roll(S1, 1, axis=0)) -
-                (S2 + numpy.roll(S2, 1, axis=1)))
-    P = 0.5 * numpy.real(S.sum())
-    assert P > 0, 'Found a mode propagating in the wrong direction! P={}'.format(P)
-
-    norm_amplitude = 1 / numpy.sqrt(P)
-    norm_angle = -numpy.angle(e[e.size//2])
-    norm_factor = norm_amplitude * numpy.exp(1j * norm_angle)
-
-    e *= norm_factor
-    h *= norm_factor
-
-    return e, h
-
-
-def v2h(v: numpy.ndarray,
-        wavenumber: complex,
-        dxes: dx_lists_t,
-        mu: vfield_t = None
-        ) -> vfield_t:
-    """
-    Given a vector v containing the vectorized H_x and H_y fields,
-     returns a vectorized H including all three H components.
-
-    :param v: Vector containing H_x and H_y fields
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Vectorized H field with all vector components
-    """
-    Dfx, Dfy = operators.deriv_forward(dxes[0])
-    op = sparse.hstack((Dfx, Dfy))
-
-    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])
-
-        op = mu_z_inv @ op @ mu_xy
-
-    w = op @ v / (1j * wavenumber)
-    return numpy.hstack((v, w)).flatten()
-
-
-def v2e(v: numpy.ndarray,
-        wavenumber: complex,
-        omega: complex,
-        dxes: dx_lists_t,
-        epsilon: vfield_t,
-        mu: vfield_t = None
-        ) -> vfield_t:
-    """
-    Given a vector v containing the vectorized H_x and H_y fields,
-     returns a vectorized E containing all three E components
-
-    :param v: Vector containing H_x and H_y fields
-    :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 fdfd_tools.operators header (2D)
-    :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Vectorized E field with all vector components.
-    """
-    h2eop = h2e(wavenumber, omega, dxes, epsilon)
-    return h2eop @ v2h(v, wavenumber, dxes, mu)
-
-
-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 fdfd_tools.operators header (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 fdfd_tools.operators header (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 fdfd_tools.operators header (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 fdfd_tools.operators header (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 fdfd_tools.operators header (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 fdfd_tools.operators header (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 fdfd_tools.operators header (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
-
-
-

fdfd_tools/waveguide_mode.py

@@ -1,340 +0,0 @@
-from typing import Dict, List
-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 solve_waveguide_mode_2d(mode_number: int,
-                            omega: complex,
-                            dxes: dx_lists_t,
-                            epsilon: vfield_t,
-                            mu: vfield_t = None,
-                            wavenumber_correction: bool = True,
-                            ) -> Dict[str, complex or field_t]:
-    """
-    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 fdfd_tools.operators header
-    :param epsilon: Dielectric constant
-    :param mu: Magnetic permeability (default 1 everywhere)
-    :param wavenumber_correction: Whether to correct the wavenumber to
-        account for numerical dispersion (default True)
-    :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.operator(numpy.real(omega), dxes_real, numpy.real(epsilon), numpy.real(mu))
-
-    eigvals, eigvecs = signed_eigensolve(A_r, mode_number+3)
-    v = 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(omega, dxes, epsilon, mu)
-    eigval, v = rayleigh_quotient_iteration(A, v)
-
-    # Calculate the wave-vector (force the real part to be positive)
-    wavenumber = numpy.sqrt(eigval)
-    wavenumber *= numpy.sign(numpy.real(wavenumber))
-
-    e, h = waveguide.normalized_fields(v, wavenumber, omega, dxes, epsilon, mu)
-
-    '''
-    Perform correction on wavenumber to account for numerical dispersion.
-
-     See Numerical Dispersion in Taflove's FDTD book.
-     This correction term reduces the error in emitted power, but additional
-      error is introduced into the E_err and H_err terms. This effect becomes
-      more pronounced as beta increases.
-    '''
-    if wavenumber_correction:
-        wavenumber -= 2 * numpy.sin(numpy.real(wavenumber / 2)) - numpy.real(wavenumber)
-
-    shape = [d.size for d in dxes[0]]
-    fields = {
-        'wavenumber': wavenumber,
-        'E': unvec(e, shape),
-        'H': unvec(h, shape),
-    }
-
-    return fields
-
-
-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,
-                         wavenumber_correction: bool = True
-                         ) -> 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 fdfd_tools.operators header
-    :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)
-    :param wavenumber_correction: Whether to correct the wavenumber to
-        account for numerical dispersion (default True)
-    :return: {'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}
-    """
-    if mu is None:
-        mu = [numpy.ones_like(epsilon[0])] * 3
-
-    '''
-    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)
-
-    # Reduce to 2D and solve the 2D problem
-    args_2d = {
-        '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]),
-        'wavenumber_correction': wavenumber_correction,
-    }
-    fields_2d = solve_waveguide_mode_2d(mode_number, omega=omega, **args_2d)
-
-    '''
-    Apply corrections and expand to 3D
-    '''
-    # Scale based on dx in propagation direction
-    dxab_forward = numpy.array([dx[order[2]][slices[order[2]]] for dx in dxes])
-
-    # Adjust for propagation direction
-    fields_2d['E'][2] *= polarity
-    fields_2d['H'][2] *= polarity
-
-    # Apply phase shift to H-field
-    d_prop = 0.5 * sum(dxab_forward)
-    for a in range(3):
-        fields_2d['H'][a] *= numpy.exp(-polarity * 1j * 0.5 * fields_2d['wavenumber'] * d_prop)
-
-    # Expand E, H to full epsilon space we were given
-    E = [None]*3
-    H = [None]*3
-    for a, o in enumerate(reverse_order):
-        E[a] = numpy.zeros_like(epsilon[0], dtype=complex)
-        H[a] = numpy.zeros_like(epsilon[0], dtype=complex)
-
-        E[a][slices] = fields_2d['E'][o][:, :, None].transpose(reverse_order)
-        H[a][slices] = fields_2d['H'][o][:, :, None].transpose(reverse_order)
-
-    results = {
-        'wavenumber': fields_2d['wavenumber'],
-        'H': H,
-        'E': E,
-    }
-
-    return results
-
-
-def compute_source(E: field_t,
-                   H: field_t,
-                   wavenumber: complex,
-                   omega: complex,
-                   dxes: dx_lists_t,
-                   axis: int,
-                   polarity: int,
-                   slices: List[slice],
-                   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 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 fdfd_tools.operators header
-    :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
-    """
-    if mu is None:
-        mu = [1] * 3
-
-    J = [None]*3
-    M = [None]*3
-
-    src_order = numpy.roll(range(3), axis)
-    exp_iphi = numpy.exp(1j * polarity * wavenumber * dxes[1][axis][slices[axis]])
-    J[src_order[0]] = numpy.zeros_like(E[0])
-    J[src_order[1]] = +exp_iphi * H[src_order[2]] * polarity
-    J[src_order[2]] = -exp_iphi * H[src_order[1]] * polarity
-
-    M[src_order[0]] = numpy.zeros_like(E[0])
-    M[src_order[1]] = +numpy.roll(E[src_order[2]], -1, axis=axis)
-    M[src_order[2]] = -numpy.roll(E[src_order[1]], -1, axis=axis)
-
-    A1f = functional.curl_h(dxes)
-
-    Jm_iw = A1f([M[k] / mu[k] for k in range(3)])
-    for k in range(3):
-        J[k] += Jm_iw[k] / (-1j * omega)
-
-    return J
-
-
-def compute_overlap_e(E: field_t,
-                      H: field_t,
-                      wavenumber: complex,
-                      omega: complex,
-                      dxes: dx_lists_t,
-                      axis: int,
-                      polarity: int,
-                      slices: List[slice],
-                      mu: field_t = None,
-                      ) -> field_t:
-    """
-    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].
-
-    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 fdfd_tools.operators header
-    :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
-    """
-    cross_plane = [slice(None)] * 3
-    cross_plane[axis] = slices[axis]
-
-    # Determine phase factors for parallel slices
-    a_shape = numpy.roll([-1, 1, 1], axis)
-    a_E = numpy.real(dxes[0][axis]).cumsum()
-    a_H = numpy.real(dxes[1][axis]).cumsum()
-    iphi = -polarity * 1j * wavenumber
-    phase_E = numpy.exp(iphi * (a_E - a_E[slices[axis]])).reshape(a_shape)
-    phase_H = numpy.exp(iphi * (a_H - a_H[slices[axis]])).reshape(a_shape)
-
-    # Expand our slice to the entire grid using the calculated phase factors
-    Ee = [None]*3
-    He = [None]*3
-    for k in range(3):
-        Ee[k] = phase_E * E[k][tuple(cross_plane)]
-        He[k] = phase_H * H[k][tuple(cross_plane)]
-
-
-    # Write out the operator product for the mode orthogonality integral
-    domain = numpy.zeros_like(E[0], dtype=int)
-    domain[slices] = 1
-
-    npts = E[0].size
-    dn = numpy.zeros(npts * 3, dtype=int)
-    dn[0:npts] = 1
-    dn = numpy.roll(dn, npts * axis)
-
-    e2h = operators.e2h(omega, dxes, mu)
-    ds = sparse.diags(vec([domain]*3))
-    h_cross_ = operators.poynting_h_cross(vec(He), dxes)
-    e_cross_ = operators.poynting_e_cross(vec(Ee), dxes)
-
-    overlap_e = dn @ ds @ (-h_cross_ + e_cross_ @ e2h)
-
-    # Normalize
-    dx_forward = dxes[0][axis][slices[axis]]
-    norm_factor = numpy.abs(overlap_e @ vec(Ee))
-    overlap_e /= norm_factor * dx_forward
-
-    return unvec(overlap_e, E[0].shape)
-
-
-def solve_waveguide_mode_cylindrical(mode_number: int,
-                                     omega: complex,
-                                     dxes: dx_lists_t,
-                                     epsilon: vfield_t,
-                                     r0: float,
-                                     wavenumber_correction: bool = True,
-                                     ) -> Dict[str, complex or field_t]:
-    """
-    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 fdfd_tools.operators header.
-        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.
-    :param wavenumber_correction: Whether to correct the wavenumber to
-        account for numerical dispersion (default True)
-    :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)
-    v = 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, v = rayleigh_quotient_iteration(A, v)
-
-    # Calculate the wave-vector (force the real part to be positive)
-    wavenumber = numpy.sqrt(eigval)
-    wavenumber *= numpy.sign(numpy.real(wavenumber))
-
-    '''
-    Perform correction on wavenumber to account for numerical dispersion.
-
-     See Numerical Dispersion in Taflove's FDTD book.
-     This correction term reduces the error in emitted power, but additional
-      error is introduced into the E_err and H_err terms. This effect becomes
-      more pronounced as the wavenumber increases.
-    '''
-    if wavenumber_correction:
-        wavenumber -= 2 * numpy.sin(numpy.real(wavenumber / 2)) - numpy.real(wavenumber)
-
-    shape = [d.size for d in dxes[0]]
-    v = numpy.hstack((v, numpy.zeros(shape[0] * shape[1])))
-    fields = {
-        'wavenumber': wavenumber,
-        'E': unvec(v, shape),
-#        'E': unvec(e, shape),
-#        'H': unvec(h, shape),
-    }
-
-    return fields

make_docs.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+cd ~/projects/meanas
+
+# Approach 1: pdf to html?
+#pdoc3 --pdf --force --template-dir pdoc_templates -o doc . | \
+#    pandoc --metadata=title:"meanas" --toc --toc-depth=4 --from=markdown+abbreviations --to=html --output=doc.html --gladtex -s -
+
+# Approach 2: pdf to html with gladtex
+pdoc3 --pdf --force --template-dir pdoc_templates -o doc . > doc.md
+pandoc --metadata=title:"meanas" --from=markdown+abbreviations --to=html --output=doc.html --gladtex -s --css pdoc_templates/pdoc.css doc.md
+gladtex -a -n -d _doc_mathimg -c white doc.html
+
+# Approach 3: html with gladtex
+#pdoc3 --html --force --template-dir pdoc_templates -o doc .
+#find doc -iname '*.html' -exec gladtex -a -n -d _mathimg -c white {} \;

meanas/VERSION

@@ -0,0 +1 @@
+0.5

meanas/__init__.py

@@ -0,0 +1,16 @@
+"""
+Electromagnetic simulation tools
+
+See the readme or `import meanas; help(meanas)` for more info.
+"""
+
+import pathlib
+
+__author__ = 'Jan Petykiewicz'
+
+with open(pathlib.Path(__file__).parent / 'VERSION', 'r') as f:
+    __version__ = f.read().strip()
+
+with open(pathlib.Path(__file__).parent.parent / 'README.md', 'r') as f:
+    __doc__ = f.read()
+

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.
-    :param guess_vector: Starting point for the eigenvector. Default is a randomly chosen vector.
-    :param iterations: Number of iterations to perform. Default 20.
-    :return: (Largest-magnitude eigenvalue, Corresponding eigenvector estimate)
+    Args:
+       operator: Matrix to analyze.
+       guess_vector: Starting point for the eigenvector. Default is a randomly chosen vector.
+       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])
@@ -37,20 +40,23 @@ def rayleigh_quotient_iteration(operator: sparse.spmatrix or spalg.LinearOperato
                                 guess_vector: numpy.ndarray,
                                 iterations: int = 40,
                                 tolerance: float = 1e-13,
-                                solver=None,
+                                solver = None,
                                 ) -> Tuple[complex, numpy.ndarray]:
     """
     Use Rayleigh quotient iteration to refine an eigenvector guess.
 
-    :param operator: Matrix to analyze.
-    :param guess_vector: Eigenvector to refine.
-    :param iterations: Maximum number of iterations to perform. Default 40.
-    :param tolerance: Stop iteration if (A - I*eigenvalue) @ v < tolerance.
-        Default 1e-13.
-    :param solver: Solver function of the form x = solver(A, b).
-        By default, use scipy.sparse.spsolve for sparse matrices and
-        scipy.sparse.bicgstab for general LinearOperator instances.
-    :return: (eigenvalue, eigenvector)
+    Args:
+        operator: Matrix to analyze.
+        guess_vector: Eigenvector to refine.
+        iterations: Maximum number of iterations to perform. Default 40.
+        tolerance: Stop iteration if `(A - I*eigenvalue) @ v < num_vectors * tolerance`,
+                    Default 1e-13.
+        solver: Solver function of the form `x = solver(A, b)`.
+                By default, use scipy.sparse.spsolve for sparse matrices and
+                scipy.sparse.bicgstab for general LinearOperator instances.
+
+    Returns:
+        (eigenvalues, eigenvectors)
     """
     try:
         _test = operator - sparse.eye(operator.shape[0])
@@ -64,7 +70,7 @@ def rayleigh_quotient_iteration(operator: sparse.spmatrix or spalg.LinearOperato
         if solver is None:
             solver = lambda A, b: spalg.bicgstab(A, b)[0]
 
-    v = guess_vector
+    v = numpy.squeeze(guess_vector)
     v /= norm(v)
     for _ in range(iterations):
         eigval = v.conj() @ (operator @ v)
@@ -85,12 +91,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.
-    :param how_many: How many eigenvalues to find.
-    :param 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
+    Args:
+        operator: Matrix to analyze.
+        how_many: How many eigenvalues to find.
+        negative: Whether to find negative-only eigenvalues.
+                  Default False (positive only).
+
+    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)
@@ -99,7 +108,7 @@ def signed_eigensolve(operator: sparse.spmatrix or spalg.LinearOperator,
     Shift by the absolute value of the largest eigenvalue, then find a few of the
      largest-magnitude (shifted) eigenvalues. A positive shift ensures that we find the
      largest _positive_ eigenvalues, since any negative eigenvalues will be shifted to the
-     range 0 >= neg_eigval + abs(lm_eigval) > abs(lm_eigval)
+     range `0 >= neg_eigval + abs(lm_eigval) > abs(lm_eigval)`
     '''
     shift = numpy.abs(lm_eigval)
     if negative:

meanas/fdfd/__init__.py

@@ -0,0 +1,38 @@
+"""
+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 the "Frequency domain" section of `meanas.fdmath`, we have
+
+$$
+ \\begin{aligned}
+ \\tilde{E}_{l, \\vec{r}} &= \\tilde{E}_{\\vec{r}} e^{-\\imath \\omega l \\Delta_t} \\\\
+ \\tilde{J}_{l, \\vec{r}} &= \\tilde{J}_{\\vec{r}} e^{-\\imath \\omega (l - \\frac{1}{2}) \\Delta_t} \\\\
+ \\hat{\\nabla} \\times (\\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{\\vec{r}})
+    -\\Omega^2 \\epsilon_{\\vec{r}} \\cdot \\tilde{E}_{\\vec{r}} &= \\imath \\Omega \\tilde{J}_{\\vec{r}} \\\\
+ \\Omega &= 2 \\sin(\\omega \\Delta_t / 2) / \\Delta_t
+ \\end{aligned}
+$$
+
+
+# TODO FDFD?
+# TODO PML
+
+
+"""
+from . import solvers, operators, functional, scpml, waveguide_2d, waveguide_3d
+# from . import farfield, bloch TODO

meanas/fdfd/bloch.py

@@ -0,0 +1,691 @@
+'''
+Bloch eigenmode solver/operators
+
+This module contains functions for generating and solving the
+ 3D Bloch eigenproblem. The approach is to transform the problem
+ into the (spatial) fourier domain, transforming the equation
+
+    1/mu * curl(1/eps * curl(H)) = (w/c)^2 H
+
+ into
+
+    conv(1/mu_k, ik x conv(1/eps_k, ik x H_k)) = (w/c)^2 H_k
+
+ where:
+
+  - the `_k` subscript denotes a 3D fourier transformed field
+  - each component of `H_k` corresponds to a plane wave with wavevector `k`
+  - `x` is the cross product
+  - `conv()` denotes convolution
+
+ Since `k` and `H` are orthogonal for each plane wave, we can use each
+ `k` to create an orthogonal basis (k, m, n), with `k x m = n`, and
+ `|m| = |n| = 1`. The cross products are then simplified with
+
+    k @ h = kx hx + ky hy + kz hz = 0 = hk
+    h = hk + hm + hn = hm + hn
+    k = kk + km + kn = kk = |k|
+
+    k x h = (ky hz - kz hy,
+             kz hx - kx hz,
+             kx hy - ky hx)
+          = ((k x h) @ k, (k x h) @ m, (k x h) @ n)_kmn
+          = (0, (m x k) @ h, (n x k) @ h)_kmn         # triple product ordering
+          = (0, kk (-n @ h), kk (m @ h))_kmn          # (m x k) = -|k| n, etc.
+          = |k| (0, -h @ n, h @ m)_kmn
+
+    k x h = (km hn - kn hm,
+             kn hk - kk hn,
+             kk hm - km hk)_kmn
+          = (0, -kk hn, kk hm)_kmn
+          = (-kk hn)(mx, my, mz) + (kk hm)(nx, ny, nz)
+          = |k| (hm * (nx, ny, nz) - hn * (mx, my, mz))
+
+ where `h` is shorthand for `H_k`, `(...)_kmn` deontes the `(k, m, n)` basis,
+ and e.g. `hm` is the component of `h` in the `m` direction.
+
+ We can also simplify `conv(X_k, Y_k)` as `fftn(X * ifftn(Y_k))`.
+
+ Using these results and storing `H_k` as `h = (hm, hn)`, we have
+
+    e_xyz = fftn(1/eps * ifftn(|k| (hm * n - hn * m)))
+    b_mn = |k| (-e_xyz @ n, e_xyz @ m)
+    h_mn = fftn(1/mu * ifftn(b_m * m + b_n * n))
+
+ which forms the operator from the left side of the equation.
+
+ We can then use a preconditioned block Rayleigh iteration algorithm, as in
+  SG Johnson and JD Joannopoulos, Block-iterative frequency-domain methods
+  for Maxwell's equations in a planewave basis, Optics Express 8, 3, 173-190 (2001)
+ (similar to that used in MPB) to find the eigenvectors for this operator.
+
+ ===
+
+ Typically you will want to do something like
+
+    recip_lattice = numpy.diag(1/numpy.array(epsilon[0].shape * dx))
+    n, v = bloch.eigsolve(5, k0, recip_lattice, epsilon)
+    f = numpy.sqrt(-numpy.real(n[0]))
+    n_eff = norm(recip_lattice @ k0) / f
+
+    v2e = bloch.hmn_2_exyz(k0, recip_lattice, epsilon)
+    e_field = v2e(v[0])
+
+    k, f = find_k(frequency=1/1550,
+                  tolerance=(1/1550 - 1/1551),
+                  direction=[1, 0, 0],
+                  G_matrix=recip_lattice,
+                  epsilon=epsilon,
+                  band=0)
+
+'''
+
+from typing import Tuple, Callable
+import logging
+import numpy
+from numpy import pi, real, trace
+from numpy.fft import fftfreq
+import scipy
+import scipy.optimize
+from scipy.linalg import norm
+import scipy.sparse.linalg as spalg
+
+from ..fdmath import fdfield_t
+
+logger = logging.getLogger(__name__)
+
+
+try:
+    import pyfftw.interfaces.numpy_fft
+    import pyfftw.interfaces
+    import multiprocessing
+    logger.info('Using pyfftw')
+
+    pyfftw.interfaces.cache.enable()
+    pyfftw.interfaces.cache.set_keepalive_time(3600)
+    fftw_args = {
+        'threads': multiprocessing.cpu_count(),
+        'overwrite_input': True,
+        'planner_effort': 'FFTW_EXHAUSTIVE',
+        }
+
+    def fftn(*args, **kwargs):
+        return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args)
+
+    def ifftn(*args, **kwargs):
+        return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)
+
+except ImportError:
+    from numpy.fft import fftn, ifftn
+    logger.info('Using numpy fft')
+
+
+def generate_kmn(k0: numpy.ndarray,
+                 G_matrix: numpy.ndarray,
+                 shape: numpy.ndarray
+                 ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+    """
+    Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.
+
+    Args:
+        k0: [k0x, k0y, k0z], Bloch wavevector, in G basis.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        shape: [nx, ny, nz] shape of the simulation grid.
+
+    Returns:
+        `(|k|, m, n)` where `|k|` has shape `tuple(shape) + (1,)`
+            and `m`, `n` have shape `tuple(shape) + (3,)`.
+            All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
+    """
+    k0 = numpy.array(k0)
+
+    Gi_grids = numpy.meshgrid(*(fftfreq(n, 1/n) for n in shape[:3]), indexing='ij')
+    Gi = numpy.stack(Gi_grids, axis=3)
+
+    k_G = k0[None, None, None, :] - Gi
+    k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2)
+
+    m = numpy.broadcast_to([0, 1, 0], tuple(shape[:3]) + (3,)).astype(float)
+    n = numpy.broadcast_to([0, 0, 1], tuple(shape[:3]) + (3,)).astype(float)
+
+    xy_non0 = numpy.any(k_xyz[:, :, :, 0:1] != 0, axis=3)
+    if numpy.any(xy_non0):
+        u = numpy.cross(k_xyz[xy_non0], [0, 0, 1])
+        m[xy_non0, :] = u / norm(u, axis=1)[:, None]
+
+    z_non0 = numpy.any(k_xyz != 0, axis=3)
+    if numpy.any(z_non0):
+        v = numpy.cross(k_xyz[z_non0], m[z_non0])
+        n[z_non0, :] = v / norm(v, axis=1)[:, None]
+
+    k_mag = norm(k_xyz, axis=3)[:, :, :, None]
+    return k_mag, m, n
+
+
+def maxwell_operator(k0: numpy.ndarray,
+                     G_matrix: numpy.ndarray,
+                     epsilon: fdfield_t,
+                     mu: fdfield_t = None
+                     ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+    """
+    Generate the Maxwell operator
+
+        conv(1/mu_k, ik x conv(1/eps_k, ik x ___))
+
+    which is the spatial-frequency-space representation of
+
+        1/mu * curl(1/eps * curl(___))
+
+    The operator is a function that acts on a vector h_mn of size `2 * epsilon[0].size`
+
+    See the `meanas.fdfd.bloch` docstring for more information.
+
+    Args:
+        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        epsilon: Dielectric constant distribution for the simulation.
+                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+        mu: Magnetic permability distribution for the simulation.
+            Default None (1 everywhere).
+
+    Returns:
+        Function which applies the maxwell operator to h_mn.
+    """
+
+    shape = epsilon[0].shape + (1,)
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    epsilon = numpy.stack(epsilon, 3)
+    if mu is not None:
+        mu = numpy.stack(mu, 3)
+
+    def operator(h: numpy.ndarray):
+        """
+        Maxwell operator for Bloch eigenmode simulation.
+
+        h is complex 2-field in (m, n) basis, vectorized
+
+        Args:
+            h: Raveled h_mn; size `2 * epsilon[0].size`.
+
+        Returns:
+            Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)).
+        """
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+
+        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
+
+        # cross product and transform into xyz basis
+        d_xyz = (n * hin_m -
+                 m * hin_n) * k_mag
+
+        # divide by epsilon
+        e_xyz = fftn(ifftn(d_xyz, axes=range(3)) / epsilon, axes=range(3))
+
+        # cross product and transform into mn basis
+        b_m = numpy.sum(e_xyz * n, axis=3)[:, :, :, None] * -k_mag
+        b_n = numpy.sum(e_xyz * m, axis=3)[:, :, :, None] * +k_mag
+
+        if mu is None:
+            h_m, h_n = b_m, b_n
+        else:
+            # transform from mn to xyz
+            b_xyz = (m * b_m[:, :, :, None] +
+                     n * b_n[:, :, :, None])
+
+            # divide by mu
+            h_xyz = fftn(ifftn(b_xyz, axes=range(3)) / mu, axes=range(3))
+
+            # transform back to mn
+            h_m = numpy.sum(h_xyz * m, axis=3)
+            h_n = numpy.sum(h_xyz * n, axis=3)
+        return numpy.hstack((h_m.ravel(), h_n.ravel()))
+
+    return operator
+
+
+def hmn_2_exyz(k0: numpy.ndarray,
+               G_matrix: numpy.ndarray,
+               epsilon: fdfield_t,
+               ) -> Callable[[numpy.ndarray], fdfield_t]:
+    """
+    Generate an operator which converts a vectorized spatial-frequency-space
+     `h_mn` into an E-field distribution, i.e.
+
+        ifft(conv(1/eps_k, ik x h_mn))
+
+    The operator is a function that acts on a vector `h_mn` of size `2 * epsilon[0].size`.
+
+    See the `meanas.fdfd.bloch` docstring for more information.
+
+    Args:
+        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        epsilon: Dielectric constant distribution for the simulation.
+                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+
+    Returns:
+        Function for converting `h_mn` into `E_xyz`
+    """
+    shape = epsilon[0].shape + (1,)
+    epsilon = numpy.stack(epsilon, 3)
+
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    def operator(h: numpy.ndarray) -> fdfield_t:
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+        d_xyz = (n * hin_m -
+                 m * hin_n) * k_mag
+
+        # divide by epsilon
+        return [ei for ei in numpy.rollaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3)]
+
+    return operator
+
+
+def hmn_2_hxyz(k0: numpy.ndarray,
+               G_matrix: numpy.ndarray,
+               epsilon: fdfield_t
+               ) -> Callable[[numpy.ndarray], fdfield_t]:
+    """
+    Generate an operator which converts a vectorized spatial-frequency-space
+     `h_mn` into an H-field distribution, i.e.
+
+        ifft(h_mn)
+
+    The operator is a function that acts on a vector `h_mn` of size `2 * epsilon[0].size`.
+
+    See the `meanas.fdfd.bloch` docstring for more information.
+
+    Args:
+        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        epsilon: Dielectric constant distribution for the simulation.
+                 Only `epsilon[0].shape` is used.
+
+    Returns:
+        Function for converting `h_mn` into `H_xyz`
+    """
+    shape = epsilon[0].shape + (1,)
+    _k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    def operator(h: numpy.ndarray):
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+        h_xyz = (m * hin_m +
+                 n * hin_n)
+        return [ifftn(hi) for hi in numpy.rollaxis(h_xyz, 3)]
+
+    return operator
+
+
+def inverse_maxwell_operator_approx(k0: numpy.ndarray,
+                                    G_matrix: numpy.ndarray,
+                                    epsilon: fdfield_t,
+                                    mu: fdfield_t = None
+                                    ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+    """
+    Generate an approximate inverse of the Maxwell operator,
+
+        ik x conv(eps_k, ik x conv(mu_k, ___))
+
+     which can be used to improve the speed of ARPACK in shift-invert mode.
+
+    See the `meanas.fdfd.bloch` docstring for more information.
+
+    Args:
+        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        epsilon: Dielectric constant distribution for the simulation.
+                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+        mu: Magnetic permability distribution for the simulation.
+            Default None (1 everywhere).
+
+    Returns:
+        Function which applies the approximate inverse of the maxwell operator to `h_mn`.
+    """
+    shape = epsilon[0].shape + (1,)
+    epsilon = numpy.stack(epsilon, 3)
+
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    if mu is not None:
+        mu = numpy.stack(mu, 3)
+
+    def operator(h: numpy.ndarray):
+        """
+        Approximate inverse Maxwell operator for Bloch eigenmode simulation.
+
+        h is complex 2-field in (m, n) basis, vectorized
+
+        Args:
+            h: Raveled h_mn; size `2 * epsilon[0].size`.
+
+        Returns:
+            Raveled ik x conv(eps_k, ik x conv(mu_k, h_mn))
+        """
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+
+        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
+
+        if mu is None:
+            b_m, b_n = hin_m, hin_n
+        else:
+            # transform from mn to xyz
+            h_xyz = (m * hin_m[:, :, :, None] +
+                     n * hin_n[:, :, :, None])
+
+            # multiply by mu
+            b_xyz = fftn(ifftn(h_xyz, axes=range(3)) * mu, axes=range(3))
+
+            # transform back to mn
+            b_m = numpy.sum(b_xyz * m, axis=3)
+            b_n = numpy.sum(b_xyz * n, axis=3)
+
+        # cross product and transform into xyz basis
+        e_xyz = (n * b_m -
+                 m * b_n) / k_mag
+
+        # multiply by epsilon
+        d_xyz = fftn(ifftn(e_xyz, axes=range(3)) * epsilon, axes=range(3))
+
+        # cross product and transform into mn basis   crossinv_t2c
+        h_m = numpy.sum(d_xyz * n, axis=3)[:, :, :, None] / +k_mag
+        h_n = numpy.sum(d_xyz * m, axis=3)[:, :, :, None] / -k_mag
+
+        return numpy.hstack((h_m.ravel(), h_n.ravel()))
+
+    return operator
+
+
+def find_k(frequency: float,
+           tolerance: float,
+           direction: numpy.ndarray,
+           G_matrix: numpy.ndarray,
+           epsilon: fdfield_t,
+           mu: fdfield_t = None,
+           band: int = 0,
+           k_min: float = 0,
+           k_max: float = 0.5,
+           solve_callback: Callable = None
+           ) -> Tuple[numpy.ndarray, float]:
+    """
+    Search for a bloch vector that has a given frequency.
+
+    Args:
+        frequency: Target frequency.
+        tolerance: Target frequency tolerance.
+        direction: k-vector direction to search along.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        epsilon: Dielectric constant distribution for the simulation.
+                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+        mu: Magnetic permability distribution for the simulation.
+            Default None (1 everywhere).
+        band: Which band to search in. Default 0 (lowest frequency).
+
+    Returns:
+        `(k, actual_frequency)`
+        The found k-vector and its frequency.
+    """
+
+    direction = numpy.array(direction) / norm(direction)
+
+    def get_f(k0_mag: float, band: int = 0):
+        k0 = direction * k0_mag
+        n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
+        f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
+        if solve_callback:
+            solve_callback(k0_mag, n, v, f)
+        return f
+
+    res = scipy.optimize.minimize_scalar(lambda x: abs(get_f(x, band) - frequency),
+                                         (k_min + k_max) / 2,
+                                         method='Bounded',
+                                         bounds=(k_min, k_max),
+                                         options={'xatol': abs(tolerance)})
+    return res.x * direction, res.fun + frequency
+
+
+def eigsolve(num_modes: int,
+             k0: numpy.ndarray,
+             G_matrix: numpy.ndarray,
+             epsilon: fdfield_t,
+             mu: fdfield_t = None,
+             tolerance: float = 1e-20,
+             max_iters: int = 10000,
+             reset_iters: int = 100,
+             ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+    """
+    Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
+     k0 of the specified structure.
+
+    Args:
+        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
+        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+        epsilon: Dielectric constant distribution for the simulation.
+                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+        mu: Magnetic permability distribution for the simulation.
+            Default `None` (1 everywhere).
+        tolerance: Solver stops when fractional change in the objective
+                   `trace(Z.H @ A @ Z @ inv(Z Z.H))` is smaller than the tolerance
+
+    Returns:
+        `(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the
+        vector `eigenvectors[i, :]`
+    """
+    h_size = 2 * epsilon[0].size
+
+    kmag = norm(G_matrix @ k0)
+
+    '''
+    Generate the operators
+    '''
+    mop = maxwell_operator(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
+    imop = inverse_maxwell_operator_approx(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
+
+    scipy_op = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=mop)
+    scipy_iop = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=imop)
+
+    y_shape = (h_size, num_modes)
+
+    prev_E = 0
+    d_scale = 1
+    prev_traceGtKG = 0
+    #prev_theta = 0.5
+    D = numpy.zeros(shape=y_shape, dtype=complex)
+
+    y0 = None
+    if y0 is None:
+        Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
+    else:
+        Z = y0
+
+    while True:
+        Z *= num_modes / norm(Z)
+        ZtZ = Z.conj().T @ Z
+        try:
+            U = numpy.linalg.inv(ZtZ)
+        except numpy.linalg.LinAlgError:
+            Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
+            continue
+
+        trace_U = real(trace(U))
+        if trace_U > 1e8 * num_modes:
+            Z = Z @ scipy.linalg.sqrtm(U).conj().T
+            prev_traceGtKG = 0
+            continue
+        break
+
+    for i in range(max_iters):
+        ZtZ = Z.conj().T @ Z
+        U = numpy.linalg.inv(ZtZ)
+        AZ = scipy_op @ Z
+        AZU = AZ @ U
+        ZtAZU = Z.conj().T @ AZU
+        E_signed = real(trace(ZtAZU))
+        sgn = numpy.sign(E_signed)
+        E = numpy.abs(E_signed)
+        G = (AZU - Z @ U @ ZtAZU) * sgn
+
+        if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7):
+            logger.info('Optimization succeded: {} - 5e-8 < {} * {} / 2'.format(abs(E - prev_E), tolerance, E + prev_E))
+            break
+
+        KG = scipy_iop @ G
+        traceGtKG = _rtrace_AtB(G, KG)
+
+        if prev_traceGtKG == 0 or i % reset_iters == 0:
+            logger.info('CG reset')
+            gamma = 0
+        else:
+            gamma = traceGtKG / prev_traceGtKG
+
+        D = gamma / d_scale * D + KG
+        d_scale = num_modes / norm(D)
+        D *= d_scale
+
+        ZtAZ = Z.conj().T @ AZ
+
+        AD = scipy_op @ D
+        DtD = D.conj().T @ D
+        DtAD = D.conj().T @ AD
+
+        symZtD = _symmetrize(Z.conj().T @ D)
+        symZtAD = _symmetrize(Z.conj().T @ AD)
+
+        Qi_memo = [None, None]
+        def Qi_func(theta):
+            nonlocal Qi_memo
+            if Qi_memo[0] == theta:
+                return Qi_memo[1]
+
+            c = numpy.cos(theta)
+            s = numpy.sin(theta)
+            Q = c*c * ZtZ + s*s * DtD + 2*s*c * symZtD
+            try:
+                Qi = numpy.linalg.inv(Q)
+            except numpy.linalg.LinAlgError:
+                logger.info('taylor Qi')
+                # if c or s small, taylor expand
+                if c < 1e-4 * s and c != 0:
+                    DtDi = numpy.linalg.inv(DtD)
+                    Qi = DtDi / (s*s) - 2*c/(s*s*s) * (DtDi @ (DtDi @ symZtD).conj().T)
+                elif s < 1e-4 * c and s != 0:
+                    ZtZi = numpy.linalg.inv(ZtZ)
+                    Qi = ZtZi / (c*c) - 2*s/(c*c*c) * (ZtZi @ (ZtZi @ symZtD).conj().T)
+                else:
+                    raise Exception('Inexplicable singularity in trace_func')
+            Qi_memo[0] = theta
+            Qi_memo[1] = Qi
+            return Qi
+
+        def trace_func(theta):
+            c = numpy.cos(theta)
+            s = numpy.sin(theta)
+            Qi = Qi_func(theta)
+            R = c*c * ZtAZ + s*s * DtAD + 2*s*c * symZtAD
+            trace = _rtrace_AtB(R, Qi)
+            return numpy.abs(trace)
+
+        '''
+        def trace_deriv(theta):
+            Qi = Qi_func(theta)
+            c2 = numpy.cos(2 * theta)
+            s2 = numpy.sin(2 * theta)
+            F = -0.5*s2 *  (ZtAZ - DtAD) + c2 * symZtAD
+            trace_deriv = _rtrace_AtB(Qi, F)
+
+            G = Qi @ F.conj().T @ Qi.conj().T
+            H = -0.5*s2 * (ZtZ - DtD) + c2 * symZtD
+            trace_deriv -= _rtrace_AtB(G, H)
+
+            trace_deriv *= 2
+            return trace_deriv * sgn
+
+        U_sZtD = U @ symZtD
+
+        dE = 2.0 * (_rtrace_AtB(U, symZtAD) -
+                    _rtrace_AtB(ZtAZU, U_sZtD))
+
+        d2E = 2 * (_rtrace_AtB(U, DtAD) -
+                   _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD)) -
+               4 * _rtrace_AtB(U, symZtAD @ U_sZtD))
+
+        # Newton-Raphson to find a root of the first derivative:
+        theta = -dE/d2E
+
+        if d2E < 0 or abs(theta) >= pi:
+            theta = -abs(prev_theta) * numpy.sign(dE)
+
+        # theta, new_E, new_dE = linmin(theta, E, dE, 0.1, min(tolerance, 1e-6), 1e-14, 0, -numpy.sign(dE) * K_PI, trace_func)
+        theta, n, _, new_E, _, _new_dE = scipy.optimize.line_search(trace_func, trace_deriv, xk=theta, pk=numpy.ones((1,1)), gfk=dE, old_fval=E, c1=min(tolerance, 1e-6), c2=0.1, amax=pi)
+        '''
+        result = scipy.optimize.minimize_scalar(trace_func, bounds=(0, pi), tol=tolerance)
+        new_E = result.fun
+        theta = result.x
+
+        improvement = numpy.abs(E - new_E) * 2 / numpy.abs(E + new_E)
+        logger.info('linmin improvement {}'.format(improvement))
+        Z *= numpy.cos(theta)
+        Z += D * numpy.sin(theta)
+
+        prev_traceGtKG = traceGtKG
+        #prev_theta = theta
+        prev_E = E
+
+    '''
+    Recover eigenvectors from Z
+    '''
+    U = numpy.linalg.inv(Z.conj().T @ Z)
+    Y = Z @ scipy.linalg.sqrtm(U)
+    W = Y.conj().T @ (scipy_op @ Y)
+
+    eigvals, W_eigvecs = numpy.linalg.eig(W)
+    eigvecs = Y @ W_eigvecs
+
+    for i in range(len(eigvals)):
+        v = eigvecs[:, i]
+        n = eigvals[i]
+        v /= norm(v)
+        eigness = norm(scipy_op @ v - (v.conj() @ (scipy_op @ v)) * v )
+        f = numpy.sqrt(-numpy.real(n))
+        df = numpy.sqrt(-numpy.real(n + eigness))
+        neff_err = kmag * (1/df - 1/f)
+        logger.info('eigness {}: {}\n neff_err: {}'.format(i, eigness, neff_err))
+
+    order = numpy.argsort(numpy.abs(eigvals))
+    return eigvals[order], eigvecs.T[order]
+
+'''
+def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_tol=1e-14, x_min=0, linmin_func):
+    if df0 > 0:
+        x0, f0, df0 = linmin(-x_guess, f0, -df0, -x_max, f_tol, df_tol, x_tol, -x_min, lambda q, dq: -linmin_func(q, dq))
+        return -x0, f0, -df0
+    elif df0 == 0:
+        return 0, f0, df0
+    else:
+        x = x_guess
+        fx = f0
+        dfx = df0
+
+        isave = numpy.zeros((2,), numpy.intc)
+        dsave = numpy.zeros((13,), float)
+
+        x, fx, dfx, task = minpack2.dsrch(x, fx, dfx, f_tol, df_tol, x_tol, task,
+                                          x_min, x_max, isave, dsave)
+        for i in range(int(1e6)):
+            if task != 'F':
+                logging.info('search converged in {} iterations'.format(i))
+                break
+            fx = f(x, dfx)
+            x, fx, dfx, task = minpack2.dsrch(x, fx, dfx, f_tol, df_tol, x_tol, task,
+                                              x_min, x_max, isave, dsave)
+
+        return x, fx, dfx
+'''
+
+def _rtrace_AtB(A, B):
+    return real(numpy.sum(A.conj() * B))
+
+def _symmetrize(A):
+    return (A + A.conj().T) * 0.5
+

meanas/fdfd/farfield.py

@@ -1,45 +1,51 @@
 """
 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
 
+from ..fdmath import fdfield_t
 
-def near_to_farfield(E_near: List[numpy.ndarray],
-                     H_near: List[numpy.ndarray],
+
+def near_to_farfield(E_near: fdfield_t,
+                     H_near: fdfield_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.
 
     The input fields should be complex phasors.
 
-    :param E_near: List of 2 ndarrays containing the 2D phasor field slices for the transverse
-        E fields (e.g. [Ex, Ey] for calculating the farfield toward the z-direction).
-    :param H_near: List of 2 ndarrays containing the 2D phasor field slices for the transverse
-        H fields (e.g. [Hx, hy] for calculating the farfield towrad the z-direction).
-    :param dx: Cell size along x-dimension, in units of wavelength.
-    :param dy: Cell size along y-dimension, in units of wavelength.
-    :param padded_size: Shape of the output. A single integer `n` will be expanded to `(n, n)`.
-        Powers of 2 are most efficient for FFT computation.
-        Default is the smallest power of 2 larger than the input, for each axis.
-    :returns: Dict with keys
-        'E_far': Normalized E-field farfield; multiply by
-            (i k exp(-i k r) / (4 pi r)) to get the actual field value.
-        'H_far': Normalized H-field farfield; multiply by
-            (i k exp(-i k r) / (4 pi r)) to get the actual field value.
-        'kx', 'ky': Wavevector values corresponding to the x- and y- axes in E_far and H_far,
-            normalized to wavelength (dimensionless).
-        'dkx', 'dky': step size for kx and ky, normalized to wavelength.
-        'theta': arctan2(ky, kx) corresponding to each (kx, ky).
-            This is the angle in the x-y plane, counterclockwise from above, starting from +x.
-        'phi': arccos(kz / k) corresponding to each (kx, ky).
-            This is the angle away from +z.
+    Args:
+        E_near: List of 2 ndarrays containing the 2D phasor field slices for the transverse
+                E fields (e.g. [Ex, Ey] for calculating the farfield toward the z-direction).
+        H_near: List of 2 ndarrays containing the 2D phasor field slices for the transverse
+                H fields (e.g. [Hx, hy] for calculating the farfield towrad the z-direction).
+        dx: Cell size along x-dimension, in units of wavelength.
+        dy: Cell size along y-dimension, in units of wavelength.
+        padded_size: Shape of the output. A single integer `n` will be expanded to `(n, n)`.
+                     Powers of 2 are most efficient for FFT computation.
+                     Default is the smallest power of 2 larger than the input, for each axis.
+
+    Returns:
+        Dict with keys
+
+        -   `E_far`: Normalized E-field farfield; multiply by
+                (i k exp(-i k r) / (4 pi r)) to get the actual field value.
+        -   `H_far`: Normalized H-field farfield; multiply by
+                (i k exp(-i k r) / (4 pi r)) to get the actual field value.
+        -   `kx`, `ky`: Wavevector values corresponding to the x- and y- axes in E_far and H_far,
+                normalized to wavelength (dimensionless).
+        -   `dkx`, `dky`: step size for kx and ky, normalized to wavelength.
+        -   `theta`: arctan2(ky, kx) corresponding to each (kx, ky).
+                This is the angle in the x-y plane, counterclockwise from above, starting from +x.
+        -   `phi`: arccos(kz / k) corresponding to each (kx, ky).
+                This is the angle away from +z.
     """
 
     if not len(E_near) == 2:
@@ -115,35 +121,39 @@ def near_to_farfield(E_near: List[numpy.ndarray],
 
 
 
-def far_to_nearfield(E_far: List[numpy.ndarray],
-                     H_far: List[numpy.ndarray],
+def far_to_nearfield(E_far: fdfield_t,
+                     H_far: fdfield_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.
 
     The input fields should be complex phasors.
 
-    :param E_far: List of 2 ndarrays containing the 2D phasor field slices for the transverse
-        E fields (e.g. [Ex, Ey] for calculating the nearfield toward the z-direction).
-        Fields should be normalized so that
-            E_far = E_far_actual / (i k exp(-i k r) / (4 pi r))
-    :param H_far: List of 2 ndarrays containing the 2D phasor field slices for the transverse
-        H fields (e.g. [Hx, hy] for calculating the nearfield toward the z-direction).
-        Fields should be normalized so that
-            H_far = H_far_actual / (i k exp(-i k r) / (4 pi r))
-    :param dkx: kx discretization, in units of wavelength.
-    :param dky: ky discretization, in units of wavelength.
-    :param padded_size: Shape of the output. A single integer `n` will be expanded to `(n, n)`.
-        Powers of 2 are most efficient for FFT computation.
-        Default is the smallest power of 2 larger than the input, for each axis.
-    :returns: Dict with keys
-        'E': E-field nearfield
-        'H': H-field nearfield
-        'dx', 'dy': spatial discretization, normalized to wavelength (dimensionless)
+    Args:
+        E_far: List of 2 ndarrays containing the 2D phasor field slices for the transverse
+                E fields (e.g. [Ex, Ey] for calculating the nearfield toward the z-direction).
+                Fields should be normalized so that
+                E_far = E_far_actual / (i k exp(-i k r) / (4 pi r))
+        H_far: List of 2 ndarrays containing the 2D phasor field slices for the transverse
+                H fields (e.g. [Hx, hy] for calculating the nearfield toward the z-direction).
+                Fields should be normalized so that
+                H_far = H_far_actual / (i k exp(-i k r) / (4 pi r))
+        dkx: kx discretization, in units of wavelength.
+        dky: ky discretization, in units of wavelength.
+        padded_size: Shape of the output. A single integer `n` will be expanded to `(n, n)`.
+                     Powers of 2 are most efficient for FFT computation.
+                     Default is the smallest power of 2 larger than the input, for each axis.
+
+    Returns:
+        Dict with keys
+
+        -   `E`: E-field nearfield
+        -   `H`: H-field nearfield
+        -   `dx`, `dy`: spatial discretization, normalized to wavelength (dimensionless)
     """
 
     if not len(E_far) == 2:

meanas/fdfd/functional.py

@@ -0,0 +1,219 @@
+"""
+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 `fdfield_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, Tuple
+import numpy
+
+from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
+from ..fdmath.functional import curl_forward, curl_back
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def e_full(omega: complex,
+           dxes: dx_lists_t,
+           epsilon: fdfield_t,
+           mu: fdfield_t = None
+           ) -> fdfield_updater_t:
+    """
+    Wave operator for use with E-field. See `operators.e_full` for details.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+        epsilon: Dielectric constant
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Return:
+        Function `f` implementing the wave operator
+        `f(E)` -> `-i * omega * J`
+    """
+    ch = curl_back(dxes[1])
+    ce = curl_forward(dxes[0])
+
+    def op_1(e):
+        curls = ch(ce(e))
+        return curls - omega ** 2 * epsilon * e
+
+    def op_mu(e):
+        curls = ch(mu * ce(e))
+        return curls - omega ** 2 * epsilon * e
+
+    if numpy.any(numpy.equal(mu, None)):
+        return op_1
+    else:
+        return op_mu
+
+
+def eh_full(omega: complex,
+            dxes: dx_lists_t,
+            epsilon: fdfield_t,
+            mu: fdfield_t = None
+            ) -> Callable[[fdfield_t, fdfield_t], Tuple[fdfield_t, fdfield_t]]:
+    """
+    Wave operator for full (both E and H) field representation.
+    See `operators.eh_full`.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+        epsilon: Dielectric constant
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Function `f` implementing the wave operator
+        `f(E, H)` -> `(J, -M)`
+    """
+    ch = curl_back(dxes[1])
+    ce = curl_forward(dxes[0])
+
+    def op_1(e, h):
+        return (ch(h) - 1j * omega * epsilon * e,
+                ce(e) + 1j * omega * h)
+
+    def op_mu(e, h):
+        return (ch(h) - 1j * omega * epsilon * e,
+                ce(e) + 1j * omega * mu * h)
+
+    if numpy.any(numpy.equal(mu, None)):
+        return op_1
+    else:
+        return op_mu
+
+
+def e2h(omega: complex,
+        dxes: dx_lists_t,
+        mu: fdfield_t = None,
+        ) -> fdfield_updater_t:
+    """
+    Utility operator for converting the `E` field into the `H` field.
+    For use with `e_full` -- assumes that there is no magnetic current `M`.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Return:
+        Function `f` for converting `E` to `H`,
+        `f(E)` -> `H`
+    """
+    ce = curl_forward(dxes[0])
+
+    def e2h_1_1(e):
+        return ce(e) / (-1j * omega)
+
+    def e2h_mu(e):
+        return ce(e) / (-1j * omega * mu)
+
+    if numpy.any(numpy.equal(mu, None)):
+        return e2h_1_1
+    else:
+        return e2h_mu
+
+
+def m2j(omega: complex,
+        dxes: dx_lists_t,
+        mu: fdfield_t = None,
+        ) -> fdfield_updater_t:
+    """
+    Utility operator for converting magnetic current `M` distribution
+    into equivalent electric current distribution `J`.
+    For use with e.g. `e_full`.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+        mu: Magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Function `f` for converting `M` to `J`,
+        `f(M)` -> `J`
+    """
+    ch = curl_back(dxes[1])
+
+    def m2j_mu(m):
+        J = ch(m / mu) / (-1j * omega)
+        return J
+
+    def m2j_1(m):
+        J = ch(m) / (-1j * omega)
+        return J
+
+    if numpy.any(numpy.equal(mu, None)):
+        return m2j_1
+    else:
+        return m2j_mu
+
+
+def e_tfsf_source(TF_region: fdfield_t,
+                  omega: complex,
+                  dxes: dx_lists_t,
+                  epsilon: fdfield_t,
+                  mu: fdfield_t = None,
+                  ) -> fdfield_updater_t:
+    """
+    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.fdmath.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)
+
+    def op(e):
+        neg_iwj = A(TF_region * e) - TF_region * A(e)
+        return neg_iwj / (-1j * omega)
+
+
+def poynting_e_cross_h(dxes: dx_lists_t) -> Callable[[fdfield_t, fdfield_t], fdfield_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.fdmath.types`
+
+    Returns:
+        Function `f` that returns E x H as required for the poynting vector.
+    """
+    def exh(e: fdfield_t, h: fdfield_t):
+        s = numpy.empty_like(e)
+        ex = e[0] * dxes[0][0][:, None, None]
+        ey = e[1] * dxes[0][1][None, :, None]
+        ez = e[2] * dxes[0][2][None, None, :]
+        hx = h[0] * dxes[1][0][:, None, None]
+        hy = h[1] * dxes[1][1][None, :, None]
+        hz = h[2] * dxes[1][2][None, None, :]
+        s[0] = numpy.roll(ey, -1, axis=0) * hz - numpy.roll(ez, -1, axis=0) * hy
+        s[1] = numpy.roll(ez, -1, axis=1) * hx - numpy.roll(ex, -1, axis=1) * hz
+        s[2] = numpy.roll(ex, -1, axis=2) * hy - numpy.roll(ey, -1, axis=2) * hx
+        return s
+    return exh

meanas/fdfd/operators.py

@@ -0,0 +1,441 @@
+"""
+Sparse matrix operators for use with electromagnetic wave equations.
+
+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 `meanas.unvec()` functions.
+
+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).
+
+Many of these functions require a `dxes` parameter, of type `dx_lists_t`; see
+the `meanas.fdmath.types` submodule for details.
+
+
+The following operators are included:
+
+- E-only wave operator
+- H-only wave operator
+- EH wave operator
+- Curl for use with E, H fields
+- E to H conversion
+- M to J conversion
+- Poynting cross products
+- Circular shifts
+- Discrete derivatives
+- Averaging operators
+- Cross product matrices
+"""
+
+from typing import List, Tuple
+import numpy
+import scipy.sparse as sparse
+
+from ..fdmath import vec, dx_lists_t, vfdfield_t
+from ..fdmath.operators import shift_with_mirror, rotation, curl_forward, curl_back
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def e_full(omega: complex,
+           dxes: dx_lists_t,
+           epsilon: vfdfield_t,
+           mu: vfdfield_t = None,
+           pec: vfdfield_t = None,
+           pmc: vfdfield_t = None,
+           ) -> sparse.spmatrix:
+    """
+    Wave operator
+     $$ \\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 preconditioners from `e_full_preconditioners()`.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.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 (i.e. `pec.size == epsilon.size`)
+        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`)
+
+    Returns:
+        Sparse matrix containing the wave operator.
+    """
+    ch = curl_back(dxes[1])
+    ce = curl_forward(dxes[0])
+
+    if numpy.any(numpy.equal(pec, None)):
+        pe = sparse.eye(epsilon.size)
+    else:
+        pe = sparse.diags(numpy.where(pec, 0, 1))     # Set pe to (not PEC)
+
+    if numpy.any(numpy.equal(pmc, None)):
+        pm = sparse.eye(epsilon.size)
+    else:
+        pm = sparse.diags(numpy.where(pmc, 0, 1))     # set pm to (not PMC)
+
+    e = sparse.diags(epsilon)
+    if numpy.any(numpy.equal(mu, None)):
+        m_div = sparse.eye(epsilon.size)
+    else:
+        m_div = sparse.diags(1 / mu)
+
+    op = pe @ (ch @ pm @ m_div @ ce - omega**2 * e) @ pe
+    return op
+
+
+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.
+
+    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`
+
+    Args:
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.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, :],
+                 dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[0][2][None, None, :]]
+
+    p_vector = numpy.sqrt(vec(p_squared))
+    P_left = sparse.diags(p_vector)
+    P_right = sparse.diags(1 / p_vector)
+    return P_left, P_right
+
+
+def h_full(omega: complex,
+           dxes: dx_lists_t,
+           epsilon: vfdfield_t,
+           mu: vfdfield_t = None,
+           pec: vfdfield_t = None,
+           pmc: vfdfield_t = None,
+           ) -> sparse.spmatrix:
+    """
+    Wave operator
+     $$ \\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu $$
+
+        del x (1/epsilon * del x) - omega**2 * mu
+
+     for use with the H-field, with wave equation
+     $$ (\\nabla \\times (\\frac{1}{\\epsilon} \\nabla \\times) - \\omega^2 \\mu) E = \\imath \\omega M $$
+
+        (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.fdmath.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 (i.e. `pec.size == epsilon.size`)
+        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`)
+
+    Returns:
+        Sparse matrix containing the wave operator.
+    """
+    ch = curl_back(dxes[1])
+    ce = curl_forward(dxes[0])
+
+    if numpy.any(numpy.equal(pec, None)):
+        pe = sparse.eye(epsilon.size)
+    else:
+        pe = sparse.diags(numpy.where(pec, 0, 1))    # set pe to (not PEC)
+
+    if numpy.any(numpy.equal(pmc, None)):
+        pm = sparse.eye(epsilon.size)
+    else:
+        pm = sparse.diags(numpy.where(pmc, 0, 1))    # Set pe to (not PMC)
+
+    e_div = sparse.diags(1 / epsilon)
+    if mu is None:
+        m = sparse.eye(epsilon.size)
+    else:
+        m = sparse.diags(mu)
+
+    A = pm @ (ce @ pe @ e_div @ ch - omega**2 * m) @ pm
+    return A
+
+
+def eh_full(omega: complex,
+            dxes: dx_lists_t,
+            epsilon: vfdfield_t,
+            mu: vfdfield_t = None,
+            pec: vfdfield_t = None,
+            pmc: vfdfield_t = None
+            ) -> sparse.spmatrix:
+    """
+    Wave operator for `[E, H]` field representation. This operator implements Maxwell's
+     equations without cancelling out either E or H. The operator is
+    $$  \\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 `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} $$
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.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 (i.e. `pec.size == epsilon.size`)
+        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`)
+
+    Returns:
+        Sparse matrix containing the wave operator.
+    """
+    if numpy.any(numpy.equal(pec, None)):
+        pe = sparse.eye(epsilon.size)
+    else:
+        pe = sparse.diags(numpy.where(pec, 0, 1))    # set pe to (not PEC)
+
+    if numpy.any(numpy.equal(pmc, None)):
+        pm = sparse.eye(epsilon.size)
+    else:
+        pm = sparse.diags(numpy.where(pmc, 0, 1))    # set pm to (not PMC)
+
+    iwe = pe @ (1j * omega * sparse.diags(epsilon)) @ pe
+    iwm = 1j * omega
+    if not numpy.any(numpy.equal(mu, None)):
+        iwm *= sparse.diags(mu)
+    iwm = pm @ iwm @ pm
+
+    A1 = pe @ curl_back(dxes[1]) @ pm
+    A2 = pm @ curl_forward(dxes[0]) @ pe
+
+    A = sparse.bmat([[-iwe, A1],
+                     [A2,  iwm]])
+    return A
+
+
+def e2h(omega: complex,
+        dxes: dx_lists_t,
+        mu: vfdfield_t = None,
+        pmc: vfdfield_t = None,
+        ) -> 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.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+        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 (i.e. `pmc.size == epsilon.size`)
+
+    Returns:
+        Sparse matrix for converting E to H.
+    """
+    op = curl_forward(dxes[0]) / (-1j * omega)
+
+    if not numpy.any(numpy.equal(mu, None)):
+        op = sparse.diags(1 / mu) @ op
+
+    if not numpy.any(numpy.equal(pmc, None)):
+        op = sparse.diags(numpy.where(pmc, 0, 1)) @ op
+
+    return op
+
+
+def m2j(omega: complex,
+        dxes: dx_lists_t,
+        mu: vfdfield_t = None
+        ) -> sparse.spmatrix:
+    """
+    Operator for converting a magnetic current M into an electric current J.
+    For use with eg. `e_full()`.
+
+    Args:
+        omega: Angular frequency of the simulation
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+        mu: Vectorized magnetic permeability (default 1 everywhere)
+
+    Returns:
+        Sparse matrix for converting M to J.
+    """
+    op = curl_back(dxes[1]) / (1j * omega)
+
+    if not numpy.any(numpy.equal(mu, None)):
+        op = op @ sparse.diags(1 / mu)
+
+    return op
+
+
+def poynting_e_cross(e: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Operator for computing the Poynting vector, containing the
+    (E x) portion of the Poynting vector.
+
+    Args:
+        e: Vectorized E-field for the ExH cross product
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+
+    Returns:
+        Sparse matrix containing (E x) portion of Poynting cross product.
+    """
+    shape = [len(dx) for dx in dxes[0]]
+
+    fx, fy, fz = [rotation(i, shape, 1) for i in range(3)]
+
+    dxag = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
+    dxbg = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
+    Ex, Ey, Ez = [ei * da for ei, da in zip(numpy.split(e, 3), dxag)]
+
+    block_diags = [[ None,     fx @ -Ez, fx @  Ey],
+                   [ fy @  Ez, None,     fy @ -Ex],
+                   [ fz @ -Ey, fz @  Ex, None]]
+    block_matrix = sparse.bmat([[sparse.diags(x) if x is not None else None for x in row]
+                                 for row in block_diags])
+    P = block_matrix @ sparse.diags(numpy.concatenate(dxag))
+    return P
+
+
+def poynting_h_cross(h: vfdfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector.
+
+    Args:
+        h: Vectorized H-field for the HxE cross product
+        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+
+    Returns:
+        Sparse matrix containing (H x) portion of Poynting cross product.
+    """
+    shape = [len(dx) for dx in dxes[0]]
+
+    fx, fy, fz = [rotation(i, shape, 1) for i in range(3)]
+
+    dxag = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
+    dxbg = [dx.ravel(order='C') for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
+    Hx, Hy, Hz = [sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg)]
+
+    P = (sparse.bmat(
+          [[ None,    -Hz @ fx,   Hy @ fx],
+           [ Hz @ fy,  None,     -Hx @ fy],
+           [-Hy @ fz,  Hx @ fz,   None]])
+          @ sparse.diags(numpy.concatenate(dxag)))
+    return P
+
+
+def e_tfsf_source(TF_region: vfdfield_t,
+                  omega: complex,
+                  dxes: dx_lists_t,
+                  epsilon: vfdfield_t,
+                  mu: vfdfield_t = None,
+                  ) -> sparse.spmatrix:
+    """
+    Operator that turns a desired E-field distribution into a
+     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.fdmath.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)
+    Q = sparse.diags(TF_region)
+    return (A @ Q - Q @ A) / (-1j * omega)
+
+
+def e_boundary_source(mask: vfdfield_t,
+                      omega: complex,
+                      dxes: dx_lists_t,
+                      epsilon: vfdfield_t,
+                      mu: vfdfield_t = None,
+                      periodic_mask_edges: bool = False,
+                      ) -> sparse.spmatrix:
+    """
+    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.
+
+    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.fdmath.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)
+
+    shape = [len(dxe) for dxe in dxes[0]]
+    jmask = numpy.zeros_like(mask, dtype=bool)
+
+    if periodic_mask_edges:
+        shift = lambda axis, polarity: rotation(axis=axis, shape=shape, shift_distance=polarity)
+    else:
+        shift = lambda axis, polarity: shift_with_mirror(axis=axis, shape=shape, shift_distance=polarity)
+
+    for axis in (0, 1, 2):
+        if shape[axis] == 1:
+            continue
+        for polarity in (-1, +1):
+            r = shift(axis, polarity) - sparse.eye(numpy.prod(shape)) # shifted minus original
+            r3 = sparse.block_diag((r, r, r))
+            jmask = numpy.logical_or(jmask, numpy.abs(r3 @ mask))
+
+#    jmask = ((numpy.roll(mask, -1, axis=0) != mask) |
+#             (numpy.roll(mask, +1, axis=0) != mask) |
+#             (numpy.roll(mask, -1, axis=1) != mask) |
+#             (numpy.roll(mask, +1, axis=1) != mask) |
+#             (numpy.roll(mask, -1, axis=2) != mask) |
+#             (numpy.roll(mask, +1, axis=2) != mask))
+
+    return sparse.diags(jmask.astype(int)) @ full
+

meanas/fdfd/scpml.py

@@ -1,29 +1,35 @@
 """
-Functions for creating stretched coordinate PMLs.
+Functions for creating stretched coordinate perfectly matched layer (PML) absorbers.
 """
 
 from typing import List, Callable
 import numpy
 
+from ..fdmath import dx_lists_t
+
+
 __author__ = 'Jan Petykiewicz'
 
 
-dx_lists_t = List[List[numpy.ndarray]]
-s_function_type = Callable[[float], float]
+s_function_t = Callable[[float], float]
+"""Typedef for s-functions, see `prepare_s_function()`"""
 
 
 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
-    :param m: Polynomial order for the PML (imaginary part increases as distance ** m)
-    :return: 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.
+    Args:
+        ln_R: Natural logarithm of the desired reflectance
+        m: Polynomial order for the PML (imaginary part increases as distance ** m)
+
+    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:
         s_max = (m + 1) * ln_R / 2  # / 2 because we assume periodic boundaries
@@ -35,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)
-    :param 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.
-    :param omega: Angular frequency for the simulation
-    :param 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
-        customization of pml parameters.
-        Default uses prepare_s_function() with no parameters.
-    :return: Complex cell widths (dx_lists)
+    Args:
+        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.
+        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.
+        s_function: created by `prepare_s_function(...)`, allowing customization of pml parameters.
+                    Default uses `prepare_s_function()` with no parameters.
+
+    Returns:
+        Complex cell widths (dx_lists_t) as discussed in `meanas.fdmath.types`.
     """
     if s_function is None:
         s_function = prepare_s_function()
@@ -87,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
-        :param axis: axis to stretch (0=x, 1=y, 2=z)
-        :param polarity: direction to stretch (-1 for -ve, +1 for +ve)
-        :param omega: Angular frequency for the simulation
-        :param 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)
-        :param s_function: s_function created by prepare_s_function(...), allowing customization
-            of pml parameters. Default uses prepare_s_function() with no parameters.
-        :return: Complex cell widths
+        Args:
+            dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
+            axis: axis to stretch (0=x, 1=y, 2=z)
+            polarity: direction to stretch (-1 for -ve, +1 for +ve)
+            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.
+            thickness: number of cells to use for pml (default 10)
+            s_function: Created by `prepare_s_function(...)`, allowing customization
+                        of pml parameters. Default uses `prepare_s_function()` with no parameters.
+
+        Returns:
+            Complex cell widths (dx_lists_t) as discussed in `meanas.fdmath.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()
@@ -146,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

@@ -0,0 +1,126 @@
+"""
+Solvers and solver interface for FDFD problems.
+"""
+
+from typing import List, Callable, Dict, Any
+import logging
+
+import numpy
+from numpy.linalg import norm
+import scipy.sparse.linalg
+
+from ..fdmath import dx_lists_t, vfdfield_t
+from . import operators
+
+
+logger = logging.getLogger(__name__)
+
+
+def _scipy_qmr(A: scipy.sparse.csr_matrix,
+               b: numpy.ndarray,
+               **kwargs
+               ) -> numpy.ndarray:
+    """
+    Wrapper for scipy.sparse.linalg.qmr
+
+    Args:
+        A: Sparse matrix
+        b: Right-hand-side vector
+        kwargs: Passed as **kwargs to the wrapped function
+
+    Returns:
+        Guess for solution (returned even if didn't converge)
+    """
+
+    '''
+    Report on our progress
+    '''
+    iter = 0
+
+    def log_residual(xk):
+        nonlocal iter
+        iter += 1
+        if iter % 100 == 0:
+            logger.info('Solver residual at iteration {} : {}'.format(iter, norm(A @ xk - b)))
+
+    if 'callback' in kwargs:
+        def augmented_callback(xk):
+            log_residual(xk)
+            kwargs['callback'](xk)
+
+        kwargs['callback'] = augmented_callback
+    else:
+        kwargs['callback'] = log_residual
+
+    '''
+    Run the actual solve
+    '''
+
+    x, _ = scipy.sparse.linalg.qmr(A, b, **kwargs)
+    return x
+
+
+def generic(omega: complex,
+            dxes: dx_lists_t,
+            J: vfdfield_t,
+            epsilon: vfdfield_t,
+            mu: vfdfield_t = None,
+            pec: vfdfield_t = None,
+            pmc: vfdfield_t = None,
+            adjoint: bool = False,
+            matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr,
+            matrix_solver_opts: Dict[str, Any] = None,
+            ) -> vfdfield_t:
+    """
+    Conjugate gradient FDFD solver using CSR sparse matrices.
+
+    All ndarray arguments should be 1D arrays, as returned by `meanas.vec()`.
+
+    Args:
+        omega: Complex frequency to solve at.
+        dxes: `[[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]]` (complex cell sizes) as
+            discussed in `meanas.fdmath.types`
+        J: Electric current distribution (at E-field locations)
+        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)
+        pmc: Perfect magnetic conductor distribution
+             (at H-field locations; non-zero value indicates PMC is present)
+        adjoint: If true, solves the adjoint problem.
+        matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`,
+                where `A`: `scipy.sparse.csr_matrix`;
+                      `b`: `numpy.ndarray`;
+                      `x`: `numpy.ndarray`;
+                Default is a wrapped version of `scipy.sparse.linalg.qmr()`
+                 which doesn't return convergence info and logs the residual
+                 every 100 iterations.
+        matrix_solver_opts: Passed as kwargs to `matrix_solver(...)`
+
+    Returns:
+        E-field which solves the system.
+    """
+
+    if matrix_solver_opts is None:
+        matrix_solver_opts = dict()
+
+    b0 = -1j * omega * J
+    A0 = operators.e_full(omega, dxes, epsilon=epsilon, mu=mu, pec=pec, pmc=pmc)
+
+    Pl, Pr = operators.e_full_preconditioners(dxes)
+
+    if adjoint:
+        A = (Pl @ A0 @ Pr).H
+        b = Pr.H @ b0
+    else:
+        A = Pl @ A0 @ Pr
+        b = Pl @ b0
+
+    x = matrix_solver(A.tocsr(), b, **matrix_solver_opts)
+
+    if adjoint:
+        x0 = Pl.H @ x
+    else:
+        x0 = Pr @ x
+
+    return x0

meanas/fdfd/waveguide_2d.py

@@ -0,0 +1,737 @@
+"""
+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], ...]]]`).
+
+
+===============
+
+Consider Maxwell's equations in continuous space, in the frequency domain. Assuming
+a structure with some (x, y) cross-section extending uniformly into the z dimension,
+with a diagonal $\\epsilon$ tensor, we have
+
+$$
+\\begin{aligned}
+\\nabla \\times \\vec{E}(x, y, z) &= -\\imath \\omega \\mu \\vec{H} \\\\
+\\nabla \\times \\vec{H}(x, y, z) &=  \\imath \\omega \\epsilon \\vec{E} \\\\
+\\vec{E}(x,y,z) = (\\vec{E}_t(x, y) + E_z(x, y)\\vec{z}) e^{-\\gamma z} \\\\
+\\vec{H}(x,y,z) = (\\vec{H}_t(x, y) + H_z(x, y)\\vec{z}) e^{-\\gamma z} \\\\
+\\end{aligned}
+$$
+
+Expanding the first two equations into vector components, we get
+
+$$
+\\begin{aligned}
+-\\imath \\omega \\mu_{xx} H_x &= \\partial_y E_z - \\partial_z E_y \\\\
+-\\imath \\omega \\mu_{yy} H_y &= \\partial_z E_x - \\partial_x E_z \\\\
+-\\imath \\omega \\mu_{zz} H_z &= \\partial_x E_y - \\partial_y E_x \\\\
+\\imath \\omega \\epsilon_{xx} E_x &= \\partial_y H_z - \\partial_z H_y \\\\
+\\imath \\omega \\epsilon_{yy} E_y &= \\partial_z H_x - \\partial_x H_z \\\\
+\\imath \\omega \\epsilon_{zz} E_z &= \\partial_x H_y - \\partial_y H_x \\\\
+\\end{aligned}
+$$
+
+Substituting in our expressions for $\\vec{E}$, $\\vec{H}$ and discretizing:
+
+$$
+\\begin{aligned}
+-\\imath \\omega \\mu_{xx} H_x &= \\tilde{\\partial}_y E_z + \\gamma E_y \\\\
+-\\imath \\omega \\mu_{yy} H_y &= -\\gamma E_x - \\tilde{\\partial}_x E_z \\\\
+-\\imath \\omega \\mu_{zz} H_z &= \\tilde{\\partial}_x E_y - \\tilde{\\partial}_y E_x \\\\
+\\imath \\omega \\epsilon_{xx} E_x &= \\hat{\\partial}_y H_z + \\gamma H_y \\\\
+\\imath \\omega \\epsilon_{yy} E_y &= -\\gamma H_x - \\hat{\\partial}_x H_z \\\\
+\\imath \\omega \\epsilon_{zz} E_z &= \\hat{\\partial}_x H_y - \\hat{\\partial}_y H_x \\\\
+\\end{aligned}
+$$
+
+Rewrite the last three equations as
+$$
+\\begin{aligned}
+\\gamma H_y &=  \\imath \\omega \\epsilon_{xx} E_x - \\hat{\\partial}_y H_z \\\\
+\\gamma H_x &= -\\imath \\omega \\epsilon_{yy} E_y - \\hat{\\partial}_x H_z \\\\
+\\imath \\omega E_z &= \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_x H_y - \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_y H_x \\\\
+\\end{aligned}
+$$
+
+Now apply $\\gamma \\tilde{\\partial}_x$ to the last equation,
+then substitute in for $\\gamma H_x$ and $\\gamma H_y$:
+
+$$
+\\begin{aligned}
+\\gamma \\tilde{\\partial}_x \\imath \\omega E_z &= \\gamma \\tilde{\\partial}_x \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_x H_y
+                                                - \\gamma \\tilde{\\partial}_x \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_y H_x \\\\
+        &= \\tilde{\\partial}_x \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_x ( \\imath \\omega \\epsilon_{xx} E_x - \\hat{\\partial}_y H_z)
+         - \\tilde{\\partial}_x \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_y (-\\imath \\omega \\epsilon_{yy} E_y - \\hat{\\partial}_x H_z)  \\\\
+        &= \\tilde{\\partial}_x \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_x ( \\imath \\omega \\epsilon_{xx} E_x)
+         - \\tilde{\\partial}_x \\frac{1}{\\epsilon_{zz}} \\hat{\\partial}_y (-\\imath \\omega \\epsilon_{yy} E_y)  \\\\
+\\gamma \\tilde{\\partial}_x E_z &= \\tilde{\\partial}_x \\frac{1}{\\epsilon_zz} \\hat{\\partial}_x (\\epsilon_{xx} E_x)
+                                    \\tilde{\\partial}_x \\frac{1}{\\epsilon_zz} \\hat{\\partial}_y (\\epsilon_{yy} E_y) \\\\
+\\end{aligned}
+$$
+
+With a similar approach (but using $\\gamma \\tilde{\\partial}_y$ instead), we can get
+
+$$
+\\begin{aligned}
+\\gamma \\tilde{\\partial}_y E_z &= \\tilde{\\partial}_y \\frac{1}{\\epsilon_zz} \\hat{\\partial}_x (\\epsilon_{xx} E_x)
+                                    \\tilde{\\partial}_y \\frac{1}{\\epsilon_zz} \\hat{\\partial}_y (\\epsilon_{yy} E_y) \\\\
+\\end{aligned}
+$$
+
+We can combine this equation for $\\gamma \\tilde{\\partial}_y E_z$ with
+the unused $\\imath \\omega \\mu_{xx} H_z$ and $\\imath \\omega \\mu_{yy} H_y$ equations to get
+
+$$
+\\begin{aligned}
+-\\imath \\omega \\mu_{xx} \\gamma H_x &=  \\gamma^2 E_y + \\tilde{\\partial}_y (
+                                      \\tilde{\\partial}_x \\frac{1}{\\epsilon_zz} \\hat{\\partial}_x (\\epsilon_{xx} E_x)
+                                    + \\tilde{\\partial}_x \\frac{1}{\\epsilon_zz} \\hat{\\partial}_y (\\epsilon_{yy} E_y)
+                                    ) \\\\
+-\\imath \\omega \\mu_{yy} \\gamma H_y &= -\\gamma^2 E_x - \\tilde{\\partial}_x (
+                                      \\tilde{\\partial}_y \\frac{1}{\\epsilon_zz} \\hat{\\partial}_x (\\epsilon_{xx} E_x)
+                                    + \\tilde{\\partial}_y \\frac{1}{\\epsilon_zz} \\hat{\\partial}_y (\\epsilon_{yy} E_y)
+                                    )\\\\
+\\end{aligned}
+$$
+
+However, based on our rewritten equation for $\\gamma H_x$ and the so-far unused
+equation for $\\imath \\omega \\mu_{zz} H_z$ we can also write
+
+$$
+\\begin{aligned}
+-\\imath \\omega \\mu_{xx} (\\gamma H_x) &= -\\imath \\omega \\mu_{xx} (-\\imath \\omega \\epsilon_{yy} E_y - \\hat{\\partial}_x H_z) \\\\
+                                         &= -\\omega^2 \\mu_{xx} \\epsilon_{yy} E_y
+                                            -\\imath \\omega \\mu_{xx} \\hat{\\partial}_x (
+                                                \\frac{1}{-\\imath \\omega \\mu_{zz}} (\\tilde{\\partial}_x E_y - \\tilde{\\partial}_y E_x)) \\\\
+                                         &= -\\omega^2 \\mu_{xx} \\epsilon_{yy} E_y
+                                            +\\mu_{xx} \\hat{\\partial}_x \\frac{1}{\\mu_{zz}} (\\tilde{\\partial}_x E_y - \\tilde{\\partial}_y E_x) \\\\
+\\end{aligned}
+$$
+
+and, similarly,
+
+$$
+\\begin{aligned}
+-\\imath \\omega \\mu_{yy} (\\gamma H_y) &= -\\omega^2 \\mu_{yy} \\epsilon_{xx} E_x
+                                            +\\mu_{yy} \\hat{\\partial}_y \\frac{1}{\\mu_{zz}} (\\tilde{\\partial}_x E_y - \\tilde{\\partial}_y E_x) \\\\
+\\end{aligned}
+$$
+
+Using these, we can construct the eigenvalue problem
+$$ \\beta^2 \\begin{bmatrix} E_x \\\\
+                             E_y \\end{bmatrix} =
+    (\\omega^2 \\begin{bmatrix} \\mu_{yy} \\epsilon_{xx} & 0 \\\\
+                                                       0 & \\mu_{xx} \\epsilon_{yy} \\end{bmatrix} +
+                 \\begin{bmatrix} -\\mu_{yy} \\hat{\\partial}_y \\\\
+                                   \\mu_{xx} \\hat{\\partial}_x \\end{bmatrix} \\mu_{zz}^{-1}
+                 \\begin{bmatrix} -\\tilde{\\partial}_y & \\tilde{\\partial}_x \\end{bmatrix} +
+      \\begin{bmatrix} \\tilde{\\partial}_x \\\\
+                       \\tilde{\\partial}_y \\end{bmatrix} \\epsilon_{zz}^{-1}
+                 \\begin{bmatrix} \\hat{\\partial}_x \\epsilon_{xx} & \\hat{\\partial}_y \\epsilon_{yy} \\end{bmatrix})
+    \\begin{bmatrix} E_x \\\\
+                     E_y \\end{bmatrix}
+$$
+
+An equivalent eigenvalue problem can be formed using the $H_x$ and $H_y$ fields, if those are more convenient.
+
+Note that $E_z$ was never discretized, so $\\gamma$ and $\\beta$ will need adjustment
+to account for numerical dispersion if the result is introduced into a space with a discretized z-axis.
+
+
+"""
+# TODO update module docs
+
+from typing import List, Tuple
+import numpy
+from numpy.linalg import norm
+import scipy.sparse as sparse
+
+from ..fdmath.operators import deriv_forward, deriv_back, curl_forward, curl_back, cross
+from ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_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: vfdfield_t,
+               mu: vfdfield_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 \\begin{bmatrix} \\mu_{yy} \\epsilon_{xx} & 0 \\\\
+                                                       0 & \\mu_{xx} \\epsilon_{yy} \\end{bmatrix} +
+                 \\begin{bmatrix} -\\mu_{yy} \\hat{\\partial}_y \\\\
+                                   \\mu_{xx} \\hat{\\partial}_x \\end{bmatrix} \\mu_{zz}^{-1}
+                 \\begin{bmatrix} -\\tilde{\\partial}_y & \\tilde{\\partial}_x \\end{bmatrix} +
+      \\begin{bmatrix} \\tilde{\\partial}_x \\\\
+                       \\tilde{\\partial}_y \\end{bmatrix} \\epsilon_{zz}^{-1}
+                 \\begin{bmatrix} \\hat{\\partial}_x \\epsilon_{xx} & \\hat{\\partial}_y \\epsilon_{yy} \\end{bmatrix}
+    $$
+
+    $\\tilde{\\partial}_x$ and $\\hat{\\partial}_x$ are the forward and backward derivatives along x,
+    and each $\\epsilon_{xx}$, $\\mu_{yy}$, etc. is a diagonal matrix containing the vectorized material
+    property distribution.
+
+    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.fdmath.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 = deriv_forward(dxes[0])
+    Dbx, Dby = 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: vfdfield_t,
+               mu: vfdfield_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 \\begin{bmatrix} \\epsilon_{yy} \\mu_{xx} & 0 \\\\
+                                                      0 & \\epsilon_{xx} \\mu_{yy} \\end{bmatrix} +
+                 \\begin{bmatrix} -\\epsilon_{yy} \\tilde{\\partial}_y \\\\
+                                   \\epsilon_{xx} \\tilde{\\partial}_x \\end{bmatrix} \\epsilon_{zz}^{-1}
+                 \\begin{bmatrix} -\\hat{\\partial}_y & \\hat{\\partial}_x \\end{bmatrix} +
+      \\begin{bmatrix} \\hat{\\partial}_x \\\\
+                       \\hat{\\partial}_y \\end{bmatrix} \\mu_{zz}^{-1}
+                 \\begin{bmatrix} \\tilde{\\partial}_x \\mu_{xx} & \\tilde{\\partial}_y \\mu_{yy} \\end{bmatrix}
+    $$
+
+    $\\tilde{\\partial}_x$ and $\\hat{\\partial}_x$ are the forward and backward derivatives along x,
+    and each $\\epsilon_{xx}$, $\\mu_{yy}$, etc. is a diagonal matrix containing the vectorized material
+    property distribution.
+
+    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.fdmath.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: vfdfield_t,
+                        mu: vfdfield_t = None,
+                        prop_phase: float = 0,
+                        ) -> Tuple[vfdfield_t, vfdfield_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.fdmath.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: vfdfield_t,
+                        mu: vfdfield_t = None,
+                        prop_phase: float = 0,
+                        ) -> Tuple[vfdfield_t, vfdfield_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.fdmath.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: vfdfield_t,
+                       mu: vfdfield_t = None,
+                       prop_phase: float = 0,
+                       ) -> Tuple[vfdfield_t, vfdfield_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: vfdfield_t,
+          mu: vfdfield_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.fdmath.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: vfdfield_t,
+          mu: vfdfield_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.fdmath.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: vfdfield_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.fdmath.types` (2D)
+        mu: Vectorized magnetic permeability grid (default 1 everywhere)
+
+    Returns:
+        Sparse matrix representing the operator.
+    """
+    Dfx, Dfy = 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: vfdfield_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.fdmath.types` (2D)
+        epsilon: Vectorized dielectric constant grid
+
+    Returns:
+        Sparse matrix representing the operator.
+    """
+    Dbx, Dby = 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: vfdfield_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.fdmath.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: vfdfield_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.fdmath.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.fdmath.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 = deriv_forward(dxes[0])
+    return 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.fdmath.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 = deriv_back(dxes[1])
+    return cross([Dbx, Dby, Bz])
+
+
+def h_err(h: vfdfield_t,
+          wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfdfield_t,
+          mu: vfdfield_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.fdmath.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: vfdfield_t,
+          wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfdfield_t,
+          mu: vfdfield_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.fdmath.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: vfdfield_t,
+                mu: vfdfield_t = None,
+                mode_margin: int = 2,
+                ) -> Tuple[List[vfdfield_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.fdmath.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 = operator_e(numpy.real(omega), dxes_real, numpy.real(epsilon), numpy.real(mu))
+
+    eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin)
+    e_xys = eigvecs[:, -(numpy.array(mode_numbers) + 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 = operator_e(omega, dxes, epsilon, mu)
+    for nn in range(len(mode_numbers)):
+        eigvals[nn], e_xys[:, nn] = rayleigh_quotient_iteration(A, e_xys[:, nn])
+
+    # 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[vfdfield_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)
+    """
+    e_xys, wavenumbers = solve_modes(mode_numbers=[mode_number], *args, **kwargs)
+    return e_xys[:, 0], wavenumbers[0]

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 ..fdmath import vec, unvec, dx_lists_t, vfdfield_t, fdfield_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: fdfield_t,
+               mu: fdfield_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.fdmath.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_2d.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: fdfield_t,
+                   wavenumber: complex,
+                   omega: complex,
+                   dxes: dx_lists_t,
+                   axis: int,
+                   polarity: int,
+                   slices: List[slice],
+                   epsilon: fdfield_t,
+                   mu: fdfield_t = None,
+                   ) -> fdfield_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.fdmath.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: fdfield_t,
+                      wavenumber: complex,
+                      dxes: dx_lists_t,
+                      axis: int,
+                      polarity: int,
+                      slices: List[slice],
+                      ) -> fdfield_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.fdmath.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: fdfield_t,
+             wavenumber: complex,
+             dxes: dx_lists_t,
+             axis: int,
+             polarity: int,
+             slices: List[slice],
+             ) -> fdfield_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.fdmath.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 ..fdmath import vec, unvec, dx_lists_t, fdfield_t, vfdfield_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: vfdfield_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.fdmath.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: vfdfield_t,
+               r0: float,
+               ) -> Dict[str, complex or fdfield_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.fdmath.types.
+              The first coordinate is assumed to be r, the second is y.
+        epsilon: Dielectric constant
+        r0: Radius of curvature for the simulation. This should be the minimum value of
+            r within the simulation domain.
+
+    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/fdmath/__init__.py

@@ -0,0 +1,747 @@
+"""
+
+Basic discrete calculus for finite difference (fd) simulations.
+
+
+Fields, Functions, and Operators
+================================
+
+Discrete fields are stored in one of two forms:
+
+- The `fdfield_t` form is a multidimensional `numpy.ndarray`
+    + For a scalar field, this is just `U[m, n, p]`, where `m`, `n`, and `p` are
+      discrete indices referring to positions on the x, y, and z axes respectively.
+    + For a vector field, the first index specifies which vector component is accessed:
+      `E[:, m, n, p] = [Ex[m, n, p], Ey[m, n, p], Ez[m, n, p]]`.
+- The `vfdfield_t` form is simply a vectorzied (i.e. 1D) version of the `field_t`,
+    as obtained by `meanas.fdmath.vectorization.vec` (effectively just `numpy.ravel`)
+
+Operators which act on fields also come in two forms:
+    + Python functions, created by the functions in `meanas.fdmath.functional`.
+        The generated functions act on fields in the `fdfield_t` form.
+    + Linear operators, usually 2D sparse matrices using `scipy.sparse`, created
+        by `meanas.fdmath.operators`. These operators act on vectorized fields in the
+        `vfdfield_t` form.
+
+The operations performed should be equivalent: `functional.op(*args)(E)` should be
+equivalent to `unvec(operators.op(*args) @ vec(E), E.shape[1:])`.
+
+Generally speaking the `field_t` form is easier to work with, but can be harder or less
+efficient to compose (e.g. it is easy to generate a single matrix by multiplying a
+series of other matrices).
+
+
+Discrete calculus
+=================
+
+This documentation and approach is roughly based on W.C. Chew's excellent
+"Electromagnetic Theory on a Lattice" (doi:10.1063/1.355770),
+which covers a superset of this material with similar notation and more detail.
+
+
+Scalar derivatives and cell shifts
+----------------------------------
+
+Define the discrete forward derivative as
+ $$ [\\tilde{\\partial}_x f]_{m + \\frac{1}{2}} = \\frac{1}{\\Delta_{x, m}} (f_{m + 1} - f_m) $$
+ where $f$ is a function defined at discrete locations on the x-axis (labeled using $m$).
+ The value at $m$ occupies a length $\\Delta_{x, m}$ along the x-axis. Note that $m$
+ is an index along the x-axis, _not_ necessarily an x-coordinate, since each length
+ $\\Delta_{x, m}, \\Delta_{x, m+1}, ...$ is independently chosen.
+
+If we treat `f` as a 1D array of values, with the `i`-th value `f[i]` taking up a length `dx[i]`
+along the x-axis, the forward derivative is
+
+    deriv_forward(f)[i] = (f[i + 1] - f[i]) / dx[i]
+
+
+Likewise, discrete reverse derivative is
+ $$ [\\hat{\\partial}_x f ]_{m - \\frac{1}{2}} = \\frac{1}{\\Delta_{x, m}} (f_{m} - f_{m - 1}) $$
+ or
+
+    deriv_back(f)[i] = (f[i] - f[i - 1]) / dx[i]
+
+The derivatives' values are shifted by a half-cell relative to the original function, and
+will have different cell widths if all the `dx[i]` ( $\\Delta_{x, m}$ ) are not
+identical:
+
+    [figure: derivatives and cell sizes]
+        dx0   dx1      dx2      dx3      cell sizes for function
+       ----- ----- ----------- -----
+       ______________________________
+            |     |           |     |
+         f0 |  f1 |     f2    |  f3 |    function
+       _____|_____|___________|_____|
+         |     |        |        |
+         | Df0 |   Df1  |   Df2  | Df3   forward derivative (periodic boundary)
+       __|_____|________|________|___
+
+     dx'3] dx'0   dx'1     dx'2  [dx'3   cell sizes for forward derivative
+       -- ----- -------- -------- ---
+     dx'0] dx'1   dx'2     dx'3  [dx'0   cell sizes for reverse derivative
+       ______________________________
+         |     |        |        |
+         | df1 |  df2   |   df3  | df0   reverse derivative (periodic boundary)
+       __|_____|________|________|___
+
+    Periodic boundaries are used here and elsewhere unless otherwise noted.
+
+In the above figure,
+ `f0 =` $f_0$, `f1 =` $f_1$
+ `Df0 =` $[\\tilde{\\partial}f]_{0 + \\frac{1}{2}}$
+ `Df1 =` $[\\tilde{\\partial}f]_{1 + \\frac{1}{2}}$
+ `df0 =` $[\\hat{\\partial}f]_{0 - \\frac{1}{2}}$
+ etc.
+
+The fractional subscript $m + \\frac{1}{2}$ is used to indicate values defined
+ at shifted locations relative to the original $m$, with corresponding lengths
+ $$ \\Delta_{x, m + \\frac{1}{2}} = \\frac{1}{2} * (\\Delta_{x, m} + \\Delta_{x, m + 1}) $$
+
+Just as $m$ is not itself an x-coordinate, neither is $m + \\frac{1}{2}$;
+carefully note the positions of the various cells in the above figure vs their labels.
+If the positions labeled with $m$ are considered the "base" or "original" grid,
+the positions labeled with $m + \\frac{1}{2}$ are said to lie on a "dual" or
+"derived" grid.
+
+For the remainder of the `Discrete calculus` section, all figures will show
+constant-length cells in order to focus on the vector derivatives themselves.
+See the `Grid description` section below for additional information on this topic
+and generalization to three dimensions.
+
+
+Gradients and fore-vectors
+--------------------------
+
+Expanding to three dimensions, we can define two gradients
+  $$ [\\tilde{\\nabla} f]_{m,n,p} = \\vec{x} [\\tilde{\\partial}_x f]_{m + \\frac{1}{2},n,p} +
+                                    \\vec{y} [\\tilde{\\partial}_y f]_{m,n + \\frac{1}{2},p} +
+                                    \\vec{z} [\\tilde{\\partial}_z f]_{m,n,p + \\frac{1}{2}}  $$
+  $$ [\\hat{\\nabla} f]_{m,n,p} = \\vec{x} [\\hat{\\partial}_x f]_{m + \\frac{1}{2},n,p} +
+                                  \\vec{y} [\\hat{\\partial}_y f]_{m,n + \\frac{1}{2},p} +
+                                  \\vec{z} [\\hat{\\partial}_z f]_{m,n,p + \\frac{1}{2}}  $$
+
+ or
+
+    [code: gradients]
+    grad_forward(f)[i,j,k] = [Dx_forward(f)[i, j, k],
+                              Dy_forward(f)[i, j, k],
+                              Dz_forward(f)[i, j, k]]
+                           = [(f[i + 1, j, k] - f[i, j, k]) / dx[i],
+                              (f[i, j + 1, k] - f[i, j, k]) / dy[i],
+                              (f[i, j, k + 1] - f[i, j, k]) / dz[i]]
+
+    grad_back(f)[i,j,k] = [Dx_back(f)[i, j, k],
+                           Dy_back(f)[i, j, k],
+                           Dz_back(f)[i, j, k]]
+                        = [(f[i, j, k] - f[i - 1, j, k]) / dx[i],
+                           (f[i, j, k] - f[i, j - 1, k]) / dy[i],
+                           (f[i, j, k] - f[i, j, k - 1]) / dz[i]]
+
+The three derivatives in the gradient cause shifts in different
+directions, so the x/y/z components of the resulting "vector" are defined
+at different points: the x-component is shifted in the x-direction,
+y in y, and z in z.
+
+We call the resulting object a "fore-vector" or "back-vector", depending
+on the direction of the shift. We write it as
+  $$ \\tilde{g}_{m,n,p} = \\vec{x} g^x_{m + \\frac{1}{2},n,p} +
+                          \\vec{y} g^y_{m,n + \\frac{1}{2},p} +
+                          \\vec{z} g^z_{m,n,p + \\frac{1}{2}} $$
+  $$ \\hat{g}_{m,n,p} = \\vec{x} g^x_{m - \\frac{1}{2},n,p} +
+                        \\vec{y} g^y_{m,n - \\frac{1}{2},p} +
+                        \\vec{z} g^z_{m,n,p - \\frac{1}{2}} $$
+
+
+    [figure: gradient / fore-vector]
+       (m, n+1, p+1) ______________ (m+1, n+1, p+1)
+                    /:            /|
+                   / :           / |
+                  /  :          /  |
+      (m, n, p+1)/_____________/   |     The forward derivatives are defined
+                 |   :         |   |     at the Dx, Dy, Dz points,
+                 |   :.........|...|     but the forward-gradient fore-vector
+     z y        Dz  /          |  /      is the set of all three
+     |/_x        | Dy          | /       and is said to be "located" at (m,n,p)
+                 |/            |/
+        (m, n, p)|_____Dx______| (m+1, n, p)
+
+
+
+Divergences
+-----------
+
+There are also two divergences,
+
+  $$ d_{n,m,p} = [\\tilde{\\nabla} \\cdot \\hat{g}]_{n,m,p}
+               = [\\tilde{\\partial}_x g^x]_{m,n,p} +
+                 [\\tilde{\\partial}_y g^y]_{m,n,p} +
+                 [\\tilde{\\partial}_z g^z]_{m,n,p}   $$
+
+  $$ d_{n,m,p} = [\\hat{\\nabla} \\cdot \\tilde{g}]_{n,m,p}
+               = [\\hat{\\partial}_x g^x]_{m,n,p} +
+                 [\\hat{\\partial}_y g^y]_{m,n,p} +
+                 [\\hat{\\partial}_z g^z]_{m,n,p}  $$
+
+ or
+
+    [code: divergences]
+    div_forward(g)[i,j,k] = Dx_forward(gx)[i, j, k] +
+                            Dy_forward(gy)[i, j, k] +
+                            Dz_forward(gz)[i, j, k]
+                          = (gx[i + 1, j, k] - gx[i, j, k]) / dx[i] +
+                            (gy[i, j + 1, k] - gy[i, j, k]) / dy[i] +
+                            (gz[i, j, k + 1] - gz[i, j, k]) / dz[i]
+
+    div_back(g)[i,j,k] = Dx_back(gx)[i, j, k] +
+                         Dy_back(gy)[i, j, k] +
+                         Dz_back(gz)[i, j, k]
+                       = (gx[i, j, k] - gx[i - 1, j, k]) / dx[i] +
+                         (gy[i, j, k] - gy[i, j - 1, k]) / dy[i] +
+                         (gz[i, j, k] - gz[i, j, k - 1]) / dz[i]
+
+where `g = [gx, gy, gz]` is a fore- or back-vector field.
+
+Since we applied the forward divergence to the back-vector (and vice-versa), the resulting scalar value
+is defined at the back-vector's (fore-vector's) location $(m,n,p)$ and not at the locations of its components
+$(m \\pm \\frac{1}{2},n,p)$ etc.
+
+    [figure: divergence]
+                                    ^^
+         (m-1/2, n+1/2, p+1/2) _____||_______ (m+1/2, n+1/2, p+1/2)
+                              /:    ||  ,,  /|
+                             / :    || //  / |      The divergence at (m, n, p) (the center
+                            /  :      //  /  |      of this cube) of a fore-vector field
+      (m-1/2, n-1/2, p+1/2)/_____________/   |      is the sum of the outward-pointing
+                           |   :         |   |      fore-vector components, which are
+         z y            <==|== :.........|.====>    located at the face centers.
+         |/_x              |  /          |  /
+                           | /    //     | /       Note that in a nonuniform grid, each
+                           |/    // ||   |/        dimension is normalized by the cell width.
+      (m-1/2, n-1/2, p-1/2)|____//_______| (m+1/2, n-1/2, p-1/2)
+                               ''   ||
+                                    VV
+
+
+Curls
+-----
+
+The two curls are then
+
+  $$ \\begin{aligned}
+     \\hat{h}_{m + \\frac{1}{2}, n + \\frac{1}{2}, p + \\frac{1}{2}} &= \\\\
+     [\\tilde{\\nabla} \\times \\tilde{g}]_{m + \\frac{1}{2}, n + \\frac{1}{2}, p + \\frac{1}{2}} &=
+        \\vec{x} (\\tilde{\\partial}_y g^z_{m,n,p + \\frac{1}{2}} - \\tilde{\\partial}_z g^y_{m,n + \\frac{1}{2},p}) \\\\
+     &+ \\vec{y} (\\tilde{\\partial}_z g^x_{m + \\frac{1}{2},n,p} - \\tilde{\\partial}_x g^z_{m,n,p + \\frac{1}{2}}) \\\\
+     &+ \\vec{z} (\\tilde{\\partial}_x g^y_{m,n + \\frac{1}{2},p} - \\tilde{\\partial}_y g^z_{m + \\frac{1}{2},n,p})
+     \\end{aligned} $$
+
+ and
+
+  $$ \\tilde{h}_{m - \\frac{1}{2}, n - \\frac{1}{2}, p - \\frac{1}{2}} =
+     [\\hat{\\nabla} \\times \\hat{g}]_{m - \\frac{1}{2}, n - \\frac{1}{2}, p - \\frac{1}{2}} $$
+
+  where $\\hat{g}$ and $\\tilde{g}$ are located at $(m,n,p)$
+  with components at $(m \\pm \\frac{1}{2},n,p)$ etc.,
+  while $\\hat{h}$ and $\\tilde{h}$ are located at $(m \\pm \\frac{1}{2}, n \\pm \\frac{1}{2}, p \\pm \\frac{1}{2})$
+  with components at $(m, n \\pm \\frac{1}{2}, p \\pm \\frac{1}{2})$ etc.
+
+
+    [code: curls]
+    curl_forward(g)[i,j,k] = [Dy_forward(gz)[i, j, k] - Dz_forward(gy)[i, j, k],
+                              Dz_forward(gx)[i, j, k] - Dx_forward(gz)[i, j, k],
+                              Dx_forward(gy)[i, j, k] - Dy_forward(gx)[i, j, k]]
+
+    curl_back(g)[i,j,k] = [Dy_back(gz)[i, j, k] - Dz_back(gy)[i, j, k],
+                           Dz_back(gx)[i, j, k] - Dx_back(gz)[i, j, k],
+                           Dx_back(gy)[i, j, k] - Dy_back(gx)[i, j, k]]
+
+
+For example, consider the forward curl, at (m, n, p), of a back-vector field `g`, defined
+ on a grid containing (m + 1/2, n + 1/2, p + 1/2).
+ The curl will be a fore-vector, so its z-component will be defined at (m, n, p + 1/2).
+ Take the nearest x- and y-components of `g` in the xy plane where the curl's z-component
+ is located; these are
+
+    [curl components]
+    (m,       n + 1/2, p + 1/2) : x-component of back-vector at (m + 1/2, n + 1/2, p + 1/2)
+    (m + 1,   n + 1/2, p + 1/2) : x-component of back-vector at (m + 3/2, n + 1/2, p + 1/2)
+    (m + 1/2, n      , p + 1/2) : y-component of back-vector at (m + 1/2, n + 1/2, p + 1/2)
+    (m + 1/2, n + 1  , p + 1/2) : y-component of back-vector at (m + 1/2, n + 3/2, p + 1/2)
+
+ These four xy-components can be used to form a loop around the curl's z-component; its magnitude and sign
+ is set by their loop-oriented sum (i.e. two have their signs flipped to complete the loop).
+
+    [figure: z-component of curl]
+                              :             |
+        z y                   :    ^^       |
+        |/_x                  :....||.<.....|  (m+1, n+1, p+1/2)
+                              /    ||      /
+                           | v     ||   | ^
+                           |/           |/
+             (m, n, p+1/2) |_____>______|  (m+1, n, p+1/2)
+
+
+
+Maxwell's Equations
+===================
+
+If we discretize both space (m,n,p) and time (l), Maxwell's equations become
+
+ $$ \\begin{aligned}
+  \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}} &= -\\tilde{\\partial}_t \\hat{B}_{l-\\frac{1}{2}, \\vec{r} + \\frac{1}{2}}
+                                                                         - \\hat{M}_{l, \\vec{r} + \\frac{1}{2}}  \\\\
+  \\hat{\\nabla} \\times \\hat{H}_{l-\\frac{1}{2},\\vec{r} + \\frac{1}{2}} &= \\hat{\\partial}_t \\tilde{D}_{l, \\vec{r}}
+                                                                             + \\tilde{J}_{l-\\frac{1}{2},\\vec{r}} \\\\
+  \\tilde{\\nabla} \\cdot \\hat{B}_{l-\\frac{1}{2}, \\vec{r} + \\frac{1}{2}} &= 0 \\\\
+  \\hat{\\nabla} \\cdot \\tilde{D}_{l,\\vec{r}} &= \\rho_{l,\\vec{r}}
+ \\end{aligned} $$
+
+ with
+
+ $$ \\begin{aligned}
+  \\hat{B}_{\\vec{r}} &= \\mu_{\\vec{r} + \\frac{1}{2}} \\cdot \\hat{H}_{\\vec{r} + \\frac{1}{2}} \\\\
+  \\tilde{D}_{\\vec{r}} &= \\epsilon_{\\vec{r}} \\cdot \\tilde{E}_{\\vec{r}}
+ \\end{aligned} $$
+
+where the spatial subscripts are abbreviated as $\\vec{r} = (m, n, p)$ and
+$\\vec{r} + \\frac{1}{2} = (m + \\frac{1}{2}, n + \\frac{1}{2}, p + \\frac{1}{2})$,
+$\\tilde{E}$ and $\\hat{H}$ are the electric and magnetic fields,
+$\\tilde{J}$ and $\\hat{M}$ are the electric and magnetic current distributions,
+and $\\epsilon$ and $\\mu$ are the dielectric permittivity and magnetic permeability.
+
+The above is Yee's algorithm, written in a form analogous to Maxwell's equations.
+The time derivatives can be expanded to form the update equations:
+
+    [code: Maxwell's equations updates]
+    H[i, j, k] -= dt * (curl_forward(E)[i, j, k] + M[t, i, j, k]) /      mu[i, j, k]
+    E[i, j, k] += dt * (curl_back(   H)[i, j, k] + J[t, i, j, k]) / epsilon[i, j, k]
+
+Note that the E-field fore-vector and H-field back-vector are offset by a half-cell, resulting
+in distinct locations for all six E- and H-field components:
+
+    [figure: Field components]
+
+            (m - 1/2,=> ____________Hx__________[H] <= r + 1/2 = (m + 1/2,
+             n + 1/2,  /:           /:          /|                n + 1/2,
+       z y   p + 1/2) / :          / :         / |                p + 1/2)
+       |/_x          /  :         /  :        /  |
+                    /   :       Ez__________Hy   |      Locations of the E- and
+                   /    :        :   :      /|   |      H-field components for the
+     (m - 1/2,    /     :        :  Ey...../.|..Hz      [E] fore-vector at r = (m,n,p)
+      n - 1/2, =>/________________________/  |  /|      (the large cube's center)
+      p + 1/2)   |      :        : /      |  | / |      and [H] back-vector at r + 1/2
+                 |      :        :/       |  |/  |      (the top right corner)
+                 |      :       [E].......|.Ex   |
+                 |      :.................|......| <= (m + 1/2, n + 1/2, p + 1/2)
+                 |     /                  |     /
+                 |    /                   |    /
+                 |   /                    |   /         This is the Yee discretization
+                 |  /                     |  /          scheme ("Yee cell").
+    r - 1/2 =    | /                      | /
+     (m - 1/2,   |/                       |/
+      n - 1/2,=> |________________________| <= (m + 1/2, n - 1/2, p - 1/2)
+      p - 1/2)
+
+Each component forms its own grid, offset from the others:
+
+    [figure: E-fields for adjacent cells]
+
+                  H1__________Hx0_________H0
+      z y        /:                       /|
+      |/_x      / :                      / |    This figure shows H back-vector locations
+               /  :                     /  |    H0, H1, etc. and their associated components
+             Hy1  :                   Hy0  |    H0 = (Hx0, Hy0, Hz0) etc.
+             /    :                   /    |
+            /    Hz1                 /     Hz0
+           H2___________Hx3_________H3     |    The equivalent drawing for E would have
+           |      :                 |      |    fore-vectors located at the cube's
+           |      :                 |      |    center (and the centers of adjacent cubes),
+           |      :                 |      |    with components on the cube's faces.
+           |      H5..........Hx4...|......H4
+           |     /                  |     /
+          Hz2   /                  Hz2   /
+           |   /                    |   /
+           | Hy6                    | Hy4
+           | /                      | /
+           |/                       |/
+           H6__________Hx7__________H7
+
+
+The divergence equations can be derived by taking the divergence of the curl equations
+and combining them with charge continuity,
+ $$ \\hat{\\nabla} \\cdot \\tilde{J} + \\hat{\\partial}_t \\rho = 0 $$
+ implying that the discrete Maxwell's equations do not produce spurious charges.
+
+
+Wave equation
+-------------
+
+Taking the backward curl of the $\\tilde{\\nabla} \\times \\tilde{E}$ equation and
+replacing the resulting $\\hat{\\nabla} \\times \\hat{H}$ term using its respective equation,
+and setting $\\hat{M}$ to zero, we can form the discrete wave equation:
+
+$$
+  \\begin{aligned}
+  \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}} &=
+      -\\tilde{\\partial}_t \\hat{B}_{l-\\frac{1}{2}, \\vec{r} + \\frac{1}{2}}
+                          - \\hat{M}_{l-1, \\vec{r} + \\frac{1}{2}}  \\\\
+  \\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}} &=
+    -\\tilde{\\partial}_t \\hat{H}_{l-\\frac{1}{2}, \\vec{r} + \\frac{1}{2}}  \\\\
+  \\hat{\\nabla} \\times (\\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}}) &=
+    \\hat{\\nabla} \\times (-\\tilde{\\partial}_t \\hat{H}_{l-\\frac{1}{2}, \\vec{r} + \\frac{1}{2}})  \\\\
+  \\hat{\\nabla} \\times (\\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}}) &=
+    -\\tilde{\\partial}_t \\hat{\\nabla} \\times \\hat{H}_{l-\\frac{1}{2}, \\vec{r} + \\frac{1}{2}}  \\\\
+  \\hat{\\nabla} \\times (\\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}}) &=
+    -\\tilde{\\partial}_t \\hat{\\partial}_t \\epsilon_{\\vec{r}} \\tilde{E}_{l, \\vec{r}} + \\hat{\\partial}_t \\tilde{J}_{l-\\frac{1}{2},\\vec{r}} \\\\
+  \\hat{\\nabla} \\times (\\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{l,\\vec{r}})
+            + \\tilde{\\partial}_t \\hat{\\partial}_t \\epsilon_{\\vec{r}} \\cdot \\tilde{E}_{l, \\vec{r}}
+            &= \\tilde{\\partial}_t \\tilde{J}_{l - \\frac{1}{2}, \\vec{r}}
+  \\end{aligned}
+$$
+
+
+Frequency domain
+----------------
+
+We can substitute in a time-harmonic fields
+
+$$
+ \\begin{aligned}
+ \\tilde{E}_{l, \\vec{r}} &= \\tilde{E}_{\\vec{r}} e^{-\\imath \\omega l \\Delta_t} \\\\
+ \\tilde{J}_{l, \\vec{r}} &= \\tilde{J}_{\\vec{r}} e^{-\\imath \\omega (l - \\frac{1}{2}) \\Delta_t}
+ \\end{aligned}
+$$
+
+resulting in
+
+$$
+ \\begin{aligned}
+ \\tilde{\\partial}_t &\\Rightarrow (e^{ \\imath \\omega \\Delta_t} - 1) / \\Delta_t = \\frac{-2 \\imath}{\\Delta_t} \\sin(\\omega \\Delta_t / 2) e^{-\\imath \\omega \\Delta_t / 2} = -\\imath \\Omega e^{-\\imath \\omega \\Delta_t / 2}\\\\
+   \\hat{\\partial}_t &\\Rightarrow (1 - e^{-\\imath \\omega \\Delta_t}) / \\Delta_t = \\frac{-2 \\imath}{\\Delta_t} \\sin(\\omega \\Delta_t / 2) e^{ \\imath \\omega \\Delta_t / 2} = -\\imath \\Omega e^{ \\imath \\omega \\Delta_t / 2}\\\\
+ \\Omega &= 2 \\sin(\\omega \\Delta_t / 2) / \\Delta_t
+ \\end{aligned}
+$$
+
+This gives the frequency-domain wave equation,
+
+$$
+ \\hat{\\nabla} \\times (\\mu^{-1}_{\\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{\\vec{r}})
+    -\\Omega^2 \\epsilon_{\\vec{r}} \\cdot \\tilde{E}_{\\vec{r}} = \\imath \\Omega \\tilde{J}_{\\vec{r}}
+$$
+
+
+Plane waves and Dispersion relation
+------------------------------------
+
+With uniform material distribution and no sources
+
+$$
+ \\begin{aligned}
+ \\mu_{\\vec{r} + \\frac{1}{2}} &= \\mu \\\\
+ \\epsilon_{\\vec{r}} &= \\epsilon \\\\
+ \\tilde{J}_{\\vec{r}} &= 0 \\\\
+ \\end{aligned}
+$$
+
+the frequency domain wave equation simplifies to
+
+$$ \\hat{\\nabla} \\times \\tilde{\\nabla} \\times \\tilde{E}_{\\vec{r}} - \\Omega^2 \\epsilon \\mu \\tilde{E}_{\\vec{r}} = 0 $$
+
+Since $\\hat{\\nabla} \\cdot \\tilde{E}_{\\vec{r}} = 0$, we can simplify
+
+$$
+ \\begin{aligned}
+ \\hat{\\nabla} \\times \\tilde{\\nabla} \\times \\tilde{E}_{\\vec{r}}
+   &= \\tilde{\\nabla}(\\hat{\\nabla} \\cdot \\tilde{E}_{\\vec{r}}) - \\hat{\\nabla} \\cdot \\tilde{\\nabla} \\tilde{E}_{\\vec{r}} \\\\
+   &= - \\hat{\\nabla} \\cdot \\tilde{\\nabla} \\tilde{E}_{\\vec{r}} \\\\
+   &= - \\tilde{\\nabla}^2 \\tilde{E}_{\\vec{r}}
+ \\end{aligned}
+$$
+
+and we get
+
+$$  \\tilde{\\nabla}^2 \\tilde{E}_{\\vec{r}} + \\Omega^2 \\epsilon \\mu \\tilde{E}_{\\vec{r}} = 0 $$
+
+We can convert this to three scalar-wave equations of the form
+
+$$ (\\tilde{\\nabla}^2 + K^2) \\phi_{\\vec{r}} = 0 $$
+
+with $K^2 = \\Omega^2 \\mu \\epsilon$. Now we let
+
+$$  \\phi_{\\vec{r}} = A e^{\\imath (k_x m \\Delta_x + k_y n \\Delta_y + k_z p \\Delta_z)}  $$
+
+resulting in
+
+$$
+ \\begin{aligned}
+ \\tilde{\\partial}_x &\\Rightarrow (e^{ \\imath k_x \\Delta_x} - 1) / \\Delta_t = \\frac{-2 \\imath}{\\Delta_x} \\sin(k_x \\Delta_x / 2) e^{ \\imath k_x \\Delta_x / 2} = \\imath K_x e^{ \\imath k_x \\Delta_x / 2}\\\\
+   \\hat{\\partial}_x &\\Rightarrow (1 - e^{-\\imath k_x \\Delta_x}) / \\Delta_t = \\frac{-2 \\imath}{\\Delta_x} \\sin(k_x \\Delta_x / 2) e^{-\\imath k_x \\Delta_x / 2} = \\imath K_x e^{-\\imath k_x \\Delta_x / 2}\\\\
+ K_x &= 2 \\sin(k_x \\Delta_x / 2) / \\Delta_x \\\\
+ \\end{aligned}
+$$
+
+with similar expressions for the y and z dimnsions (and $K_y, K_z$).
+
+This implies
+
+$$
+  \\tilde{\\nabla}^2 = -(K_x^2 + K_y^2 + K_z^2) \\phi_{\\vec{r}} \\\\
+  K_x^2 + K_y^2 + K_z^2 = \\Omega^2 \\mu \\epsilon = \\Omega^2 / c^2
+$$
+
+where $c = \\sqrt{\\mu \\epsilon}$.
+
+Assuming real $(k_x, k_y, k_z), \\omega$ will be real only if
+
+$$ c^2 \\Delta_t^2 = \\frac{\\Delta_t^2}{\\mu \\epsilon} < 1/(\\frac{1}{\\Delta_x^2} + \\frac{1}{\\Delta_y^2} + \\frac{1}{\\Delta_z^2}) $$
+
+If $\\Delta_x = \\Delta_y = \\Delta_z$, this simplifies to $c \\Delta_t < \\Delta_x / \\sqrt{3}$.
+This last form can be interpreted as enforcing causality; the distance that light
+travels in one timestep (i.e., $c \\Delta_t$) must be less than the diagonal
+of the smallest cell ( $\\Delta_x / \\sqrt{3}$ when on a uniform cubic grid).
+
+
+Grid description
+================
+
+As described in the section on scalar discrete derivatives above, cell widths
+(`dx[i]`, `dy[j]`, `dz[k]`) along each axis can be arbitrary and independently
+defined. Moreover, all field components are actually defined at "derived" or "dual"
+positions, in-between the "base" grid points on one or more axes.
+
+To get a better sense of how this works, let's start by drawing a grid with uniform
+`dy` and `dz` and nonuniform `dx`. We will only draw one cell in the y and z dimensions
+to make the illustration simpler; we need at least two cells in the x dimension to
+demonstrate how nonuniform `dx` affects the various components.
+
+Place the E fore-vectors at integer indices $r = (m, n, p)$ and the H back-vectors
+at fractional indices $r + \\frac{1}{2} = (m + \\frac{1}{2}, n + \\frac{1}{2},
+p + \\frac{1}{2})$. Remember that these are indices and not coordinates; they can
+correspond to arbitrary (monotonically increasing) coordinates depending on the cell widths.
+
+Draw lines to denote the planes on which the H components and back-vectors are defined.
+For simplicity, don't draw the equivalent planes for the E components and fore-vectors,
+except as necessary to show their locations -- it's easiest to just connect them to their
+associated H-equivalents.
+
+The result looks something like this:
+
+    [figure: Component centers]
+                                                                 p=
+              [H]__________Hx___________[H]_____Hx______[H]   __ +1/2
+      z y     /:           /:           /:      /:      /|     |      |
+      |/_x   / :          / :          / :     / :     / |     |      |
+            /  :         /  :         /  :    /  :    /  |     |      |
+          Hy   :       Ez...........Hy   :  Ez......Hy   |     |      |
+          /:   :        :   :       /:   :   :   :  /|   |     |      |
+         / :  Hz        :  Ey....../.:..Hz   :  Ey./.|..Hz    __ 0    | dz[0]
+        /  :  /:        :  /      /  :  /:   :  / /  |  /|     |      |
+       /_________________________/_______________/   | / |     |      |
+       |   :/  :        :/       |   :/  :   :/  |   |/  |     |      |
+       |  Ex   :       [E].......|..Ex   :  [E]..|..Ex   |     |      |
+       |       :                 |       :       |       |     |      |
+       |      [H]..........Hx....|......[H].....H|x.....[H]   __ --------- (n=+1/2, p=-1/2)
+       |      /                  |      /        |      /     /       /
+      Hz     /                  Hz     /        Hz     /     /       /
+       |    /                    |    /          |    /     /       /
+       |  Hy                     |  Hy           |  Hy    __ 0     / dy[0]
+       |  /                      |  /            |  /     /       /
+       | /                       | /             | /     /       /
+       |/                        |/              |/     /       /
+      [H]__________Hx___________[H]_____Hx______[H]   __ -1/2  /
+                                                           =n
+       |------------|------------|-------|-------|
+     -1/2           0          +1/2     +1     +3/2 = m
+
+        ------------------------- ----------------
+                  dx[0]                  dx[1]
+
+      Part of a nonuniform "base grid", with labels specifying
+      positions of the various field components. [E] fore-vectors
+      are at the cell centers, and [H] back-vectors are at the
+      vertices. H components along the near (-y) top (+z) edge
+      have been omitted to make the insides of the cubes easier
+      to visualize.
+
+The above figure shows where all the components are located; however, it is also useful to show
+what volumes those components correspond to. Consider the Ex component at `m = +1/2`: it is
+shifted in the x-direction by a half-cell from the E fore-vector at `m = 0` (labeled `[E]`
+in the figure). It corresponds to a volume between `m = 0` and `m = +1` (the other
+dimensions are not shifted, i.e. they are still bounded by `n, p = +-1/2`). (See figure
+below). Since `m` is an index and not an x-coordinate, the Ex component is not necessarily
+at the center of the volume it represents, and the x-length of its volume is the derived
+quantity `dx'[0] = (dx[0] + dx[1]) / 2` rather than the base `dx`.
+(See also `Scalar derivatives and cell shifts`).
+
+    [figure: Ex volumes]
+                                                                 p=
+               <_________________________________________>   __ +1/2
+      z y     <<           /:           /       /:      >>    |      |
+      |/_x   < <          / :          /       / :     > >    |      |
+            <  <         /  :         /       /  :    >  >    |      |
+           <   <        /   :        /       /   :   >   >    |      |
+          <:   <       /    :        :      /    :  >:   >    |      |
+         < :   <      /     :        :     /     : > :   >   __ 0    | dz[0]
+        <  :   <     /      :        :    /      :>  :   >    |      |
+       <____________/____________________/_______>   :   >    |      |
+       <   :   <    |       :        :   |       >   :   >    |      |
+       <  Ex   <    |       :       Ex   |       >  Ex   >    |      |
+       <   :   <    |       :        :   |       >   :   >    |      |
+       <   :   <....|.......:........:...|.......>...:...>   __ --------- (n=+1/2, p=-1/2)
+       <   :  <     |      /         :  /|      />   :  >    /       /
+       <   : <      |     /          : / |     / >   : >    /       /
+       <   :<       |    /           :/  |    /  >   :>    /       /
+       <   <        |   /            :   |   /   >   >    _ 0     / dy[0]
+       <  <         |  /                 |  /    >  >    /       /
+       < <          | /                  | /     > >    /       /
+       <<           |/                   |/      >>    /       /
+       <____________|____________________|_______>   __ -1/2  /
+                                                         =n
+       |------------|------------|-------|-------|
+     -1/2           0          +1/2      +1    +3/2 = m
+
+       ~------------ -------------------- -------~
+         dx'[-1]          dx'[0]           dx'[1]
+
+     The Ex values are positioned on the x-faces of the base
+     grid. They represent the Ex field in volumes shifted by
+     a half-cell in the x-dimension, as shown here. Only the
+     center cell (with width dx'[0]) is fully shown; the
+     other two are truncated (shown using >< markers).
+
+     Note that the Ex positions are the in the same positions
+     as the previous figure; only the cell boundaries have moved.
+     Also note that the points at which Ex is defined are not
+     necessarily centered in the volumes they represent; non-
+     uniform cell sizes result in off-center volumes like the
+     center cell here.
+
+The next figure shows the volumes corresponding to the Hy components, which
+are shifted in two dimensions (x and z) compared to the base grid.
+
+    [figure: Hy volumes]
+                                                                 p=
+      z y     mmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmm   __ +1/2  s
+      |/_x   <<           m:                    m:      >>    |       |
+            < <          m :                   m :     > >    |       | dz'[1]
+           <  <         m  :                  m  :    >  >    |       |
+         Hy........... m........Hy...........m......Hy   >    |       |
+         <    <       m    :                m    :  >    >    |       |
+        <     ______ m_____:_______________m_____:_>______   __ 0
+       <      <     m     /:              m     / >      >    |       |
+      mmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmm       >    |       |
+      <       <    |    /  :             |    /  >       >    |       | dz'[0]
+      <       <    |   /   :             |   /   >       >    |       |
+      <       <    |  /    :             |  /    >       >    |       |
+      <       wwwww|w/wwwwwwwwwwwwwwwwwww|w/wwwww>wwwwwwww   __       s
+      <      <     |/     w              |/     w>      >    /         /
+      _____________|_____________________|________     >    /         /
+      <    <       |    w                |    w  >    >    /         /
+      <  Hy........|...w........Hy.......|...w...>..Hy    _ 0       / dy[0]
+      < <          |  w                  |  w    >  >    /         /
+      <<           | w                   | w     > >    /         /
+      <            |w                    |w      >>    /         /
+      wwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwwww   __ -1/2    /
+
+      |------------|------------|--------|-------|
+    -1/2           0          +1/2      +1     +3/2 = m
+
+      ~------------ --------------------- -------~
+         dx'[-1]            dx'[0]         dx'[1]
+
+     The Hy values are positioned on the y-edges of the base
+     grid. Again here, the 'Hy' labels represent the same points
+     as in the basic grid figure above; the edges have shifted
+     by a half-cell along the x- and z-axes.
+
+     The grid lines _|:/ are edges of the area represented by
+     each Hy value, and the lines drawn using <m>.w represent
+     edges where a cell's faces extend beyond the drawn area
+     (i.e. where the drawing is truncated in the x- or z-
+     directions).
+
+
+Datastructure: dx_lists_t
+-------------------
+
+In this documentation, the E fore-vectors are placed on the base grid. An
+equivalent formulation could place the H back-vectors on the base grid instead.
+However, in the case of a non-uniform grid, the operation to get from the "base"
+cell widths to "derived" ones is not its own inverse.
+
+The base grid's cell sizes could be fully described by a list of three 1D arrays,
+specifying the cell widths along all three axes:
+
+    [dx, dy, dz] = [[dx[0], dx[1], ...], [dy[0], ...], [dz[0], ...]]
+
+Note that this is a list-of-arrays rather than a 2D array, as the simulation domain
+may have a different number of cells along each axis.
+
+Knowing the base grid's cell widths and the boundary conditions (periodic unless
+otherwise noted) is enough information to calculate the cell widths  `dx'`, `dy'`,
+and `dz'` for the derived grids.
+
+However, since most operations are trivially generalized to allow either E or H
+to be defined on the base grid, they are written to take the a full set of base
+and derived cell widths, distinguished by which field they apply to rather than
+their "base" or "derived" status. This removes the need for each function to
+generate the derived widths, and makes the "base" vs "derived" distinction
+unnecessary in the code.
+
+The resulting data structure containing all the cell widths takes the form of a
+list-of-lists-of-arrays. The first list-of-arrays provides the cell widths for
+the E-field fore-vectors, while the second list-of-arrays does the same for the
+H-field back-vectors:
+
+     [[[dx_e[0], dx_e[1], ...], [dy_e[0], ...], [dz_e[0], ...]],
+      [[dx_h[0], dx_h[1], ...], [dy_h[0], ...], [dz_h[0], ...]]]
+
+   where `dx_e[0]` is the x-width of the `m=0` cells, as used when calculating dE/dx,
+   and `dy_h[0]` is  the y-width of the `n=0` cells, as used when calculating dH/dy, etc.
+
+
+Permittivity and Permeability
+=============================
+
+Since each vector component of E and H is defined in a different location and represents
+a different volume, the value of the spatially-discrete `epsilon` and `mu` can also be
+different for all three field components, even when representing a simple planar interface
+between two isotropic materials.
+
+As a result, `epsilon` and `mu` are taken to have the same dimensions as the field, and
+composed of the three diagonal tensor components:
+
+    [equations: epsilon_and_mu]
+    epsilon = [epsilon_xx, epsilon_yy, epsilon_zz]
+    mu = [mu_xx, mu_yy, mu_zz]
+
+or
+
+$$
+ \\epsilon = \\begin{bmatrix} \\epsilon_{xx} & 0 & 0 \\\\
+                              0 & \\epsilon_{yy} & 0 \\\\
+                              0 & 0 & \\epsilon_{zz} \\end{bmatrix}
+$$
+$$
+ \\mu = \\begin{bmatrix} \\mu_{xx} & 0 & 0 \\\\
+                         0 & \\mu_{yy} & 0 \\\\
+                         0 & 0 & \\mu_{zz} \\end{bmatrix}
+$$
+
+where the off-diagonal terms (e.g. `epsilon_xy`) are assumed to be zero.
+
+High-accuracy volumetric integration of shapes on multiple grids can be performed
+by the [gridlock](https://mpxd.net/code/jan/gridlock) module.
+
+The values of the vacuum permittivity and permability effectively become scaling
+factors that appear in several locations (e.g. between the E and H fields). In
+order to limit floating-point inaccuracy and simplify calculations, they are often
+set to 1 and relative permittivities and permeabilities are used in their places;
+the true values can be multiplied back in after the simulation is complete if non-
+normalized results are needed.
+"""
+
+from .types import fdfield_t, vfdfield_t, dx_lists_t, fdfield_updater_t
+from .vectorization import vec, unvec
+from . import operators, functional, types, vectorization
+

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 .types import fdfield_t, fdfield_updater_t
+
+
+def deriv_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t:
+    """
+    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) -> fdfield_updater_t:
+    """
+    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) -> fdfield_updater_t:
+    """
+    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: fdfield_t) -> fdfield_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) -> fdfield_updater_t:
+    """
+    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: fdfield_t) -> fdfield_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 .types import fdfield_t, vfdfield_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: vfdfield_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/fdmath/types.py

@@ -0,0 +1,43 @@
+"""
+Types shared across multiple submodules
+"""
+import numpy
+from typing import List, Callable
+
+
+# Field types
+# TODO: figure out a better way to set the docstrings without creating actual subclasses?
+#   Probably not a big issue since they're only used for type hinting
+class fdfield_t(numpy.ndarray):
+    """
+    Vector field with shape (3, X, Y, Z) (e.g. `[E_x, E_y, E_z]`)
+
+    This is actually is just an unaltered `numpy.ndarray`
+    """
+    pass
+
+class vfdfield_t(numpy.ndarray):
+    """
+    Linearized vector field (single vector of length 3*X*Y*Z)
+
+    This is actually just an unaltered `numpy.ndarray`
+    """
+    pass
+
+
+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_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,
+   and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
+'''
+
+
+fdfield_updater_t = Callable[..., fdfield_t]
+'''
+ Convenience type for functions which take and return an fdfield_t
+'''

meanas/fdmath/vectorization.py

@@ -0,0 +1,52 @@
+"""
+Functions for moving between a vector field (list of 3 ndarrays, `[f_x, f_y, f_z]`)
+and a 1D array representation of that field `[f_x0, f_x1, f_x2,... f_y0,... f_z0,...]`.
+Vectorized versions of the field use row-major (ie., C-style) ordering.
+"""
+
+from typing import List
+import numpy
+
+from .types import fdfield_t, vfdfield_t
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def vec(f: fdfield_t) -> vfdfield_t:
+    """
+    Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
+
+    Returns `None` if called with `f=None`.
+
+    Args:
+        f: A vector field, `[f_x, f_y, f_z]` where each `f_` component is a 1- to
+           3-D ndarray (`f_*` should all be the same size). Doesn't fail with `f=None`.
+
+    Returns:
+        1D ndarray containing the linearized field (or `None`)
+    """
+    if numpy.any(numpy.equal(f, None)):
+        return None
+    return numpy.ravel(f, order='C')
+
+
+def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t:
+    """
+    Perform the inverse of vec(): take a 1D ndarray and output a 3D field
+     of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional
+     ndarray.
+
+    Returns `None` if called with `v=None`.
+
+    Args:
+        v: 1D ndarray representing a 3D vector field of shape shape (or None)
+        shape: shape of the vector field
+
+    Returns:
+        `[f_x, f_y, f_z]` where each `f_` is a `len(shape)` dimensional ndarray (or `None`)
+    """
+    if numpy.any(numpy.equal(v, None)):
+        return None
+    return v.reshape((3, *shape), order='C')
+

meanas/fdtd/__init__.py

@@ -0,0 +1,166 @@
+"""
+Utilities for running finite-difference time-domain (FDTD) simulations
+
+See the discussion of `Maxwell's Equations` in `meanas.fdmath` for basic
+mathematical background.
+
+
+Timestep
+========
+
+From the discussion of "Plane waves and the Dispersion relation" in `meanas.fdmath`,
+we have
+
+$$ c^2 \\Delta_t^2 = \\frac{\\Delta_t^2}{\\mu \\epsilon} < 1/(\\frac{1}{\\Delta_x^2} + \\frac{1}{\\Delta_y^2} + \\frac{1}{\\Delta_z^2}) $$
+
+or, if $\\Delta_x = \\Delta_y = \\Delta_z$, then $c \\Delta_t < \\frac{\\Delta_x}{\\sqrt{3}}$.
+
+Based on this, we can set
+
+    dt = sqrt(mu.min() * epsilon.min()) / sqrt(1/dx_min**2 + 1/dy_min**2 + 1/dz_min**2)
+
+The `dx_min`, `dy_min`, `dz_min` should be the minimum value across both the base and derived grids.
+
+
+Poynting Vector and Energy Conservation
+=======================================
+
+Let
+
+$$ \\begin{aligned}
+  \\tilde{S}_{l, l', \\vec{r}} &=& &\\tilde{E}_{l, \\vec{r}} \\otimes \\hat{H}_{l', \\vec{r} + \\frac{1}{2}}  \\\\
+  &=&  &\\vec{x} (\\tilde{E}^y_{l,m+1,n,p} \\hat{H}^z_{l',\\vec{r} + \\frac{1}{2}} - \\tilde{E}^z_{l,m+1,n,p} \\hat{H}^y_{l', \\vec{r} + \\frac{1}{2}}) \\\\
+  & &+ &\\vec{y} (\\tilde{E}^z_{l,m,n+1,p} \\hat{H}^x_{l',\\vec{r} + \\frac{1}{2}} - \\tilde{E}^x_{l,m,n+1,p} \\hat{H}^z_{l', \\vec{r} + \\frac{1}{2}}) \\\\
+  & &+ &\\vec{z} (\\tilde{E}^x_{l,m,n,p+1} \\hat{H}^y_{l',\\vec{r} + \\frac{1}{2}} - \\tilde{E}^y_{l,m,n,p+1} \\hat{H}^z_{l', \\vec{r} + \\frac{1}{2}})
+   \\end{aligned}
+$$
+
+where $\\vec{r} = (m, n, p)$ and $\\otimes$ is a modified cross product
+in which the $\\tilde{E}$ terms are shifted as indicated.
+
+By taking the divergence and rearranging terms, we can show that
+
+$$
+  \\begin{aligned}
+  \\hat{\\nabla} \\cdot \\tilde{S}_{l, l', \\vec{r}}
+   &= \\hat{\\nabla} \\cdot (\\tilde{E}_{l, \\vec{r}} \\otimes \\hat{H}_{l', \\vec{r} + \\frac{1}{2}})  \\\\
+   &= \\hat{H}_{l', \\vec{r} + \\frac{1}{2}} \\cdot \\tilde{\\nabla} \\times \\tilde{E}_{l, \\vec{r}} -
+      \\tilde{E}_{l, \\vec{r}} \\cdot \\hat{\\nabla} \\times \\hat{H}_{l', \\vec{r} + \\frac{1}{2}} \\\\
+   &= \\hat{H}_{l', \\vec{r} + \\frac{1}{2}} \\cdot
+            (-\\tilde{\\partial}_t \\mu_{\\vec{r} + \\frac{1}{2}} \\hat{H}_{l - \\frac{1}{2}, \\vec{r} + \\frac{1}{2}} -
+                \\hat{M}_{l-1, \\vec{r} + \\frac{1}{2}}) -
+      \\tilde{E}_{l, \\vec{r}} \\cdot (\\hat{\\partial}_t \\tilde{\\epsilon}_{\\vec{r}} \\tilde{E}_{l'+\\frac{1}{2}, \\vec{r}} +
+                \\tilde{J}_{l', \\vec{r}}) \\\\
+   &= \\hat{H}_{l'} \\cdot (-\\mu / \\Delta_t)(\\hat{H}_{l + \\frac{1}{2}} - \\hat{H}_{l - \\frac{1}{2}}) -
+      \\tilde{E}_l \\cdot (\\epsilon / \\Delta_t )(\\tilde{E}_{l'+\\frac{1}{2}} - \\tilde{E}_{l'-\\frac{1}{2}})
+      - \\hat{H}_{l'} \\cdot \\hat{M}_{l-1} - \\tilde{E}_l \\cdot \\tilde{J}_{l'} \\\\
+  \\end{aligned}
+$$
+
+where in the last line the spatial subscripts have been dropped to emphasize
+the time subscripts $l, l'$, i.e.
+
+$$
+  \\begin{aligned}
+  \\tilde{E}_l &= \\tilde{E}_{l, \\vec{r}} \\\\
+  \\hat{H}_l &= \\tilde{H}_{l, \\vec{r} + \\frac{1}{2}}  \\\\
+  \\tilde{\\epsilon} &= \\tilde{\\epsilon}_{\\vec{r}}  \\\\
+  \\end{aligned}
+$$
+
+etc.
+For $l' = l + \\frac{1}{2}$ we get
+
+$$
+  \\begin{aligned}
+  \\hat{\\nabla} \\cdot \\tilde{S}_{l, l + \\frac{1}{2}}
+   &= \\hat{H}_{l + \\frac{1}{2}} \\cdot
+            (-\\mu / \\Delta_t)(\\hat{H}_{l + \\frac{1}{2}} - \\hat{H}_{l - \\frac{1}{2}}) -
+      \\tilde{E}_l \\cdot (\\epsilon / \\Delta_t)(\\tilde{E}_{l+1} - \\tilde{E}_l)
+      - \\hat{H}_{l'} \\cdot \\hat{M}_l - \\tilde{E}_l \\cdot \\tilde{J}_{l + \\frac{1}{2}} \\\\
+   &= (-\\mu / \\Delta_t)(\\hat{H}^2_{l + \\frac{1}{2}} - \\hat{H}_{l + \\frac{1}{2}} \\cdot \\hat{H}_{l - \\frac{1}{2}}) -
+      (\\epsilon / \\Delta_t)(\\tilde{E}_{l+1} \\cdot \\tilde{E}_l - \\tilde{E}^2_l)
+      - \\hat{H}_{l'} \\cdot \\hat{M}_l - \\tilde{E}_l \\cdot \\tilde{J}_{l + \\frac{1}{2}} \\\\
+   &= -(\\mu \\hat{H}^2_{l + \\frac{1}{2}}
+       +\\epsilon \\tilde{E}_{l+1} \\cdot \\tilde{E}_l) / \\Delta_t \\ \\
+      +(\\mu \\hat{H}_{l + \\frac{1}{2}} \\cdot \\hat{H}_{l - \\frac{1}{2}}
+       +\\epsilon \\tilde{E}^2_l) / \\Delta_t \\ \\
+      - \\hat{H}_{l+\\frac{1}{2}} \\cdot \\hat{M}_l \\ \\
+      - \\tilde{E}_l \\cdot \\tilde{J}_{l+\\frac{1}{2}} \\\\
+  \\end{aligned}
+$$
+
+and for $l' = l - \\frac{1}{2}$,
+
+$$
+  \\begin{aligned}
+  \\hat{\\nabla} \\cdot \\tilde{S}_{l, l - \\frac{1}{2}}
+   &=  (\\mu \\hat{H}^2_{l - \\frac{1}{2}}
+       +\\epsilon \\tilde{E}_{l-1} \\cdot \\tilde{E}_l) / \\Delta_t \\ \\
+      -(\\mu \\hat{H}_{l + \\frac{1}{2}} \\cdot \\hat{H}_{l - \\frac{1}{2}}
+       +\\epsilon \\tilde{E}^2_l) / \\Delta_t \\ \\
+      - \\hat{H}_{l-\\frac{1}{2}} \\cdot \\hat{M}_l \\ \\
+      - \\tilde{E}_l \\cdot \\tilde{J}_{l-\\frac{1}{2}} \\\\
+  \\end{aligned}
+$$
+
+These two results form the discrete time-domain analogue to Poynting's theorem.
+They hint at the expressions for the energy, which can be calculated at the same
+time-index as either the E or H field:
+
+$$
+ \\begin{aligned}
+ U_l &= \\epsilon \\tilde{E}^2_l + \\mu \\hat{H}_{l + \\frac{1}{2}} \\cdot \\hat{H}_{l - \\frac{1}{2}} \\\\
+ U_{l + \\frac{1}{2}} &= \\epsilon \\tilde{E}_l \\cdot \\tilde{E}_{l + 1} + \\mu \\hat{H}^2_{l + \\frac{1}{2}} \\\\
+ \\end{aligned}
+$$
+
+Rewriting the Poynting theorem in terms of the energy expressions,
+
+$$
+  \\begin{aligned}
+  (U_{l+\\frac{1}{2}} - U_l) / \\Delta_t
+   &= -\\hat{\\nabla} \\cdot \\tilde{S}_{l, l + \\frac{1}{2}} \\ \\
+      - \\hat{H}_{l+\\frac{1}{2}} \\cdot \\hat{M}_l \\ \\
+      - \\tilde{E}_l \\cdot \\tilde{J}_{l+\\frac{1}{2}} \\\\
+  (U_l - U_{l-\\frac{1}{2}}) / \\Delta_t
+   &= -\\hat{\\nabla} \\cdot \\tilde{S}_{l, l - \\frac{1}{2}} \\ \\
+      - \\hat{H}_{l-\\frac{1}{2}} \\cdot \\hat{M}_l \\ \\
+      - \\tilde{E}_l \\cdot \\tilde{J}_{l-\\frac{1}{2}} \\\\
+ \\end{aligned}
+$$
+
+This result is exact an should practically hold to within numerical precision. No time-
+or spatial-averaging is necessary.
+
+Note that each value of $J$ contributes to the energy twice (i.e. once per field update)
+despite only causing the value of $E$ to change once (same for $M$ and $H$).
+
+
+Sources
+=============
+
+It is often useful to excite the simulation with an arbitrary broadband pulse and then
+extract the frequency-domain response by performing an on-the-fly Fourier transform
+of the time-domain fields.
+
+The Ricker wavelet (normalized second derivative of a Gaussian) is commonly used for the pulse
+shape. It can be written
+
+$$ f_r(t) = (1 - \\frac{1}{2} (\\omega (t - \\tau))^2) e^{-(\\frac{\\omega (t - \\tau)}{2})^2} $$
+
+with $\\tau > \\frac{2 * \\pi}{\\omega}$ as a minimum delay to avoid a discontinuity at
+t=0 (assuming the source is off for t<0 this gives $\\sim 10^{-3}$ error at t=0).
+
+
+
+Boundary conditions
+===================
+# TODO notes about boundaries / PMLs
+"""
+
+from .base import maxwell_e, maxwell_h
+from .pml import cpml
+from .energy import (poynting, poynting_divergence, energy_hstep, energy_estep,
+                     delta_energy_h2e, delta_energy_h2e, delta_energy_j)
+from .boundaries import conducting_boundary

meanas/fdtd/base.py

@@ -0,0 +1,122 @@
+"""
+Basic FDTD field updates
+
+
+"""
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
+from ..fdmath.functional import curl_forward, curl_back
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def maxwell_e(dt: float, dxes: dx_lists_t = None) -> fdfield_updater_t:
+    """
+    Build a function which performs a portion the time-domain E-field update,
+
+        E += curl_back(H[t]) / epsilon
+
+    The full update should be
+
+        E += (curl_back(H[t]) + J) / epsilon
+
+    which requires an additional step of `E += J / epsilon` which is not performed
+    by the generated function.
+
+    See `meanas.fdmath` for descriptions of
+
+    - This update step: "Maxwell's equations" section
+    - `dxes`: "Datastructure: dx_lists_t" section
+    - `epsilon`: "Permittivity and Permeability" section
+
+    Also see the "Timestep" section of `meanas.fdtd` for a discussion of
+    the `dt` parameter.
+
+    Args:
+        dt: Timestep. See `meanas.fdtd` for details.
+        dxes: Grid description; see `meanas.fdmath`.
+
+    Returns:
+        Function `f(E_old, H_old, epsilon) -> E_new`.
+    """
+    if dxes is not None:
+        curl_h_fun = curl_back(dxes[1])
+    else:
+        curl_h_fun = curl_back()
+
+    def me_fun(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t):
+        """
+        Update the E-field.
+
+        Args:
+            e: E-field at time t=0
+            h: H-field at time t=0.5
+            epsilon: Dielectric constant distribution.
+
+        Returns:
+            E-field at time t=1
+        """
+        e += dt * curl_h_fun(h) / epsilon
+        return e
+
+    return me_fun
+
+
+def maxwell_h(dt: float, dxes: dx_lists_t = None) -> fdfield_updater_t:
+    """
+    Build a function which performs part of the time-domain H-field update,
+
+        H -= curl_forward(E[t]) / mu
+
+    The full update should be
+
+        H -= (curl_forward(E[t]) + M) / mu
+
+    which requires an additional step of `H -= M / mu` which is not performed
+    by the generated function; this step can be omitted if there is no magnetic
+    current `M`.
+
+    See `meanas.fdmath` for descriptions of
+
+    - This update step: "Maxwell's equations" section
+    - `dxes`: "Datastructure: dx_lists_t" section
+    - `mu`: "Permittivity and Permeability" section
+
+    Also see the "Timestep" section of `meanas.fdtd` for a discussion of
+    the `dt` parameter.
+
+    Args:
+        dt: Timestep. See `meanas.fdtd` for details.
+        dxes: Grid description; see `meanas.fdmath`.
+
+    Returns:
+        Function `f(E_old, H_old, epsilon) -> E_new`.
+    """
+    if dxes is not None:
+        curl_e_fun = curl_forward(dxes[0])
+    else:
+        curl_e_fun = curl_forward()
+
+    def mh_fun(e: fdfield_t, h: fdfield_t, mu: fdfield_t = None):
+        """
+        Update the H-field.
+
+        Args:
+            e: E-field at time t=1
+            h: H-field at time t=0.5
+            mu: Magnetic permeability. Default is 1 everywhere.
+
+        Returns:
+            H-field at time t=1.5
+        """
+        if mu is not None:
+            h -= dt * curl_e_fun(e) / mu
+        else:
+            h -= dt * curl_e_fun(e)
+
+        return h
+
+    return mh_fun

meanas/fdtd/boundaries.py

@@ -0,0 +1,67 @@
+"""
+Boundary conditions
+
+#TODO conducting boundary documentation
+"""
+
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
+
+
+def conducting_boundary(direction: int,
+                        polarity: int
+                        ) -> Tuple[fdfield_updater_t, fdfield_updater_t]:
+    dirs = [0, 1, 2]
+    if direction not in dirs:
+        raise Exception('Invalid direction: {}'.format(direction))
+    dirs.remove(direction)
+    u, v = dirs
+
+    if polarity < 0:
+        boundary_slice = [slice(None)] * 3
+        shifted1_slice = [slice(None)] * 3
+        boundary_slice[direction] = 0
+        shifted1_slice[direction] = 1
+
+        def en(e: fdfield_t):
+            e[direction][boundary_slice] = 0
+            e[u][boundary_slice] = e[u][shifted1_slice]
+            e[v][boundary_slice] = e[v][shifted1_slice]
+            return e
+
+        def hn(h: fdfield_t):
+            h[direction][boundary_slice] = h[direction][shifted1_slice]
+            h[u][boundary_slice] = 0
+            h[v][boundary_slice] = 0
+            return h
+
+        return en, hn
+
+    if polarity > 0:
+        boundary_slice = [slice(None)] * 3
+        shifted1_slice = [slice(None)] * 3
+        shifted2_slice = [slice(None)] * 3
+        boundary_slice[direction] = -1
+        shifted1_slice[direction] = -2
+        shifted2_slice[direction] = -3
+
+        def ep(e: fdfield_t):
+            e[direction][boundary_slice] = -e[direction][shifted2_slice]
+            e[direction][shifted1_slice] = 0
+            e[u][boundary_slice] = e[u][shifted1_slice]
+            e[v][boundary_slice] = e[v][shifted1_slice]
+            return e
+
+        def hp(h: fdfield_t):
+            h[direction][boundary_slice] = h[direction][shifted1_slice]
+            h[u][boundary_slice] = -h[u][shifted2_slice]
+            h[u][shifted1_slice] = 0
+            h[v][boundary_slice] = -h[v][shifted2_slice]
+            h[v][shifted1_slice] = 0
+            return h
+
+        return ep, hp
+
+    raise Exception('Bad polarity: {}'.format(polarity))

meanas/fdtd/energy.py

@@ -0,0 +1,140 @@
+# pylint: disable=unsupported-assignment-operation
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
+from ..fdmath.functional import deriv_back, deriv_forward
+
+
+def poynting(e: fdfield_t,
+             h: fdfield_t,
+             dxes: dx_lists_t = None,
+             ) -> fdfield_t:
+    """
+    Calculate the poynting vector
+    """
+    if dxes is None:
+        dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
+
+    ex = e[0] * dxes[0][0][:, None, None]
+    ey = e[1] * dxes[0][1][None, :, None]
+    ez = e[2] * dxes[0][2][None, None, :]
+    hx = h[0] * dxes[1][0][:, None, None]
+    hy = h[1] * dxes[1][1][None, :, None]
+    hz = h[2] * dxes[1][2][None, None, :]
+
+    s = numpy.empty_like(e)
+    s[0] = numpy.roll(ey, -1, axis=0) * hz - numpy.roll(ez, -1, axis=0) * hy
+    s[1] = numpy.roll(ez, -1, axis=1) * hx - numpy.roll(ex, -1, axis=1) * hz
+    s[2] = numpy.roll(ex, -1, axis=2) * hy - numpy.roll(ey, -1, axis=2) * hx
+    return s
+
+
+def poynting_divergence(s: fdfield_t = None,
+                        *,
+                        e: fdfield_t = None,
+                        h: fdfield_t = None,
+                        dxes: dx_lists_t = None,
+                        ) -> fdfield_t:
+    """
+    Calculate the divergence of the poynting vector
+    """
+    if s is None:
+        s = poynting(e, h, dxes=dxes)
+
+    Dx, Dy, Dz = deriv_back()
+    ds = Dx(s[0]) + Dy(s[1]) + Dz(s[2])
+    return ds
+
+
+def energy_hstep(e0: fdfield_t,
+                 h1: fdfield_t,
+                 e2: fdfield_t,
+                 epsilon: fdfield_t = None,
+                 mu: fdfield_t = None,
+                 dxes: dx_lists_t = None,
+                 ) -> fdfield_t:
+    u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes)
+    return u
+
+
+def energy_estep(h0: fdfield_t,
+                 e1: fdfield_t,
+                 h2: fdfield_t,
+                 epsilon: fdfield_t = None,
+                 mu: fdfield_t = None,
+                 dxes: dx_lists_t = None,
+                 ) -> fdfield_t:
+    u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes)
+    return u
+
+
+def delta_energy_h2e(dt: float,
+                     e0: fdfield_t,
+                     h1: fdfield_t,
+                     e2: fdfield_t,
+                     h3: fdfield_t,
+                     epsilon: fdfield_t = None,
+                     mu: fdfield_t = None,
+                     dxes: dx_lists_t = None,
+                     ) -> fdfield_t:
+    """
+    This is just from (e2 * e2 + h3 * h1) - (h1 * h1 + e0 * e2)
+    """
+    de = e2 * (e2 - e0) / dt
+    dh = h1 * (h3 - h1) / dt
+    du = dxmul(de, dh, epsilon, mu, dxes)
+    return du
+
+
+def delta_energy_e2h(dt: float,
+                     h0: fdfield_t,
+                     e1: fdfield_t,
+                     h2: fdfield_t,
+                     e3: fdfield_t,
+                     epsilon: fdfield_t = None,
+                     mu: fdfield_t = None,
+                     dxes: dx_lists_t = None,
+                     ) -> fdfield_t:
+    """
+    This is just from (h2 * h2 + e3 * e1) - (e1 * e1 + h0 * h2)
+    """
+    de = e1 * (e3 - e1) / dt
+    dh = h2 * (h2 - h0) / dt
+    du = dxmul(de, dh, epsilon, mu, dxes)
+    return du
+
+
+def delta_energy_j(j0: fdfield_t, e1: fdfield_t, dxes: dx_lists_t = None) -> fdfield_t:
+    if dxes is None:
+        dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
+
+    du = ((j0 * e1).sum(axis=0) *
+          dxes[0][0][:, None, None] *
+          dxes[0][1][None, :, None] *
+          dxes[0][2][None, None, :])
+    return du
+
+
+def dxmul(ee: fdfield_t,
+          hh: fdfield_t,
+          epsilon: fdfield_t = None,
+          mu: fdfield_t = None,
+          dxes: dx_lists_t = None
+          ) -> fdfield_t:
+    if epsilon is None:
+        epsilon = 1
+    if mu is None:
+        mu = 1
+    if dxes is None:
+        dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
+
+    result = ((ee * epsilon).sum(axis=0) *
+              dxes[0][0][:, None, None] *
+              dxes[0][1][None, :, None] *
+              dxes[0][2][None, None, :] +
+              (hh * mu).sum(axis=0) *
+              dxes[1][0][:, None, None] *
+              dxes[1][1][None, :, None] *
+              dxes[1][2][None, None, :])
+    return result

meanas/fdtd/pml.py

@@ -0,0 +1,127 @@
+"""
+PML implementations
+
+#TODO discussion of PMLs
+#TODO cpml documentation
+
+"""
+# TODO retest pmls!
+
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def cpml(direction: int,
+         polarity: int,
+         dt: float,
+         epsilon: fdfield_t,
+         thickness: int = 8,
+         ln_R_per_layer: float = -1.6,
+         epsilon_eff: float = 1,
+         mu_eff: float = 1,
+         m: float = 3.5,
+         ma: float = 1,
+         cfs_alpha: float = 0,
+         dtype: numpy.dtype = numpy.float32,
+         ) -> Tuple[Callable, Callable, Dict[str, fdfield_t]]:
+
+    if direction not in range(3):
+        raise Exception('Invalid direction: {}'.format(direction))
+
+    if polarity not in (-1, 1):
+        raise Exception('Invalid polarity: {}'.format(polarity))
+
+    if thickness <= 2:
+        raise Exception('It would be wise to have a pml with 4+ cells of thickness')
+
+    if epsilon_eff <= 0:
+        raise Exception('epsilon_eff must be positive')
+
+    sigma_max = -ln_R_per_layer / 2 * (m + 1)
+    kappa_max = numpy.sqrt(epsilon_eff * mu_eff)
+    alpha_max = cfs_alpha
+    transverse = numpy.delete(range(3), direction)
+    u, v = transverse
+
+    xe = numpy.arange(1, thickness+1, dtype=float)
+    xh = numpy.arange(1, thickness+1, dtype=float)
+    if polarity > 0:
+        xe -= 0.5
+    elif polarity < 0:
+        xh -= 0.5
+        xe = xe[::-1]
+        xh = xh[::-1]
+    else:
+        raise Exception('Bad polarity!')
+
+    expand_slice = [None] * 3
+    expand_slice[direction] = slice(None)
+    expand_slice = tuple(expand_slice)
+
+    def par(x):
+        scaling = (x / thickness) ** m
+        sigma = scaling * sigma_max
+        kappa = 1 + scaling * (kappa_max - 1)
+        alpha = ((1 - x / thickness) ** ma) * alpha_max
+        p0 = numpy.exp(-(sigma / kappa + alpha) * dt)
+        p1 = sigma / (sigma + kappa * alpha) * (p0 - 1)
+        p2 = 1 / kappa
+        return p0[expand_slice], p1[expand_slice], p2[expand_slice]
+
+    p0e, p1e, p2e = par(xe)
+    p0h, p1h, p2h = par(xh)
+
+    region = [slice(None)] * 3
+    if polarity < 0:
+        region[direction] = slice(None, thickness)
+    elif polarity > 0:
+        region[direction] = slice(-thickness, None)
+    else:
+        raise Exception('Bad polarity!')
+    region = tuple(region)
+
+    se = 1 if direction == 1 else -1
+
+    # TODO check if epsilon is uniform in pml region?
+    shape = list(epsilon[0].shape)
+    shape[direction] = thickness
+    psi_e = [numpy.zeros(shape, dtype=dtype), numpy.zeros(shape, dtype=dtype)]
+    psi_h = [numpy.zeros(shape, dtype=dtype), numpy.zeros(shape, dtype=dtype)]
+
+    fields = {
+        'psi_e_u': psi_e[0],
+        'psi_e_v': psi_e[1],
+        'psi_h_u': psi_h[0],
+        'psi_h_v': psi_h[1],
+    }
+
+    # Note that this is kinda slow -- would be faster to reuse dHv*p2h for the original
+    #  H update, but then you have multiple arrays and a monolithic (field + pml) update operation
+    def pml_e(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t) -> Tuple[fdfield_t, fdfield_t]:
+        dHv = h[v][region] - numpy.roll(h[v], 1, axis=direction)[region]
+        dHu = h[u][region] - numpy.roll(h[u], 1, axis=direction)[region]
+        psi_e[0] *= p0e
+        psi_e[0] += p1e * dHv * p2e
+        psi_e[1] *= p0e
+        psi_e[1] += p1e * dHu * p2e
+        e[u][region] += se * dt / epsilon[u][region] * (psi_e[0] + (p2e - 1) * dHv)
+        e[v][region] -= se * dt / epsilon[v][region] * (psi_e[1] + (p2e - 1) * dHu)
+        return e, h
+
+    def pml_h(e: fdfield_t, h: fdfield_t) -> Tuple[fdfield_t, fdfield_t]:
+        dEv = (numpy.roll(e[v], -1, axis=direction)[region] - e[v][region])
+        dEu = (numpy.roll(e[u], -1, axis=direction)[region] - e[u][region])
+        psi_h[0] *= p0h
+        psi_h[0] += p1h * dEv * p2h
+        psi_h[1] *= p0h
+        psi_h[1] += p1h * dEu * p2h
+        h[u][region] -= se * dt * (psi_h[0] + (p2h - 1) * dEv)
+        h[v][region] += se * dt * (psi_h[1] + (p2h - 1) * dEu)
+        return e, h
+
+    return pml_e, pml_h, fields

meanas/test/__init__.py

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

meanas/test/conftest.py

@@ -0,0 +1,79 @@
+from typing import List, Tuple
+import numpy
+import pytest
+
+from .utils import PRNG
+
+#####################################
+#      Test fixtures
+#####################################
+
+@pytest.fixture(scope='module',
+                params=[(5, 5, 1),
+                        (5, 1, 5),
+                        (5, 5, 5),
+                        #(7, 7, 7),
+                       ])
+def shape(request):
+    yield (3, *request.param)
+
+
+@pytest.fixture(scope='module', params=[1.0, 1.5])
+def epsilon_bg(request):
+    yield request.param
+
+
+@pytest.fixture(scope='module', params=[1.0, 2.5])
+def epsilon_fg(request):
+    yield request.param
+
+
+@pytest.fixture(scope='module', params=['center', '000', 'random'])
+def epsilon(request, shape, epsilon_bg, epsilon_fg):
+    is3d = (numpy.array(shape) == 1).sum() == 0
+    if is3d:
+        if request.param == '000':
+            pytest.skip('Skipping 000 epsilon because test is 3D (for speed)')
+        if epsilon_bg != 1:
+            pytest.skip('Skipping epsilon_bg != 1 because test is 3D (for speed)')
+        if epsilon_fg not in (1.0, 2.0):
+            pytest.skip('Skipping epsilon_fg not in (1, 2) because test is 3D (for speed)')
+
+    epsilon = numpy.full(shape, epsilon_bg, dtype=float)
+    if request.param == 'center':
+        epsilon[:, shape[1]//2, shape[2]//2, shape[3]//2] = epsilon_fg
+    elif request.param == '000':
+        epsilon[:, 0, 0, 0] = epsilon_fg
+    elif request.param == 'random':
+        epsilon[:] = PRNG.uniform(low=min(epsilon_bg, epsilon_fg),
+                                  high=max(epsilon_bg, epsilon_fg),
+                                  size=shape)
+
+    yield epsilon
+
+
+@pytest.fixture(scope='module', params=[1.0])#, 1.5])
+def j_mag(request):
+    yield request.param
+
+
+@pytest.fixture(scope='module', params=[1.0, 1.5])
+def dx(request):
+    yield request.param
+
+
+@pytest.fixture(scope='module', params=['uniform', 'centerbig'])
+def dxes(request, shape, dx):
+    if request.param == 'uniform':
+        dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
+    elif request.param == 'centerbig':
+        dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
+        for eh in (0, 1):
+            for ax in (0, 1, 2):
+               dxes[eh][ax][dxes[eh][ax].size // 2] *= 1.1
+    elif request.param == 'random':
+        dxe = [PRNG.uniform(low=1.0 * dx, high=1.1 * dx, size=s) for s in shape[1:]]
+        dxh = [(d + numpy.roll(d, -1)) / 2 for d in dxe]
+        dxes = [dxe, dxh]
+    yield dxes
+

meanas/test/test_fdfd.py

@@ -0,0 +1,135 @@
+# pylint: disable=redefined-outer-name
+from typing import List, Tuple
+import dataclasses
+import pytest
+import numpy
+#from numpy.testing import assert_allclose, assert_array_equal
+
+from .. import fdfd
+from ..fdmath import vec, unvec
+from .utils import assert_close, assert_fields_close
+
+
+def test_residual(sim):
+    A = fdfd.operators.e_full(sim.omega, sim.dxes, vec(sim.epsilon)).tocsr()
+    b = -1j * sim.omega * vec(sim.j)
+    residual = A @ vec(sim.e) - b
+    assert numpy.linalg.norm(residual) < 1e-10
+
+
+def test_poynting_planes(sim):
+    mask = (sim.j != 0).any(axis=0)
+    if mask.sum() != 2:
+        pytest.skip(f'test_poynting_planes will only test 2-point sources, got {mask.sum()}')
+#    for dxg in sim.dxes:
+#        for dxa in dxg:
+#            if not (dxa == sim.dxes[0][0][0]).all():
+#                pytest.skip('test_poynting_planes skips nonuniform dxes')
+    points = numpy.where(mask)
+    mask[points[0][0], points[1][0], points[2][0]] = 0
+
+    mx = numpy.roll(mask, -1, axis=0)
+    my = numpy.roll(mask, -1, axis=1)
+    mz = numpy.roll(mask, -1, axis=2)
+
+    e2h = fdfd.operators.e2h(omega=sim.omega, dxes=sim.dxes, pmc=sim.pmc)
+    ev = vec(sim.e)
+    hv = e2h @ ev
+
+    exh = fdfd.operators.poynting_e_cross(e=ev, dxes=sim.dxes) @ hv.conj()
+    s = unvec(exh.real / 2, sim.shape[1:])
+    planes = [s[0, mask].sum(), -s[0, mx].sum(),
+              s[1, mask].sum(), -s[1, my].sum(),
+              s[2, mask].sum(), -s[2, mz].sum()]
+
+    e_dot_j = sim.e * sim.j * sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :]
+    src_energy = -e_dot_j[:, mask].real / 2
+
+    assert_close(sum(planes), src_energy.sum())
+
+
+#####################################
+#      Test fixtures
+#####################################
+# Also see conftest.py
+
+@pytest.fixture(params=[1/1500])
+def omega(request):
+    yield request.param
+
+
+@pytest.fixture(params=[None])
+def pec(request):
+    yield request.param
+
+
+@pytest.fixture(params=[None])
+def pmc(request):
+    yield request.param
+
+
+#@pytest.fixture(scope='module',
+#                params=[(25, 5, 5)])
+#def shape(request):
+#    yield (3, *request.param)
+
+
+@pytest.fixture(params=['diag'])        #'center'
+def j_distribution(request, shape, j_mag):
+    j = numpy.zeros(shape, dtype=complex)
+    center_mask = numpy.zeros(shape, dtype=bool)
+    center_mask[:, shape[1]//2, shape[2]//2, shape[3]//2] = True
+
+    if request.param == 'center':
+        j[center_mask] = j_mag
+    elif request.param == 'diag':
+        j[numpy.roll(center_mask, [1, 1, 1], axis=(1, 2, 3))] = j_mag
+        j[numpy.roll(center_mask, [-1, -1, -1], axis=(1, 2, 3))] = -1j * j_mag
+    yield j
+
+
+@dataclasses.dataclass()
+class FDResult:
+    shape: Tuple[int]
+    dxes: List[List[numpy.ndarray]]
+    epsilon: numpy.ndarray
+    omega: complex
+    j: numpy.ndarray
+    e: numpy.ndarray
+    pmc: numpy.ndarray
+    pec: numpy.ndarray
+
+
+@pytest.fixture()
+def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
+#    is3d = (numpy.array(shape) == 1).sum() == 0
+#    if is3d:
+#        pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)')
+
+#    # If no edge currents, add pmls
+#    src_mask = j_distribution.any(axis=0)
+#    th = 10
+#    #if src_mask.sum() - src_mask[th:-th, th:-th, th:-th].sum() == 0:
+#    if src_mask.sum() - src_mask[th:-th, :, :].sum() == 0:
+#        for axis in (0,):
+#            for polarity in (-1, 1):
+#                dxes = fdfd.scpml.stretch_with_scpml(dxes, axis=axis, polarity=polarity,
+
+    j_vec = vec(j_distribution)
+    eps_vec = vec(epsilon)
+    e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec,
+                                 matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11})
+    e = unvec(e_vec, shape[1:])
+
+    sim = FDResult(
+        shape=shape,
+        dxes=dxes,
+        epsilon=epsilon,
+        j=j_distribution,
+        e=e,
+        pec=pec,
+        pmc=pmc,
+        omega=omega,
+        )
+
+    return sim

meanas/test/test_fdfd_pml.py

@@ -0,0 +1,148 @@
+#####################################
+# pylint: disable=redefined-outer-name
+from typing import List, Tuple
+import dataclasses
+import pytest
+import numpy
+from numpy.testing import assert_allclose, assert_array_equal
+
+from .. import fdfd
+from ..fdmath import vec, unvec
+from .utils import assert_close, assert_fields_close
+from .test_fdfd import FDResult
+
+
+def test_pml(sim, src_polarity):
+    dim = numpy.where(numpy.array(sim.shape[1:]) > 1)[0][0]    # Propagation axis
+
+    e_sqr = numpy.squeeze((sim.e.conj() * sim.e).sum(axis=0))
+
+#    from matplotlib import pyplot
+#    pyplot.figure()
+#    pyplot.plot(numpy.squeeze(e_sqr))
+#    pyplot.show(block=True)
+
+    e_sqr_tgt = e_sqr[16:19]
+    e_sqr_rev = e_sqr[10:13]
+    if src_polarity < 0:
+        e_sqr_tgt, e_sqr_rev = e_sqr_rev, e_sqr_tgt
+
+    assert_allclose(e_sqr_rev, 0, atol=1e-12)
+    assert_allclose(e_sqr_tgt, 1, rtol=3e-6)
+
+
+#    pyplot.figure()
+#    pyplot.plot(numpy.squeeze(sim.e[0].real), label='Ex_real')
+#    pyplot.plot(numpy.squeeze(sim.e[0].imag), label='Ex_imag')
+#    pyplot.plot(numpy.squeeze(sim.e[1].real), label='Ey_real')
+#    pyplot.plot(numpy.squeeze(sim.e[1].imag), label='Ey_imag')
+#    pyplot.plot(numpy.squeeze(sim.e[2].real), label='Ez_real')
+#    pyplot.plot(numpy.squeeze(sim.e[2].imag), label='Ez_imag')
+#    pyplot.legend()
+#    pyplot.show(block=True)
+
+
+#      Test fixtures
+#####################################
+# Also see conftest.py
+
+@pytest.fixture(params=[1/1500])
+def omega(request):
+    yield request.param
+
+
+@pytest.fixture(params=[None])
+def pec(request):
+    yield request.param
+
+
+@pytest.fixture(params=[None])
+def pmc(request):
+    yield request.param
+
+
+
+@pytest.fixture(params=[(30, 1, 1),
+                        (1, 30, 1),
+                        (1, 1, 30)])
+def shape(request):
+    yield (3, *request.param)
+
+
+@pytest.fixture(params=[+1, -1])
+def src_polarity(request):
+    yield request.param
+
+
+@pytest.fixture()
+def j_distribution(request, shape, epsilon, dxes, omega, src_polarity):
+    j = numpy.zeros(shape, dtype=complex)
+
+    dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0]    # Propagation axis
+    other_dims = [0, 1, 2]
+    other_dims.remove(dim)
+
+    dx_prop = (dxes[0][dim][shape[dim + 1] // 2] +
+               dxes[1][dim][shape[dim + 1] // 2]) / 2       #TODO is this right for nonuniform dxes?
+
+    # Mask only contains components orthogonal to propagation direction
+    center_mask = numpy.zeros(shape, dtype=bool)
+    center_mask[other_dims, shape[1]//2, shape[2]//2, shape[3]//2] = True
+    if (epsilon[center_mask] != epsilon[center_mask][0]).any():
+        center_mask[other_dims[1]] = False          # If epsilon is not isotropic, pick only one dimension
+
+
+    wavenumber = omega * numpy.sqrt(epsilon[center_mask].mean())
+    wavenumber_corrected = 2 / dx_prop * numpy.arcsin(wavenumber * dx_prop / 2)
+
+    e = numpy.zeros_like(epsilon, dtype=complex)
+    e[center_mask] = 1 / numpy.linalg.norm(center_mask[:])
+
+    slices = [slice(None), slice(None), slice(None)]
+    slices[dim] = slice(shape[dim + 1] // 2,
+                        shape[dim + 1] // 2 + 1)
+
+    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
+
+
+@pytest.fixture()
+def epsilon(request, shape, epsilon_bg, epsilon_fg):
+    epsilon = numpy.full(shape, epsilon_fg, dtype=float)
+    yield epsilon
+
+
+@pytest.fixture(params=['uniform'])
+def dxes(request, shape, dx, omega, epsilon_fg):
+    if request.param == 'uniform':
+        dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
+    dim = numpy.where(numpy.array(shape[1:]) > 1)[0][0]    # Propagation axis
+    for axis in (dim,):
+        for polarity in (-1, 1):
+            dxes = fdfd.scpml.stretch_with_scpml(dxes, axis=axis, polarity=polarity,
+                                                 omega=omega, epsilon_effective=epsilon_fg,
+                                                 thickness=10)
+    yield dxes
+
+
+@pytest.fixture()
+def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
+    j_vec = vec(j_distribution)
+    eps_vec = vec(epsilon)
+    e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec,
+                                 matrix_solver_opts={'atol': 1e-15, 'tol': 1e-11})
+    e = unvec(e_vec, shape[1:])
+
+    sim = FDResult(
+        shape=shape,
+        dxes=dxes,
+        epsilon=epsilon,
+        j=j_distribution,
+        e=e,
+        pec=pec,
+        pmc=pmc,
+        omega=omega,
+        )
+
+    return sim

meanas/test/test_fdtd.py

@@ -0,0 +1,214 @@
+# pylint: disable=redefined-outer-name, no-member
+from typing import List, Tuple
+import dataclasses
+import pytest
+import numpy
+from numpy.testing import assert_allclose, assert_array_equal
+
+from .. import fdtd
+from .utils import assert_close, assert_fields_close, PRNG
+
+
+def test_initial_fields(sim):
+    # Make sure initial fields didn't change
+    e0 = sim.es[0]
+    h0 = sim.hs[0]
+    j0 = sim.js[0]
+    mask = (j0 != 0)
+
+    assert_fields_close(e0[mask], j0[mask] / sim.epsilon[mask])
+    assert not e0[~mask].any()
+    assert not h0.any()
+
+
+def test_initial_energy(sim):
+    """
+    Assumes fields start at 0 before J0 is added
+    """
+    j0 = sim.js[0]
+    e0 = sim.es[0]
+    h0 = sim.hs[0]
+    h1 = sim.hs[1]
+    mask = (j0 != 0)
+    dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0)
+    u0 = (j0 * j0.conj() / sim.epsilon * dV).sum(axis=0)
+    args = {'dxes': sim.dxes,
+            'epsilon': sim.epsilon}
+
+    # Make sure initial energy and E dot J are correct
+    energy0 = fdtd.energy_estep(h0=h0, e1=e0, h2=h1, **args)
+    e0_dot_j0 = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes)
+    assert_fields_close(energy0, u0)
+    assert_fields_close(e0_dot_j0, u0)
+
+
+def test_energy_conservation(sim):
+    """
+    Assumes fields start at 0 before J0 is added
+    """
+    e0 = sim.es[0]
+    j0 = sim.js[0]
+    u = fdtd.delta_energy_j(j0=j0, e1=e0, dxes=sim.dxes).sum()
+    args = {'dxes': sim.dxes,
+            'epsilon': sim.epsilon}
+
+    for ii in range(1, 8):
+        u_hstep = fdtd.energy_hstep(e0=sim.es[ii-1], h1=sim.hs[ii], e2=sim.es[ii],     **args)  # pylint: disable=bad-whitespace
+        u_estep = fdtd.energy_estep(h0=sim.hs[ii],   e1=sim.es[ii], h2=sim.hs[ii + 1], **args)  # pylint: disable=bad-whitespace
+        delta_j_A = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii-1], dxes=sim.dxes)
+        delta_j_B = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii],   dxes=sim.dxes)  # pylint: disable=bad-whitespace
+
+        u += delta_j_A.sum()
+        assert_close(u_hstep.sum(), u)
+        u += delta_j_B.sum()
+        assert_close(u_estep.sum(), u)
+
+
+def test_poynting_divergence(sim):
+    args = {'dxes': sim.dxes,
+            'epsilon': sim.epsilon}
+
+    u_eprev = None
+    for ii in range(1, 8):
+        u_hstep = fdtd.energy_hstep(e0=sim.es[ii-1], h1=sim.hs[ii], e2=sim.es[ii],     **args)  # pylint: disable=bad-whitespace
+        u_estep = fdtd.energy_estep(h0=sim.hs[ii],   e1=sim.es[ii], h2=sim.hs[ii + 1], **args)  # pylint: disable=bad-whitespace
+        delta_j_B = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii], dxes=sim.dxes)
+
+        du_half_h2e = u_estep - u_hstep - delta_j_B
+        div_s_h2e = sim.dt * fdtd.poynting_divergence(e=sim.es[ii], h=sim.hs[ii], dxes=sim.dxes)
+        assert_fields_close(du_half_h2e, -div_s_h2e)
+
+        if u_eprev is None:
+            u_eprev = u_estep
+            continue
+
+        # previous half-step
+        delta_j_A = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii-1], dxes=sim.dxes)
+
+        du_half_e2h = u_hstep - u_eprev - delta_j_A
+        div_s_e2h = sim.dt * fdtd.poynting_divergence(e=sim.es[ii-1], h=sim.hs[ii], dxes=sim.dxes)
+        assert_fields_close(du_half_e2h, -div_s_e2h)
+        u_eprev = u_estep
+
+
+def test_poynting_planes(sim):
+    mask = (sim.js[0] != 0).any(axis=0)
+    if mask.sum() > 1:
+        pytest.skip('test_poynting_planes can only test single point sources, got {}'.format(mask.sum()))
+
+    args = {'dxes': sim.dxes,
+            'epsilon': sim.epsilon}
+
+    mx = numpy.roll(mask, -1, axis=0)
+    my = numpy.roll(mask, -1, axis=1)
+    mz = numpy.roll(mask, -1, axis=2)
+
+    u_eprev = None
+    for ii in range(1, 8):
+        u_hstep = fdtd.energy_hstep(e0=sim.es[ii-1], h1=sim.hs[ii], e2=sim.es[ii],     **args)  # pylint: disable=bad-whitespace
+        u_estep = fdtd.energy_estep(h0=sim.hs[ii],   e1=sim.es[ii], h2=sim.hs[ii + 1], **args)  # pylint: disable=bad-whitespace
+        delta_j_B = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii], dxes=sim.dxes)
+        du_half_h2e = u_estep - u_hstep - delta_j_B
+
+        s_h2e = -fdtd.poynting(e=sim.es[ii], h=sim.hs[ii], dxes=sim.dxes) * sim.dt
+        planes = [s_h2e[0, mask].sum(), -s_h2e[0, mx].sum(),
+                  s_h2e[1, mask].sum(), -s_h2e[1, my].sum(),
+                  s_h2e[2, mask].sum(), -s_h2e[2, mz].sum()]
+
+        assert_close(sum(planes), du_half_h2e[mask])
+
+        if u_eprev is None:
+            u_eprev = u_estep
+            continue
+
+        delta_j_A = fdtd.delta_energy_j(j0=sim.js[ii], e1=sim.es[ii-1], dxes=sim.dxes)
+        du_half_e2h = u_hstep - u_eprev - delta_j_A
+
+        s_e2h = -fdtd.poynting(e=sim.es[ii - 1], h=sim.hs[ii], dxes=sim.dxes) * sim.dt
+        planes = [s_e2h[0, mask].sum(), -s_e2h[0, mx].sum(),
+                  s_e2h[1, mask].sum(), -s_e2h[1, my].sum(),
+                  s_e2h[2, mask].sum(), -s_e2h[2, mz].sum()]
+        assert_close(sum(planes), du_half_e2h[mask])
+
+        # previous half-step
+        u_eprev = u_estep
+
+
+#####################################
+#      Test fixtures
+#####################################
+# Also see conftest.py
+
+
+@pytest.fixture(params=[0.3])
+def dt(request):
+    yield request.param
+
+
+@dataclasses.dataclass()
+class TDResult:
+    shape: Tuple[int]
+    dt: float
+    dxes: List[List[numpy.ndarray]]
+    epsilon: numpy.ndarray
+    j_distribution: numpy.ndarray
+    j_steps: Tuple[int]
+    es: List[numpy.ndarray] = dataclasses.field(default_factory=list)
+    hs: List[numpy.ndarray] = dataclasses.field(default_factory=list)
+    js: List[numpy.ndarray] = dataclasses.field(default_factory=list)
+
+
+@pytest.fixture(params=[(0, 4, 8),]) #(0,)])
+def j_steps(request):
+    yield request.param
+
+
+@pytest.fixture(params=['center', 'random'])
+def j_distribution(request, shape, j_mag):
+    j = numpy.zeros(shape)
+    if request.param == 'center':
+        j[:, shape[1]//2, shape[2]//2, shape[3]//2] = j_mag
+    elif request.param == '000':
+        j[:, 0, 0, 0] = j_mag
+    elif request.param == 'random':
+        j[:] = PRNG.uniform(low=-j_mag, high=j_mag, size=shape)
+    yield j
+
+
+@pytest.fixture()
+def sim(request, shape, epsilon, dxes, dt, j_distribution, j_steps):
+    is3d = (numpy.array(shape) == 1).sum() == 0
+    if is3d:
+        if dt != 0.3:
+            pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)')
+
+    sim = TDResult(
+        shape=shape,
+        dt=dt,
+        dxes=dxes,
+        epsilon=epsilon,
+        j_distribution=j_distribution,
+        j_steps=j_steps,
+        )
+
+    e = numpy.zeros_like(epsilon)
+    h = numpy.zeros_like(epsilon)
+
+    assert 0 in j_steps
+    j_zeros = numpy.zeros_like(j_distribution)
+
+    eh2h = fdtd.maxwell_h(dt=dt, dxes=dxes)
+    eh2e = fdtd.maxwell_e(dt=dt, dxes=dxes)
+    for tt in range(10):
+        e = e.copy()
+        h = h.copy()
+        eh2h(e, h)
+        eh2e(e, h, epsilon)
+        if tt in j_steps:
+            e += j_distribution / epsilon
+            sim.js.append(j_distribution)
+        else:
+            sim.js.append(j_zeros)
+        sim.es.append(e)
+        sim.hs.append(h)
+    return sim

meanas/test/utils.py

@@ -0,0 +1,12 @@
+import numpy
+
+PRNG = numpy.random.RandomState(12345)
+
+def assert_fields_close(x, y, *args, **kwargs):
+    numpy.testing.assert_allclose(x, y, verbose=False,
+                err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(x, -1),
+                                                               numpy.rollaxis(y, -1)), *args, **kwargs)
+
+def assert_close(x, y, *args, **kwargs):
+    numpy.testing.assert_allclose(x, y, *args, **kwargs)
+

pdoc_templates/config.mako

@@ -0,0 +1,47 @@
+<%!
+    # 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
+    git_link_template = 'https://mpxd.net/code/jan/fdfd_tools/src/commit/{commit}/{path}#L{start_line}-L{end_line}'
+
+    # 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,445 @@
+<%
+  import os
+
+  import pdoc
+  from pdoc.html_helpers import extract_toc, glimpse, to_html as _to_html, format_git_link, _md, to_markdown
+
+  from markdown.inlinepatterns import InlineProcessor
+  from markdown.util import AtomicString, etree
+
+
+  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)
+
+
+  # Altered latex delimeters (allow inline $...$, wrap in <eq></eq>)
+  class _MathPattern(InlineProcessor):
+      NAME = 'pdoc-math'
+      PATTERN = r'(?<!\S|\\)(?:\\\((.+?)\\\)|\\\[(.+?)\\\]|\$\$(.+?)\$\$|\$(\S.*?)\$)'
+      PRIORITY = 181  # Larger than that of 'escape' pattern
+
+      def handleMatch(self, m, data):
+          for value, is_block in zip(m.groups(), (False, True, True, False)):
+              if value:
+                  break
+          wrapper = etree.Element('eq')
+          wrapper.text = AtomicString(value)
+          return wrapper, m.start(0), m.end(0)
+
+  def to_html(text: str):
+      if not latex_math and _MathPattern.NAME in _md.inlinePatterns:
+          _md.inlinePatterns.deregister(_MathPattern.NAME)
+      elif latex_math and _MathPattern.NAME not in _md.inlinePatterns:
+          _md.inlinePatterns.register(_MathPattern(_MathPattern.PATTERN),
+                                      _MathPattern.NAME,
+                                      _MathPattern.PRIORITY)
+      md = to_markdown(text, docformat='numpy,google', module=module, link=link)
+      return _md.reset().convert(md)
+
+
+#  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
+
+  <%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>

pdoc_templates/pdf.mako

@@ -0,0 +1,185 @@
+<%!
+    import re
+    import pdoc
+    from pdoc.html_helpers import to_markdown, format_git_link
+
+    def link(d, fmt='{}'):
+        name = fmt.format(d.qualname + ('()' if isinstance(d, pdoc.Function) else ''))
+        if isinstance(d, pdoc.External):
+            return name
+        return '[{}](#{})'.format(name, d.refname)
+
+    def _to_md(text, module):
+        text = to_markdown(text, module=module, link=link)
+        # Setext H2 headings to atx H2 headings
+        text = re.sub(r'\n(.+)\n-{3,}\n', r'\n## \1\n\n', text)
+        # Convert admonitions into simpler paragraphs, dedent contents
+        text = re.sub(r'^(?P<indent>( *))!!! \w+ \"([^\"]*)\"(.*(?:\n(?P=indent) +.*)*)',
+                      lambda m: '{}**{}:** {}'.format(m.group(2), m.group(3),
+                                                      re.sub('\n {,4}', '\n', m.group(4))),
+                      text, flags=re.MULTILINE)
+        return text
+
+    def subh(text, level=2):
+        # Deepen heading levels so H2 becomes H4 etc.
+        return re.sub(r'\n(#+) +(.+)\n', r'\n%s\1 \2\n' % ('#' * level), text)
+%>
+
+<%def name="title(level, string, id=None)">
+    <% id = ' {#%s}' % id if id is not None else '' %>
+${('#' * level) + ' ' + string + id}
+</%def>
+
+<%def name="funcdef(f)">
+    <%
+        returns = show_type_annotations and f.return_annotation() or ''
+        if returns:
+            returns = ' -> ' + returns
+    %>
+> `${f.funcdef()} ${f.name}(${', '.join(f.params(annotate=show_type_annotations))})${returns}`
+</%def>
+
+<%def name="classdef(c)">
+> `class ${c.name}(${', '.join(c.params(annotate=show_type_annotations))})`
+</%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) %>
+[[view code]](${git_link})
+  %endif
+</%def>
+
+---
+description: |
+    API documentation for modules: ${', '.join(m.name for m in modules)}.
+
+lang: en
+
+classoption: oneside
+geometry: margin=1in
+papersize: a4
+
+linkcolor: blue
+links-as-notes: true
+...
+% for module in modules:
+<%
+    submodules = module.submodules()
+    variables = module.variables()
+    functions = module.functions()
+    classes = module.classes()
+
+    def to_md(text):
+        return _to_md(text, module)
+%>
+
+-------------------------------------------
+
+${title(1, ('Namespace' if module.is_namespace else 'Module') + ' `%s`' % module.name, module.refname)}
+${module.docstring | to_md}
+
+% if submodules:
+${title(2, 'Sub-modules')}
+    % for m in submodules:
+* [${m.name}](#${m.refname})
+    % endfor
+% endif
+
+% if variables:
+${title(2, 'Variables')}
+    % for v in variables:
+${title(3, 'Variable `%s`' % v.name, v.refname)}
+${show_source(v)}
+${v.docstring | to_md, subh, subh}
+    % endfor
+% endif
+
+% if functions:
+${title(2, 'Functions')}
+    % for f in functions:
+${title(3, 'Function `%s`' % f.name, f.refname)}
+${show_source(f)}
+
+${funcdef(f)}
+
+${f.docstring | to_md, subh, subh}
+    % endfor
+% endif
+
+% if classes:
+${title(2, 'Classes')}
+    % for cls in classes:
+${title(3, 'Class `%s`' % cls.name, cls.refname)}
+${show_source(cls)}
+
+${classdef(cls)}
+
+${cls.docstring | to_md, subh}
+<%
+    class_vars = cls.class_variables(show_inherited_members, sort=sort_identifiers)
+    static_methods = cls.functions(show_inherited_members, sort=sort_identifiers)
+    inst_vars = cls.instance_variables(show_inherited_members, sort=sort_identifiers)
+    methods = cls.methods(show_inherited_members, sort=sort_identifiers)
+    mro = cls.mro()
+    subclasses = cls.subclasses()
+%>
+        % if mro:
+${title(4, 'Ancestors (in MRO)')}
+            % for c in mro:
+* [${c.refname}](#${c.refname})
+            % endfor
+        % endif
+
+        % if subclasses:
+${title(4, 'Descendants')}
+            % for c in subclasses:
+* [${c.refname}](#${c.refname})
+            % endfor
+        % endif
+
+        % if class_vars:
+${title(4, 'Class variables')}
+            % for v in class_vars:
+${title(5, 'Variable `%s`' % v.name, v.refname)}
+${v.docstring | to_md, subh, subh}
+            % endfor
+        % endif
+
+        % if inst_vars:
+${title(4, 'Instance variables')}
+            % for v in inst_vars:
+${title(5, 'Variable `%s`' % v.name, v.refname)}
+${v.docstring | to_md, subh, subh}
+            % endfor
+        % endif
+
+        % if static_methods:
+${title(4, 'Static methods')}
+            % for f in static_methods:
+${title(5, '`Method %s`' % f.name, f.refname)}
+
+${funcdef(f)}
+
+${f.docstring | to_md, subh, subh}
+            % endfor
+        % endif
+
+        % if methods:
+${title(4, 'Methods')}
+            % for f in methods:
+${title(5, 'Method `%s`' % f.name, f.refname)}
+
+${funcdef(f)}
+
+${f.docstring | to_md, subh, subh}
+            % endfor
+        % endif
+    % endfor
+% endif
+
+##\## for module in modules:
+% endfor
+
+-----
+Generated by *pdoc* ${pdoc.__version__} (<https://pdoc3.github.io>).

pdoc_templates/pdoc.css

@@ -0,0 +1,381 @@
+  .flex {
+    display: flex !important;
+  }
+
+  body {
+    line-height: 1.5em;
+    background: black;
+    color: #DDD;
+    max-width: 140ch;
+  }
+
+  #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;
+    border-top: 20px white;
+  }
+  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;
+    }
+    blockquote code {
+      background: #111;
+      font-weight: bold;
+      font-size: .85em;
+      padding: 5px 10px;
+      display: inline-block;
+      min-width: 40%;
+    }
+      blockquote code: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;
+    }
+
+  @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;
+    }
+  }
+
+@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;
+    }
+}

setup.py

@@ -1,18 +1,44 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 from setuptools import setup, find_packages
 
-setup(name='fdfd_tools',
-      version='0.4',
-      description='FDFD Electromagnetic simulation tools',
+with open('README.md', 'r') as f:
+    long_description = f.read()
+
+with open('meanas/VERSION', 'r') as f:
+    version = f.read().strip()
+
+setup(name='meanas',
+      version=version,
+      description='Electromagnetic simulation tools',
+      long_description=long_description,
+      long_description_content_type='text/markdown',
       author='Jan Petykiewicz',
       author_email='anewusername@gmail.com',
-      url='https://mpxd.net/gogs/jan/fdfd_tools',
+      url='https://mpxd.net/code/jan/meanas',
       packages=find_packages(),
+      package_data={
+          'meanas': ['VERSION']
+      },
       install_requires=[
             'numpy',
             'scipy',
       ],
       extras_require={
+            'test': [
+                'pytest',
+                'dataclasses',
+                ],
+            'examples': [
+                'gridlock',
+                ],
       },
+      classifiers=[
+            'Programming Language :: Python :: 3',
+            'Development Status :: 4 - Beta',
+            'Intended Audience :: Developers',
+            'Intended Audience :: Science/Research',
+            'License :: OSI Approved :: GNU Affero General Public License v3',
+            'Topic :: Scientific/Engineering :: Physics',
+      ],
       )