Remove unused waveguide_mode functions

currently unused but useful for reusing solve_wgmode arguments

README.md

@@ -1,46 +1,56 @@
-# 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
+
+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
+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.
 
+
 ## Installation
 
 **Requirements:**
-* python 3 (written and tested with 3.5)
+* python 3 (tests require 3.7)
 * numpy
 * scipy
 
 
 Install with pip, via git:
 ```bash
-pip install git+https://mpxd.net/gogs/jan/fdfd_tools.git@release
+pip install git+https://mpxd.net/code/jan/meanas.git@release
 ```
 
 ## 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,11 +2,10 @@ 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 vec, unvec, fdtd
+from meanas.fdfd import waveguide_mode, functional, scpml, operators
+from meanas.fdfd.solvers import generic as generic_solver
 
 import gridlock
 
@@ -57,18 +56,23 @@ 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,24 +117,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,
+        'epsilon': grid.grids,
     }
 
-    wg_results = waveguide_mode.solve_waveguide_mode(mode_number=0, **wg_args, epsilon=grid.grids)
-    J = waveguide_mode.compute_source(**wg_args, **wg_results)
+    wg_results = waveguide_mode.solve_waveguide_mode(mode_number=0, **wg_args)
+    J = waveguide_mode.compute_source(**wg_args, E=wg_results['E'], wavenumber=wg_results['wavenumber'])
     H_overlap = waveguide_mode.compute_overlap_e(**wg_args, **wg_results)
 
     pecg = gridlock.Grid(edge_coords, initial=0.0, num_grids=3)
@@ -141,6 +147,12 @@ 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()
+
     '''
     Solve!
     '''
@@ -155,7 +167,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,27 +175,21 @@ 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.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
+        h = operators.e2h(omega, dxes) @ e
+        cross1 = operators.poynting_e_cross(e, dxes) @ h.conj()
+        cross2 = 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
@@ -209,7 +215,7 @@ def module_available(name):
 
 
 if __name__ == '__main__':
-    # test0()
+    #test0()
 
     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

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/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

meanas/__init__.py

@@ -0,0 +1,48 @@
+"""
+Electromagnetic simulation 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**
+- 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
+```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) 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.
+
+
+Dependencies:
+- numpy
+- scipy
+
+"""
+
+from .types import dx_lists_t, field_t, vfield_t, field_updater
+from .vectorization import vec, unvec
+
+__author__ = 'Jan Petykiewicz'
+version = '0.5'

meanas/fdfd/bloch.py

@@ -73,21 +73,46 @@ This module contains functions for generating and solving the
 
 '''
 
-from typing import List, Tuple, Callable, Dict
+from typing import Tuple, Callable
 import logging
 import numpy
-from numpy.fft import fftn, ifftn, fftfreq
+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 .eigensolvers import rayleigh_quotient_iteration
 from . import field_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
@@ -255,7 +280,7 @@ def hmn_2_hxyz(k0: numpy.ndarray,
     :return: Function for converting h_mn into H_xyz
     """
     shape = epsilon[0].shape + (1,)
-    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+    _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)]
@@ -329,147 +354,14 @@ def inverse_maxwell_operator_approx(k0: numpy.ndarray,
         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
+        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 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,
@@ -479,6 +371,7 @@ def find_k(frequency: float,
            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.
@@ -499,8 +392,10 @@ def find_k(frequency: float,
 
     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)
+        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),
@@ -511,3 +406,244 @@ def find_k(frequency: float,
     return res.x * direction, res.fun + frequency
 
 
+def eigsolve(num_modes: int,
+             k0: numpy.ndarray,
+             G_matrix: numpy.ndarray,
+             epsilon: field_t,
+             mu: field_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.
+
+    :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).
+    :param tolerance: Solver stops when fractional change in the objective
+        trace(Z.H @ A @ Z @ inv(Z Z.H)) is smaller than the tolerance
+    :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)
+
+    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

@@ -6,9 +6,11 @@ import numpy
 from numpy.fft import fft2, fftshift, fftfreq, ifft2, ifftshift
 from numpy import pi
 
+from .. import field_t
 
-def near_to_farfield(E_near: List[numpy.ndarray],
-                     H_near: List[numpy.ndarray],
+
+def near_to_farfield(E_near: field_t,
+                     H_near: field_t,
                      dx: float,
                      dy: float,
                      padded_size: List[int] = None
@@ -115,8 +117,8 @@ 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: field_t,
+                     H_far: field_t,
                      dkx: float,
                      dky: float,
                      padded_size: List[int] = None

meanas/fdfd/functional.py

@@ -2,13 +2,13 @@
 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.
+The functions generated here expect field 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
 import numpy
 
-from . import dx_lists_t, field_t
+from .. import dx_lists_t, field_t
 
 __author__ = 'Jan Petykiewicz'
 
@@ -20,7 +20,7 @@ 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
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Function for taking the discretized curl of the H-field, F(H) -> curlH
     """
     dxyz_b = numpy.meshgrid(*dxes[1], indexing='ij')
@@ -29,9 +29,13 @@ def curl_h(dxes: dx_lists_t) -> functional_matrix:
         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)]
+        e = numpy.empty_like(h)
+        e[0] = dh(h[2], 1)
+        e[0] -= dh(h[1], 2)
+        e[1] = dh(h[0], 2)
+        e[1] -= dh(h[2], 0)
+        e[2] = dh(h[1], 0)
+        e[2] -= dh(h[0], 1)
         return e
 
     return ch_fun
@@ -41,7 +45,7 @@ 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
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Function for taking the discretized curl of the E-field, F(E) -> curlE
     """
     dxyz_a = numpy.meshgrid(*dxes[0], indexing='ij')
@@ -50,9 +54,13 @@ def curl_e(dxes: dx_lists_t) -> functional_matrix:
         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)]
+        h = numpy.empty_like(e)
+        h[0] = de(e[2], 1)
+        h[0] -= de(e[1], 2)
+        h[1] = de(e[0], 2)
+        h[1] -= de(e[2], 0)
+        h[2] = de(e[1], 0)
+        h[2] -= de(e[0], 1)
         return h
 
     return ce_fun
@@ -69,7 +77,7 @@ def e_full(omega: complex,
     (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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :param epsilon: Dielectric constant
     :param mu: Magnetic permeability (default 1 everywhere)
     :return: Function implementing the wave operator A(E) -> E
@@ -79,11 +87,11 @@ def e_full(omega: complex,
 
     def op_1(e):
         curls = ch(ce(e))
-        return [c - omega ** 2 * e * x for c, e, x in zip(curls, epsilon, e)]
+        return curls - omega ** 2 * 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)]
+        curls = ch(mu * ce(e))
+        return curls - omega ** 2 * epsilon * e
 
     if numpy.any(numpy.equal(mu, None)):
         return op_1
@@ -100,7 +108,7 @@ def eh_full(omega: complex,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :param epsilon: Dielectric constant
     :param mu: Magnetic permeability (default 1 everywhere)
     :return: Function implementing the wave operator A(E, H) -> (E, H)
@@ -109,12 +117,12 @@ def eh_full(omega: complex,
     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)])
+        return (ch(h) - 1j * omega * epsilon * e,
+                ce(e) + 1j * omega * 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)])
+        return (ch(h) - 1j * omega * epsilon * e,
+                ce(e) + 1j * omega * mu * h)
 
     if numpy.any(numpy.equal(mu, None)):
         return op_1
@@ -131,19 +139,68 @@ def e2h(omega: complex,
    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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
    :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)]
+        return A2(e) / (-1j * omega)
 
     def e2h_mu(e):
-        return [y / (-1j * omega * m) for y, m in zip(A2(e), mu)]
+        return A2(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: field_t = None,
+        ) -> functional_matrix:
+    """
+   Utility operator for converting magnetic current (M) distribution
+   into equivalent electric current distribution (J).
+   For use with e.g. e_full().
+
+   :param omega: Angular frequency of the simulation
+   :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
+   :param mu: Magnetic permeability (default 1 everywhere)
+   :return: Function for converting M to J
+   """
+    ch = curl_h(dxes)
+
+    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: field_t,
+                  omega: complex,
+                  dxes: dx_lists_t,
+                  epsilon: field_t,
+                  mu: field_t = None,
+                  ) -> functional_matrix:
+    """
+    Operator that turuns an E-field distribution into a total-field/scattered-field
+    (TFSF) source.
+    """
+    # 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)
+

meanas/fdfd/operators.py

@@ -3,17 +3,13 @@ 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).
+ meanas.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.
+Many of these functions require a 'dxes' parameter, of type meanas.dx_lists_type; see
+the meanas.types submodule for details.
 
 
 The following operators are included:
@@ -36,7 +32,7 @@ from typing import List, Tuple
 import numpy
 import scipy.sparse as sparse
 
-from . import vec, dx_lists_t, vfield_t
+from .. import vec, dx_lists_t, vfield_t
 
 
 __author__ = 'Jan Petykiewicz'
@@ -57,7 +53,7 @@ def e_full(omega: complex,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :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
@@ -101,7 +97,7 @@ def e_full_preconditioners(dxes: dx_lists_t
 
     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
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Preconditioner matrices (Pl, Pr)
     """
     p_squared = [dxes[0][0][:, None, None] * dxes[1][1][None, :, None] * dxes[1][2][None, None, :],
@@ -127,7 +123,7 @@ def h_full(omega: complex,
     (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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :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
@@ -177,7 +173,7 @@ def eh_full(omega: complex,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :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
@@ -216,7 +212,7 @@ 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
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Sparse matrix for taking the discretized curl of the H-field
     """
     return cross(deriv_back(dxes[1]))
@@ -226,7 +222,7 @@ 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
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Sparse matrix for taking the discretized curl of the E-field
     """
     return cross(deriv_forward(dxes[0]))
@@ -242,7 +238,7 @@ def e2h(omega: complex,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :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).
@@ -270,7 +266,7 @@ def m2j(omega: complex,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :param mu: Vectorized magnetic permeability (default 1 everywhere)
     :return: Sparse matrix for converting E to H
     """
@@ -345,9 +341,7 @@ def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> spa
 
     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]
+    j_ind = numpy.ravel_multi_index(ijk, shape, order='C')
 
     vij = (numpy.ones(n), (i_ind, j_ind.ravel(order='C')))
 
@@ -451,10 +445,10 @@ def avgb(axis: int, shape: List[int]) -> sparse.spmatrix:
 
 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.
+    Operator for computing the Poynting vector, containing the (E x) portion of the Poynting vector.
 
     :param e: Vectorized E-field for the ExH cross product
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Sparse matrix containing (E x) portion of Poynting cross product
     """
     shape = [len(dx) for dx in dxes[0]]
@@ -483,7 +477,7 @@ 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
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :return: Sparse matrix containing (H x) portion of Poynting cross product
     """
     shape = [len(dx) for dx in dxes[0]]
@@ -505,3 +499,50 @@ def poynting_h_cross(h: vfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
          [ bx @ Hz @ fy @ dagx,  zero,                -bz @ Hx @ fy @ dagz],
          [-bx @ Hy @ fz @ dagx,  by @ Hx @ fz @ dagy,  zero]])
     return P
+
+
+def e_tfsf_source(TF_region: vfield_t,
+                  omega: complex,
+                  dxes: dx_lists_t,
+                  epsilon: vfield_t,
+                  mu: vfield_t = None,
+                  ) -> sparse.spmatrix:
+    """
+    Operator that turns an E-field distribution into a total-field/scattered-field
+    (TFSF) source.
+    """
+    # 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: vfield_t,
+                      omega: complex,
+                      dxes: dx_lists_t,
+                      epsilon: vfield_t,
+                      mu: vfield_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.
+    """
+    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):
+        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))
+
+    return sparse.diags(jmask.astype(int)) @ full

meanas/fdfd/scpml.py

@@ -5,10 +5,11 @@ Functions for creating stretched coordinate PMLs.
 from typing import List, Callable
 import numpy
 
+from .. import dx_lists_t
+
 __author__ = 'Jan Petykiewicz'
 
 
-dx_lists_t = List[List[numpy.ndarray]]
 s_function_type = Callable[[float], float]
 
 

meanas/fdfd/solvers.py

@@ -70,7 +70,7 @@ def generic(omega: complex,
     """
     Conjugate gradient FDFD solver using CSR sparse matrices.
 
-    All ndarray arguments should be 1D array, as returned by fdfd_tools.vec().
+    All ndarray arguments should be 1D array, as returned by meanas.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)

meanas/fdfd/waveguide.py

@@ -17,24 +17,50 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
  (ie. dxes = [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...]], [[dx_h_0, ...], [dy_h_0, ...]]])
  with propagation along the z axis.
 """
+# TODO update module docs
 
 from typing import List, Tuple
 import numpy
 from numpy.linalg import norm
 import scipy.sparse as sparse
 
-from . import vec, unvec, dx_lists_t, field_t, vfield_t
+from .. import vec, unvec, dx_lists_t, field_t, vfield_t
 from . import operators
 
 
 __author__ = 'Jan Petykiewicz'
 
 
-def operator(omega: complex,
+def operator_e(omega: complex,
              dxes: dx_lists_t,
              epsilon: vfield_t,
              mu: vfield_t = None,
              ) -> sparse.spmatrix:
+    if numpy.any(numpy.equal(mu, None)):
+        mu = numpy.ones_like(epsilon)
+
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    Dbx, Dby = operators.deriv_back(dxes[1])
+
+    eps_parts = numpy.split(epsilon, 3)
+    eps_xy = sparse.diags(numpy.hstack((eps_parts[0], eps_parts[1])))
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    mu_parts = numpy.split(mu, 3)
+    mu_yx = sparse.diags(numpy.hstack((mu_parts[1], mu_parts[0])))
+    mu_z_inv = sparse.diags(1 / mu_parts[2])
+
+    op = omega * omega * mu_yx @ eps_xy + \
+        mu_yx @ sparse.vstack((-Dby, Dbx)) @ mu_z_inv @ sparse.hstack((-Dfy, Dfx)) + \
+        sparse.vstack((Dfx, Dfy)) @ eps_z_inv @ sparse.hstack((Dbx, Dby)) @ eps_xy
+    return op
+
+
+def operator_h(omega: complex,
+               dxes: dx_lists_t,
+               epsilon: vfield_t,
+               mu: vfield_t = None,
+               ) -> sparse.spmatrix:
     """
     Waveguide operator of the form
 
@@ -51,7 +77,7 @@ def operator(omega: complex,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
     :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
     :return: Sparse matrix representation of the operator
@@ -70,51 +96,101 @@ def operator(omega: complex,
     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 + \
+    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(v: numpy.ndarray,
-                      wavenumber: complex,
-                      omega: complex,
-                      dxes: dx_lists_t,
-                      epsilon: vfield_t,
-                      mu: vfield_t = None
-                      ) -> Tuple[vfield_t, vfield_t]:
+def normalized_fields_e(e_xy: numpy.ndarray,
+                        wavenumber: complex,
+                        omega: complex,
+                        dxes: dx_lists_t,
+                        epsilon: vfield_t,
+                        mu: vfield_t = None,
+                        dx_prop: float = 0,
+                        ) -> Tuple[vfield_t, vfield_t]:
     """
-    Given a vector v containing the vectorized H_x and H_y fields,
+    Given a vector e_xy containing the vectorized E_x and E_y fields,
      returns normalized, vectorized E and H fields for the system.
 
-    :param v: Vector containing H_x and H_y fields
-    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param e_xy: Vector containing E_x and E_y fields
+    :param wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
     :param omega: The angular frequency of the system
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
     :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :param dxes_prop: Grid cell width in the propagation direction. Default 0 (continuous).
     :return: Normalized, vectorized (e, h) containing all vector components.
     """
-    e = v2e(v, wavenumber, omega, dxes, epsilon, mu=mu)
-    h = v2h(v, wavenumber, dxes, mu=mu)
+    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, wavenumber=wavenumber, omega=omega,
+                                        dxes=dxes, epsilon=epsilon, mu=mu, dx_prop=dx_prop)
+    return e_norm, h_norm
 
+
+def normalized_fields_h(h_xy: numpy.ndarray,
+                        wavenumber: complex,
+                        omega: complex,
+                        dxes: dx_lists_t,
+                        epsilon: vfield_t,
+                        mu: vfield_t = None,
+                        dx_prop: float = 0,
+                        ) -> Tuple[vfield_t, vfield_t]:
+    """
+    Given a vector e_xy containing the vectorized E_x and E_y fields,
+     returns normalized, vectorized E and H fields for the system.
+
+    :param e_xy: Vector containing E_x and E_y fields
+    :param wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :param dxes_prop: Grid cell width in the propagation direction. Default 0 (continuous).
+    :return: Normalized, vectorized (e, h) containing all vector components.
+    """
+    e = hxy2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ h_xy
+    h = hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu) @ h_xy
+    e_norm, h_norm = _normalized_fields(e=e, h=h, wavenumber=wavenumber, omega=omega,
+                                        dxes=dxes, epsilon=epsilon, mu=mu, dx_prop=dx_prop)
+    return e_norm, h_norm
+
+
+def _normalized_fields(e: numpy.ndarray,
+                       h: numpy.ndarray,
+                       wavenumber: complex,
+                       omega: complex,
+                       dxes: dx_lists_t,
+                       epsilon: vfield_t,
+                       mu: vfield_t = None,
+                       dx_prop: float = 0,
+                       ) -> Tuple[vfield_t, vfield_t]:
+    # TODO documentation
     shape = [s.size for s in dxes[0]]
     dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)]
 
     E = unvec(e, shape)
     H = unvec(h, shape)
 
-    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())
+    phase = numpy.exp(-1j * wavenumber * dx_prop / 2)
+    S1 = E[0] * numpy.conj(H[1] * phase) * dxes_real[0][1] * dxes_real[1][0]
+    S2 = E[1] * numpy.conj(H[0] * phase) * dxes_real[0][0] * dxes_real[1][1]
+    P = numpy.real(S1.sum() - S2.sum())
     assert P > 0, 'Found a mode propagating in the wrong direction! P={}'.format(P)
 
+    energy = epsilon * e.conj() * e
+
     norm_amplitude = 1 / numpy.sqrt(P)
-    norm_angle = -numpy.angle(e[e.size//2])
-    norm_factor = norm_amplitude * numpy.exp(1j * norm_angle)
+    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]
+    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
@@ -122,56 +198,104 @@ def normalized_fields(v: numpy.ndarray,
     return e, h
 
 
-def v2h(v: numpy.ndarray,
-        wavenumber: complex,
-        dxes: dx_lists_t,
-        mu: vfield_t = None
-        ) -> vfield_t:
+def exy2h(wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> sparse.spmatrix:
     """
-    Given a vector v containing the vectorized H_x and H_y fields,
-     returns a vectorized H including all three H components.
+    Operator which transforms the vector e_xy containing the vectorized E_x and E_y fields,
+     into a vectorized H containing all three H components
 
-    :param 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 wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
+    :param epsilon: Vectorized dielectric constant grid
     :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Vectorized H field with all vector components
+    :return: Sparse matrix representing the operator
+    """
+    e2hop = e2h(wavenumber=wavenumber, omega=omega, dxes=dxes, mu=mu)
+    return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)
+
+
+def hxy2e(wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> sparse.spmatrix:
+    """
+    Operator which transforms the vector h_xy containing the vectorized H_x and H_y fields,
+     into a vectorized E containing all three E components
+
+    :param wavenumber: Wavenumber satisfying `operator_h(...) @ h_xy == wavenumber**2 * h_xy`
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Sparse matrix representing the operator
+    """
+    h2eop = h2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon)
+    return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)
+
+
+def hxy2h(wavenumber: complex,
+          dxes: dx_lists_t,
+          mu: vfield_t = None
+          ) -> sparse.spmatrix:
+    """
+    Operator which transforms the vector h_xy containing the vectorized H_x and H_y fields,
+     into a vectorized H containing all three H components
+
+    :param wavenumber: Wavenumber satisfying `operator_h(...) @ h_xy == wavenumber**2 * h_xy`
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Sparse matrix representing the operator
     """
     Dfx, Dfy = operators.deriv_forward(dxes[0])
-    op = sparse.hstack((Dfx, Dfy))
+    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])
 
-        op = mu_z_inv @ op @ mu_xy
+        hxy2hz = mu_z_inv @ hxy2hz @ mu_xy
 
-    w = op @ v / (1j * wavenumber)
-    return numpy.hstack((v, w)).flatten()
+    n_pts = dxes[1][0].size * dxes[1][1].size
+    op = sparse.vstack((sparse.eye(2 * n_pts),
+                        hxy2hz))
+    return op
 
 
-def v2e(v: numpy.ndarray,
-        wavenumber: complex,
-        omega: complex,
-        dxes: dx_lists_t,
-        epsilon: vfield_t,
-        mu: vfield_t = None
-        ) -> vfield_t:
+def exy2e(wavenumber: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          ) -> sparse.spmatrix:
     """
-    Given a vector v containing the vectorized H_x and H_y fields,
-     returns a vectorized E containing all three E components
+    Operator which transforms the vector e_xy containing the vectorized E_x and E_y fields,
+     into a vectorized E containing all three E components
 
-    :param 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 wavenumber: Wavenumber satisfying `operator_e(...) @ e_xy == wavenumber**2 * e_xy`
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
-    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
-    :return: Vectorized E field with all vector components.
+    :return: Sparse matrix representing the operator
     """
-    h2eop = h2e(wavenumber, omega, dxes, epsilon)
-    return h2eop @ v2h(v, wavenumber, dxes, mu)
+    Dbx, Dby = operators.deriv_back(dxes[1])
+    exy2ez = sparse.hstack((Dbx, Dby)) / (1j * wavenumber)
+
+    if not numpy.any(numpy.equal(epsilon, None)):
+        epsilon_parts = numpy.split(epsilon, 3)
+        epsilon_xy = sparse.diags(numpy.hstack((epsilon_parts[0], epsilon_parts[1])))
+        epsilon_z_inv = sparse.diags(1 / epsilon_parts[2])
+
+        exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy
+
+    n_pts = dxes[0][0].size * dxes[0][1].size
+    op = sparse.vstack((sparse.eye(2 * n_pts),
+                        exy2ez))
+    return op
 
 
 def e2h(wavenumber: complex,
@@ -185,7 +309,7 @@ def e2h(wavenumber: complex,
 
     :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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
     :return: Sparse matrix representation of the operator
     """
@@ -206,7 +330,7 @@ def h2e(wavenumber: complex,
 
     :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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
     :return: Sparse matrix representation of the operator
     """
@@ -219,7 +343,7 @@ 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)
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :return: Sparse matrix representation of the operator
     """
     n = 1
@@ -236,7 +360,7 @@ 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)
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :return: Sparse matrix representation of the operator
     """
     n = 1
@@ -261,7 +385,7 @@ def h_err(h: vfield_t,
     :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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
     :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
     :return: Relative error norm(OP @ h) / norm(h)
@@ -292,7 +416,7 @@ def e_err(e: vfield_t,
     :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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
     :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
     :return: Relative error norm(OP @ e) / norm(e)
@@ -328,7 +452,7 @@ def cylindrical_operator(omega: complex,
      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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types (2D)
     :param epsilon: Vectorized dielectric constant grid
     :param r0: Radius of curvature for the simulation. This should be the minimum value of
         r within the simulation domain.

meanas/fdfd/waveguide_mode.py

@@ -2,63 +2,53 @@ 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 vec, unvec, dx_lists_t, vfield_t, field_t
 from . import operators, waveguide, functional
-from .eigensolvers import signed_eigensolve, rayleigh_quotient_iteration
+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,
+                            epsilon: field_t,
+                            mu: field_t = None,
+                            mode_margin: int = 2,
                             ) -> 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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :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}
+    :param mode_margin: The eigensolver will actually solve for (mode_number + mode_margin)
+        modes, but only return the target mode. Increasing this value can improve the solver's
+        ability to find the correct mode. Default 2.
+    :return: {'E': numpy.ndarray, 'H': 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))
+    A_r = waveguide.operator_e(numpy.real(omega), dxes_real, vec(numpy.real(epsilon)), vec(numpy.real(mu)))
 
-    eigvals, eigvecs = signed_eigensolve(A_r, mode_number+3)
-    v = eigvecs[:, -(mode_number + 1)]
+    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + mode_margin)
+    exy = 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)
+    A = waveguide.operator_e(omega, dxes, vec(epsilon), vec(mu))
+    eigval, exy = rayleigh_quotient_iteration(A, exy)
 
     # 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)
+    e, h = waveguide.normalized_fields_e(exy, wavenumber, omega, dxes, vec(epsilon), vec(mu))
 
     shape = [d.size for d in dxes[0]]
     fields = {
@@ -78,7 +68,6 @@ def solve_waveguide_mode(mode_number: 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
@@ -86,19 +75,19 @@ def solve_waveguide_mode(mode_number: int,
 
     :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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :param axis: Propagation axis (0=x, 1=y, 2=z)
     :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
     :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
         as the waveguide cross-section. slices[axis] should select only one
     :param epsilon: Dielectric constant
     :param mu: Magnetic permeability (default 1 everywhere)
-    :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
+        mu = numpy.ones_like(epsilon)
+
+    slices = tuple(slices)
 
     '''
     Solve the 2D problem in the specified plane
@@ -107,57 +96,54 @@ def solve_waveguide_mode(mode_number: int,
     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)
+
     # 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,
+        'epsilon': [epsilon[i][slices].transpose(order) for i in order],
+        'mu': [mu[i][slices].transpose(order) for i in order],
     }
     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])
+    # Correct wavenumber to account for numerical dispersion.
+    fields_2d['wavenumber'] = 2/dx_prop * numpy.arcsin(fields_2d['wavenumber'] * dx_prop/2)
 
     # Adjust for propagation direction
-    fields_2d['E'][2] *= polarity
-    fields_2d['H'][2] *= polarity
+    fields_2d['H'] *= 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)
+    fields_2d['H'][:2] *= numpy.exp(-1j * polarity * 0.5 * fields_2d['wavenumber'] * dx_prop)
+    fields_2d['E'][2]  *= numpy.exp(-1j * polarity * 0.5 * fields_2d['wavenumber'] * dx_prop)
 
     # Expand E, H to full epsilon space we were given
-    E = [None]*3
-    H = [None]*3
+    E = numpy.zeros_like(epsilon, dtype=complex)
+    H = numpy.zeros_like(epsilon, dtype=complex)
     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)
+        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],
+                   epsilon: field_t,
                    mu: field_t = None,
                    ) -> field_t:
     """
@@ -165,10 +151,9 @@ def compute_source(E: field_t,
     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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :param axis: Propagation axis (0=x, 1=y, 2=z)
     :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
     :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
@@ -176,28 +161,20 @@ def compute_source(E: field_t,
     :param mu: Magnetic permeability (default 1 everywhere)
     :return: J distribution for the unidirectional source
     """
-    if mu is None:
-        mu = [1] * 3
+    E_expanded = expand_wgmode_e(E=E, dxes=dxes, wavenumber=wavenumber, axis=axis,
+                                 polarity=polarity, slices=slices)
 
-    J = [None]*3
-    M = [None]*3
+    smask = [slice(None)] * 4
+    if polarity > 0:
+        smask[axis + 1] = slice(slices[axis].start, None)
+    else:
+        smask[axis + 1] = slice(None, slices[axis].stop)
 
-    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)
+    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
 
 
@@ -209,6 +186,7 @@ def compute_overlap_e(E: field_t,
                       axis: int,
                       polarity: int,
                       slices: List[slice],
+                      epsilon: field_t,         # TODO unused??
                       mu: field_t = None,
                       ) -> field_t:
     """
@@ -223,7 +201,7 @@ def compute_overlap_e(E: field_t,
     :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 dxes: Grid parameters [dx_e, dx_h] as described in meanas.types
     :param axis: Propagation axis (0=x, 1=y, 2=z)
     :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
     :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
@@ -231,8 +209,11 @@ def compute_overlap_e(E: field_t,
     :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]
+    slices = tuple(slices)
+
+    cross_plane = [slice(None)] * 4
+    cross_plane[axis + 1] = slices[axis]
+    cross_plane = tuple(cross_plane)
 
     # Determine phase factors for parallel slices
     a_shape = numpy.roll([-1, 1, 1], axis)
@@ -243,11 +224,8 @@ def compute_overlap_e(E: field_t,
     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)]
+    Ee = phase_E * E[cross_plane]
+    He = phase_H * H[cross_plane]
 
 
     # Write out the operator product for the mode orthogonality integral
@@ -279,21 +257,19 @@ def solve_waveguide_mode_cylindrical(mode_number: int,
                                      dxes: dx_lists_t,
                                      epsilon: vfield_t,
                                      r0: float,
-                                     wavenumber_correction: bool = True,
                                      ) -> Dict[str, complex or field_t]:
     """
+    TODO: fixup
     Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
      of the bent waveguide with the specified mode number.
 
     :param mode_number: Number of the mode, 0-indexed
     :param omega: Angular frequency of the simulation
-    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header.
+    :param dxes: Grid parameters [dx_e, dx_h] as described in meanas.types.
         The first coordinate is assumed to be r, the second is y.
     :param epsilon: Dielectric constant
     :param r0: Radius of curvature for the simulation. This should be the minimum value of
         r within the simulation domain.
-    :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}
     """
 
@@ -317,16 +293,7 @@ def solve_waveguide_mode_cylindrical(mode_number: int,
     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)
+    # TODO: Perform correction on wavenumber to account for numerical dispersion.
 
     shape = [d.size for d in dxes[0]]
     v = numpy.hstack((v, numpy.zeros(shape[0] * shape[1])))
@@ -338,3 +305,60 @@ def solve_waveguide_mode_cylindrical(mode_number: int,
     }
 
     return fields
+
+
+def compute_overlap_ce(E: field_t,
+                       wavenumber: complex,
+                       dxes: dx_lists_t,
+                       axis: int,
+                       polarity: int,
+                       slices: List[slice],
+                       ) -> field_t:
+    slices = tuple(slices)
+
+    Ee = expand_wgmode_e(E=E, wavenumber=wavenumber,
+                         dxes=dxes, axis=axis, polarity=polarity,
+                         slices=slices)
+
+    start, stop = sorted((slices[axis].start, slices[axis].start - 2 * polarity))
+
+    slices2 = list(slices)
+    slices2[axis] = slice(start, stop)
+    slices2 = (slice(None), *slices2)
+
+    Etgt = numpy.zeros_like(Ee)
+    Etgt[slices2] = Ee[slices2]
+
+    Etgt /= (Etgt.conj() * Etgt).sum()
+    return Etgt, slices2
+
+
+def expand_wgmode_e(E: field_t,
+                    wavenumber: complex,
+                    dxes: dx_lists_t,
+                    axis: int,
+                    polarity: int,
+                    slices: List[slice],
+                    ) -> field_t:
+    slices = tuple(slices)
+
+    # Determine phase factors for parallel slices
+    a_shape = numpy.roll([1, -1, 1, 1], axis)
+    a_E = numpy.real(dxes[0][axis]).cumsum()
+    r_E = a_E - a_E[slices[axis]]
+    iphi = polarity * -1j * wavenumber
+    phase_E = numpy.exp(iphi * r_E).reshape(a_shape)
+
+    # Expand our slice to the entire grid using the phase factors
+    E_expanded = numpy.zeros_like(E)
+
+    slices_exp = list(slices)
+    slices_exp[axis] = slice(E.shape[axis + 1])
+    slices_exp = (slice(None), *slices_exp)
+
+    slices_in = (slice(None), *slices)
+
+    E_expanded[slices_exp] = phase_E * numpy.array(E)[slices_in]
+    return E_expanded
+
+

meanas/fdtd/__init__.py

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

meanas/fdtd/boundaries.py

@@ -0,0 +1,68 @@
+"""
+Boundary conditions
+"""
+
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from .. import dx_lists_t, field_t, field_updater
+
+
+def conducting_boundary(direction: int,
+                        polarity: int
+                        ) -> Tuple[field_updater, field_updater]:
+    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))
+
+

meanas/fdtd/energy.py

@@ -0,0 +1,84 @@
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from .. import dx_lists_t, field_t, field_updater
+
+
+def poynting(e, h):
+    s = (numpy.roll(e[1], -1, axis=0) * h[2] - numpy.roll(e[2], -1, axis=0) * h[1],
+         numpy.roll(e[2], -1, axis=1) * h[0] - numpy.roll(e[0], -1, axis=1) * h[2],
+         numpy.roll(e[0], -1, axis=2) * h[1] - numpy.roll(e[1], -1, axis=2) * h[0])
+    return numpy.array(s)
+
+
+def poynting_divergence(s=None, *, e=None, h=None, dxes=None): # TODO dxes
+    if dxes is None:
+        dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
+
+    if s is None:
+        s = poynting(e, h)
+
+    ds = ((s[0] - numpy.roll(s[0], 1, axis=0)) / numpy.sqrt(dxes[0][0] * dxes[1][0])[:, None, None] +
+          (s[1] - numpy.roll(s[1], 1, axis=1)) / numpy.sqrt(dxes[0][1] * dxes[1][1])[None, :, None] +
+          (s[2] - numpy.roll(s[2], 1, axis=2)) / numpy.sqrt(dxes[0][2] * dxes[1][2])[None, None, :] )
+    return ds
+
+
+def energy_hstep(e0, h1, e2, epsilon=None, mu=None, dxes=None):
+    u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes)
+    return u
+
+
+def energy_estep(h0, e1, h2, epsilon=None, mu=None, dxes=None):
+    u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes)
+    return u
+
+
+def delta_energy_h2e(dt, e0, h1, e2, h3, epsilon=None, mu=None, dxes=None):
+    """
+    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, h0, e1, h2, e3, epsilon=None, mu=None, dxes=None):
+    """
+    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, e1, dxes=None):
+    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, hh, epsilon=None, mu=None, dxes=None):
+    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,122 @@
+"""
+PML implementations
+
+"""
+# TODO retest pmls!
+
+from typing import List, Callable, Tuple, Dict
+import numpy
+
+from .. import dx_lists_t, field_t, field_updater
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def cpml(direction:int,
+         polarity: int,
+         dt: float,
+         epsilon: field_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, 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')
+
+    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)
+
+    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!')
+
+    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: field_t, h: field_t, epsilon: field_t) -> Tuple[field_t, field_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: field_t, h: field_t) -> Tuple[field_t, field_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/test_fdtd.py

@@ -0,0 +1,293 @@
+import numpy
+import pytest
+import dataclasses
+from typing import List, Tuple
+from numpy.testing import assert_allclose, assert_array_equal
+
+from meanas import fdtd
+
+
+prng = numpy.random.RandomState(12345)
+
+
+def assert_fields_close(a, b, *args, **kwargs):
+    numpy.testing.assert_allclose(a, b, verbose=False, err_msg='Fields did not match:\n{}\n{}'.format(numpy.rollaxis(a, -1),
+                                                                                                      numpy.rollaxis(b, -1)), *args, **kwargs)
+
+def assert_close(a, b, *args, **kwargs):
+    numpy.testing.assert_allclose(a, b, *args, **kwargs)
+
+
+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)
+        u_estep = fdtd.energy_estep(h0=sim.hs[ii],   e1=sim.es[ii], h2=sim.hs[ii + 1], **args)
+        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)
+
+        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}
+
+    dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0)
+
+    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)
+        u_estep = fdtd.energy_estep(h0=sim.hs[ii],   e1=sim.es[ii], h2=sim.hs[ii + 1], **args)
+        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) * dV
+        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) * dV
+        assert_fields_close(du_half_e2h, -div_s_e2h)
+        u_eprev = u_estep
+
+
+def test_poynting_planes(sim):
+    mask = (sim.js[0] != 0)
+    if mask.sum() > 1:
+        pytest.skip('test_poynting_planes can only test single point sources')
+
+    args = {'dxes': sim.dxes,
+            'epsilon': sim.epsilon}
+    dV = numpy.prod(numpy.meshgrid(*sim.dxes[0], indexing='ij'), axis=0)
+
+    mx = numpy.roll(mask, (-1, -1), axis=(0, 1))
+    my = numpy.roll(mask, -1, axis=2)
+    mz = numpy.roll(mask, (+1, -1), axis=(0, 3))
+    px = numpy.roll(mask, -1, axis=0)
+    py = mask.copy()
+    pz = numpy.roll(mask, +1, axis=0)
+
+    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)
+        u_estep = fdtd.energy_estep(h0=sim.hs[ii],   e1=sim.es[ii], h2=sim.hs[ii + 1], **args)
+
+        s_h2e = -fdtd.poynting(e=sim.es[ii], h=sim.hs[ii]) * sim.dt
+        s_h2e[0] *= sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :]
+        s_h2e[1] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][2][None, None, :]
+        s_h2e[2] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None]
+        planes = [s_h2e[px].sum(), -s_h2e[mx].sum(),
+                  s_h2e[py].sum(), -s_h2e[my].sum(),
+                  s_h2e[pz].sum(), -s_h2e[mz].sum()]
+        assert_close(sum(planes), (u_estep - u_hstep).sum())
+        if u_eprev is None:
+            u_eprev = u_estep
+            continue
+
+        s_e2h = -fdtd.poynting(e=sim.es[ii - 1], h=sim.hs[ii]) * sim.dt
+        s_e2h[0] *= sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :]
+        s_e2h[1] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][2][None, None, :]
+        s_e2h[2] *= sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None]
+        planes = [s_e2h[px].sum(), -s_e2h[mx].sum(),
+                  s_e2h[py].sum(), -s_e2h[my].sum(),
+                  s_e2h[pz].sum(), -s_e2h[mz].sum()]
+        assert_close(sum(planes), (u_hstep - u_eprev).sum())
+
+        # previous half-step
+        u_eprev = u_estep
+
+
+#####################################
+#      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=[0.3])
+def dt(request):
+    yield 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=['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(scope='module', params=[1.0, 1.5])
+def dx(request):
+    yield request.param
+
+
+@pytest.fixture(scope='module', params=['uniform'])
+def dxes(request, shape, dx):
+    if request.param == 'uniform':
+        dxes = [[numpy.full(s, dx) for s in shape[1:]] for _ in range(2)]
+    yield dxes
+
+
+@pytest.fixture(scope='module',
+                params=[(0,),
+                        (0, 4, 8),
+                       ]
+                )
+def j_steps(request):
+    yield request.param
+
+
+@dataclasses.dataclass()
+class SimResult:
+    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(scope='module')
+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 = SimResult(
+        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/types.py

@@ -0,0 +1,22 @@
+"""
+Types shared across multiple submodules
+"""
+import numpy
+from typing import List, Callable
+
+
+# Field types
+field_t = numpy.ndarray        # vector field with shape (3, X, Y, Z) (e.g. [E_x, E_y, E_z])
+vfield_t = numpy.ndarray       # linearized vector field (vector of length 3*X*Y*Z)
+
+'''
+ '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.
+'''
+dx_lists_t = List[List[numpy.ndarray]]
+
+
+field_updater = Callable[[field_t], field_t]

meanas/vectorization.py

@@ -4,15 +4,13 @@ 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'
+from .types import field_t, vfield_t
 
-# Types
-field_t = List[numpy.ndarray]  # vector field (eg. [E_x, E_y, E_z]
-vfield_t = numpy.ndarray       # linearized vector field
+
+__author__ = 'Jan Petykiewicz'
 
 
 def vec(f: field_t) -> vfield_t:
@@ -27,7 +25,7 @@ def vec(f: field_t) -> vfield_t:
     """
     if numpy.any(numpy.equal(f, None)):
         return None
-    return numpy.hstack(tuple((fi.ravel(order='C') for fi in f)))
+    return numpy.ravel(f, order='C')
 
 
 def unvec(v: vfield_t, shape: numpy.ndarray) -> field_t:
@@ -45,5 +43,5 @@ def unvec(v: vfield_t, shape: numpy.ndarray) -> field_t:
     """
     if numpy.any(numpy.equal(v, None)):
         return None
-    return [vi.reshape(shape, order='C') for vi in numpy.split(v, 3)]
+    return v.reshape((3, *shape), order='C')
 

setup.py

@@ -1,18 +1,36 @@
-#!/usr/bin/env python
+#!/usr/bin/env python3
 
 from setuptools import setup, find_packages
+import meanas
 
-setup(name='fdfd_tools',
-      version='0.4',
-      description='FDFD Electromagnetic simulation tools',
+with open('README.md', 'r') as f:
+    long_description = f.read()
+
+setup(name='meanas',
+      version=meanas.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/fdfd_tools',
       packages=find_packages(),
       install_requires=[
             'numpy',
             'scipy',
       ],
       extras_require={
+            'test': [
+                'pytest',
+                'dataclasses',
+                ],
       },
+      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',
+      ],
       )