initial development

.gitignore

@@ -58,3 +58,5 @@ docs/_build/
 # PyBuilder
 target/
 
+
+.idea/

README.md

@@ -1,3 +1,33 @@
-# fdfd-tools
+# fdfd_tools
 
-Python tools for creating finite-difference frequency-domain (FDFD) electromagnetic simulations.
+**fdfd_tools** is a python package containing utilities for
+creating and analyzing 2D and 3D finite-difference frequency-domain (FDFD)
+electromagnetic simulations.
+
+
+**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, ...)
+
+This package does *not* provide a matrix solver. The waveguide mode solver
+uses scipy's eigenvalue solver; I recommend a GPU-based iterative solver (eg.
+those included in [MAGMA](http://icl.cs.utk.edu/magma/index.html)). You 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)
+* numpy
+* scipy
+
+
+Install with pip, via git:
+```bash
+pip install git+https://mpxd.net/gogs/jan/fdfd_tools.git@release
+```

examples/test.py

@@ -0,0 +1,234 @@
+import numpy
+from numpy.ctypeslib import ndpointer
+import ctypes
+
+# h5py used by (uncalled) h5_write(); not used in currently-called code
+
+from fdfd_tools import vec, unvec, waveguide_mode
+import fdfd_tools, fdfd_tools.functional, fdfd_tools.grid
+import gridlock
+
+from matplotlib import pyplot
+
+__author__ = 'Jan Petykiewicz'
+
+
+def complex_to_alternating(x: numpy.ndarray) -> numpy.ndarray:
+    stacked = numpy.vstack((numpy.real(x), numpy.imag(x)))
+    return stacked.T.astype(numpy.float64).flatten()
+
+
+def solve_A(A, b: numpy.ndarray) -> numpy.ndarray:
+    A_vals = complex_to_alternating(A.data)
+    b_vals = complex_to_alternating(b)
+    x_vals = numpy.zeros_like(b_vals)
+
+    args = ['dummy',
+            '--solver', 'QMR',
+            '--maxiter', '40000',
+            '--atol', '1e-6',
+            '--verbose', '100']
+    argc = ctypes.c_int(len(args))
+    argv_arr_t = ctypes.c_char_p * len(args)
+    argv_arr = argv_arr_t()
+    argv_arr[:] = [s.encode('ascii') for s in args]
+
+    A_dim = ctypes.c_int(A.shape[0])
+    A_nnz = ctypes.c_int(A.nnz)
+    npdouble = ndpointer(ctypes.c_double)
+    npint = ndpointer(ctypes.c_int)
+
+    lib = ctypes.cdll.LoadLibrary('/home/jan/magma_solve/zsolve_shared.so')
+    c_solver = lib.zsolve
+    c_solver.argtypes = [ctypes.c_int, argv_arr_t,
+                         ctypes.c_int, ctypes.c_int,
+                         npdouble, npint, npint, npdouble, npdouble]
+
+    c_solver(argc, argv_arr, A_dim, A_nnz, A_vals,
+             A.indptr.astype(numpy.intc),
+             A.indices.astype(numpy.intc),
+             b_vals, x_vals)
+
+    x = (x_vals[::2] + 1j * x_vals[1::2]).flatten()
+    return x
+
+
+def write_h5(filename, A, b):
+    import h5py
+    # dtype=np.dtype([('real', 'float64'), ('imag', 'float64')])
+    h5py.get_config().complex_names = ('real', 'imag')
+    with h5py.File(filename, 'w') as mat_file:
+        mat_file.create_group('/A')
+        mat_file['/A/ir'] = A.indices.astype(numpy.intc)
+        mat_file['/A/jc'] = A.indptr.astype(numpy.intc)
+        mat_file['/A/data'] = A.data
+        mat_file['/b'] = b
+        mat_file['/x'] = numpy.zeros_like(b)
+
+
+def test0():
+    dx = 50                # discretization (nm/cell)
+    pml_thickness = 10     # (number of cells)
+
+    wl = 1550               # Excitation wavelength
+    omega = 2 * numpy.pi / wl
+
+    # Device design parameters
+    radii = (1, 0.6)
+    th = 220
+    center = [0, 0, 0]
+
+    # refractive indices
+    n_ring = numpy.sqrt(12.6)  # ~Si
+    n_air = 4.0   # air
+
+    # Half-dimensions of the simulation grid
+    xyz_max = numpy.array([1.2, 1.2, 0.3]) * 1000 + pml_thickness * dx
+
+    # Coordinates of the edges of the cells.
+    half_edge_coords = [numpy.arange(dx/2, m + dx, step=dx) for m in xyz_max]
+    edge_coords = [numpy.hstack((-h[::-1], h)) for h in half_edge_coords]
+
+    # #### Create the grid, mask, and draw the device ####
+    grid = gridlock.Grid(edge_coords, initial=n_air**2, num_grids=3)
+    grid.draw_cylinder(surface_normal=gridlock.Direction.z,
+                       center=center,
+                       radius=max(radii),
+                       thickness=th,
+                       eps=n_ring**2,
+                       num_points=24)
+    grid.draw_cylinder(surface_normal=gridlock.Direction.z,
+                       center=center,
+                       radius=min(radii),
+                       thickness=th*1.1,
+                       eps=n_air ** 2,
+                       num_points=24)
+
+    dx0_a = grid.dxyz
+    dx0_b = [grid.shifted_dxyz(which_shifts=a)[a] for a in range(3)]
+    dxes = [dx0_a, dx0_b]
+    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)
+
+    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
+
+    A = fdfd_tools.functional.e_full(omega, dxes, vec(grid.grids)).tocsr()
+    b = -1j * omega * vec(J)
+
+    x = solve_A(A, b)
+    E = unvec(x, grid.shape)
+
+    print('Norm of the residual is {}'.format(numpy.linalg.norm(A.dot(x) - b)/numpy.linalg.norm(b)))
+
+    pyplot.figure()
+    pyplot.pcolor(numpy.real(E[1][:, :, grid.shape[2]//2]), cmap='seismic')
+    pyplot.axis('equal')
+    pyplot.show()
+
+
+def test1():
+    dx = 40                 # discretization (nm/cell)
+    pml_thickness = 10      # (number of cells)
+
+    wl = 1550               # Excitation wavelength
+    omega = 2 * numpy.pi / wl
+
+    # Device design parameters
+    w = 600
+    th = 220
+    center = [0, 0, 0]
+
+    # refractive indices
+    n_wg = numpy.sqrt(12.6)  # ~Si
+    n_air = 1.0              # air
+
+    # Half-dimensions of the simulation grid
+    xyz_max = numpy.array([0.8, 0.9, 0.6]) * 1000 + (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]
+
+    # #### 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)
+
+    dx0_a = grid.dxyz
+    dx0_b = [grid.shifted_dxyz(which_shifts=a)[a] for a in range(3)]
+    dxes = [dx0_a, dx0_b]
+    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)
+
+    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))
+    wg_args = {
+        'omega': omega,
+        'slices': [slice(i, f+1) for i, f in zip(*ind_dims)],
+        'dxes': dxes,
+        'axis': 0,
+        'polarity': +1,
+    }
+
+    wg_results = waveguide_mode.solve_waveguide_mode(mode_number=0, **wg_args, epsilon=grid.grids)
+    J = waveguide_mode.compute_source(**wg_args, **wg_results)
+    H_overlap = waveguide_mode.compute_overlap_e(**wg_args, **wg_results)
+
+    A = fdfd_tools.operators.e_full(omega, dxes, vec(grid.grids)).tocsr()
+    b = -1j * omega * vec(J)
+    x = solve_A(A, b)
+    E = unvec(x, grid.shape)
+
+    print('Norm of the residual is ', numpy.linalg.norm(A @ x - b))
+
+    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], :, :]))
+    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]]))
+    pyplot.subplot(2, 2, 4)
+
+    def poyntings(E):
+        e = vec(E)
+        h = fdfd_tools.operators.e2h(omega, dxes) @ e
+        cross1 = fdfd_tools.operators.poynting_e_cross(e, dxes) @ h.conj()
+        cross2 = fdfd_tools.operators.poynting_h_cross(h.conj(), dxes) @ e
+        s1 = unvec(0.5 * numpy.real(cross1), grid.shape)
+        s2 = unvec(0.5 * numpy.real(-cross2), grid.shape)
+        return s1, s2
+
+    s1x, s2x = poyntings(E)
+    pyplot.plot(s1x[0].sum(axis=2).sum(axis=1))
+    pyplot.hold(True)
+    pyplot.plot(s2x[0].sum(axis=2).sum(axis=1))
+    pyplot.show()
+
+    q = []
+    for i in range(-5, 30):
+        H_rolled = [numpy.roll(h, i, axis=0) for h in H_overlap]
+        q += [numpy.abs(vec(E) @ vec(H_rolled))]
+    pyplot.figure()
+    pyplot.plot(q)
+    pyplot.title('Overlap with mode')
+    pyplot.show()
+    print('Average overlap with mode:', sum(q)/len(q))
+
+if __name__ == '__main__':
+    # test0()
+    test1()

fdfd_tools/__init__.py

@@ -0,0 +1,25 @@
+"""
+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,
+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/functional.py

@@ -0,0 +1,149 @@
+"""
+Functional versions of many FDFD operators. These can be useful for performing
+ FDFD calculations without needing to construct large matrices in memory.
+
+The functions generated here expect inputs in the form E = [E_x, E_y, E_z], where each
+ component E_* is an ndarray of equal shape.
+"""
+from typing import List, Callable
+import numpy
+
+from . import dx_lists_t, field_t
+
+__author__ = 'Jan Petykiewicz'
+
+
+functional_matrix = Callable[[List[numpy.ndarray]], List[numpy.ndarray]]
+
+
+def curl_h(dxes: dx_lists_t) -> functional_matrix:
+    """
+    Curl operator for use with the H field.
+
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Function for taking the discretized curl of the H-field, F(H) -> curlH
+    """
+    dxyz_b = numpy.meshgrid(*dxes[1], indexing='ij')
+
+    def dH(f, ax):
+        return (f - numpy.roll(f, 1, axis=ax)) / dxyz_b[ax]
+
+    def ch_fun(H: List[numpy.ndarray]) -> List[numpy.ndarray]:
+        E = [dH(H[2], 1) - dH(H[1], 2),
+             dH(H[0], 2) - dH(H[2], 0),
+             dH(H[1], 0) - dH(H[0], 1)]
+        return E
+
+    return ch_fun
+
+
+def curl_e(dxes: dx_lists_t) -> functional_matrix:
+    """
+    Curl operator for use with the E field.
+
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Function for taking the discretized curl of the E-field, F(E) -> curlE
+    """
+    dxyz_a = numpy.meshgrid(*dxes[0], indexing='ij')
+
+    def dE(f, ax):
+        return (numpy.roll(f, -1, axis=ax) - f) / dxyz_a[ax]
+
+    def ce_fun(E: List[numpy.ndarray]) -> List[numpy.ndarray]:
+        H = [dE(E[2], 1) - dE(E[1], 2),
+             dE(E[0], 2) - dE(E[2], 0),
+             dE(E[1], 0) - dE(E[0], 1)]
+        return H
+
+    return ce_fun
+
+
+def e_full(omega: complex,
+           dxes: dx_lists_t,
+           epsilon: field_t,
+           mu: field_t = None
+           ) -> functional_matrix:
+    """
+    Wave operator del x (1/mu * del x) - omega**2 * epsilon, for use with E-field,
+     with wave equation
+    (del x (1/mu * del x) - omega**2 * epsilon) E = -i * omega * J
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param epsilon: Dielectric constant
+    :param mu: Magnetic permeability (default 1 everywhere)
+    :return: Function implementing the wave operator A(E) -> E
+    """
+    ch = curl_h(dxes)
+    ce = curl_e(dxes)
+
+    def op_1(E):
+        curls = ch(ce(E))
+        return [c - omega ** 2 * e * x for c, e, x in zip(curls, epsilon, E)]
+
+    def op_mu(E):
+        curls = ch([m * y for m, y in zip(mu, ce(E))])
+        return [c - omega ** 2 * e * x for c, e, x in zip(curls, epsilon, E)]
+
+    if numpy.any(numpy.equal(mu, None)):
+        return op_1
+    else:
+        return op_mu
+
+
+def eh_full(omega: complex,
+            dxes: dx_lists_t,
+            epsilon: field_t,
+            mu: field_t = None
+            ) -> functional_matrix:
+    """
+    Wave operator for full (both E and H) field representation.
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param epsilon: Dielectric constant
+    :param mu: Magnetic permeability (default 1 everywhere)
+    :return: Function implementing the wave operator A(E, H) -> (E, H)
+    """
+    ch = curl_h(dxes)
+    ce = curl_e(dxes)
+
+    def op_1(E, H):
+        return ([c - 1j * omega * e * x for c, e, x in zip(ch(H), epsilon, E)],
+                [c + 1j * omega * y for c, y in zip(ce(E), H)])
+
+    def op_mu(E, H):
+        return ([c - 1j * omega * e * x for c, e, x in zip(ch(H), epsilon, E)],
+                [c + 1j * omega * m * y for c, m, y in zip(ce(E), mu, H)])
+
+    if numpy.any(numpy.equal(mu, None)):
+        return op_1
+    else:
+        return op_mu
+
+
+def e2h(omega: complex,
+        dxes: dx_lists_t,
+        mu: field_t = None,
+        ) -> functional_matrix:
+    """
+   Utility operator for converting the E field into the H field.
+   For use with e_full -- assumes that there is no magnetic current M.
+
+   :param omega: Angular frequency of the simulation
+   :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+   :param mu: Magnetic permeability (default 1 everywhere)
+   :return: Function for converting E to H
+   """
+    A2 = curl_e(dxes)
+
+    def e2h_1_1(E):
+        return [y / (-1j * omega) for y in A2(E)]
+
+    def e2h_mu(E):
+        return [y / (-1j * omega * m) for y, m in zip(A2(E), mu)]
+
+    if numpy.any(numpy.equal(mu, None)):
+        return e2h_1_1
+    else:
+        return e2h_mu

fdfd_tools/grid.py

@@ -0,0 +1,167 @@
+"""
+Functions for creating stretched coordinate PMLs.
+"""
+
+from typing import List, Tuple, Callable
+import numpy
+
+__author__ = 'Jan Petykiewicz'
+
+
+dx_lists_t = List[List[numpy.ndarray]]
+s_function_type = Callable[[float], float]
+
+
+def prepare_s_function(ln_R: float = -16,
+                       m: float = 4
+                       ) -> s_function_type:
+    """
+    Create an s_function to pass to the SCPML functions. This is used when you would like to
+    customize the PML parameters.
+
+    :param ln_R: Natural logarithm of the desired reflectance
+    :param m: Polynomial order for the PML (imaginary part increases as distance ** m)
+    :return: An s_function, which takes an ndarray (distances) and returns an ndarray (complex part of
+                the cell width; needs to be divided by sqrt(epilon_effective) * real(omega)) before
+                use.
+    """
+    def s_factor(distance: numpy.ndarray) -> numpy.ndarray:
+        s_max = (m + 1) * ln_R / 2  # / 2 because we have assume boundaries
+        return s_max * (distance ** m)
+    return s_factor
+
+
+def uniform_grid_scpml(shape: numpy.ndarray or List[int],
+                       thicknesses: numpy.ndarray or List[int],
+                       omega: float,
+                       epsilon_effective: float = 1.0,
+                       s_function: s_function_type = None,
+                       ) -> dx_lists_t:
+    """
+    Create dx arrays for a uniform grid with a cell width of 1 and a pml.
+
+    If you want something more fine-grained, check out stretch_with_scpml(...).
+
+    :param shape: Shape of the grid, including the PMLs (which are 2*thicknesses thick)
+    :param thicknesses: [th_x, th_y, th_z] Thickness of the PML in each direction. Both polarities are added.
+        Each th_ of pml is applied twice, once on each edge of the grid along the given axis.
+        th_* may be zero, in which case no pml is added.
+    :param omega: Angular frequency for the simulation
+    :param epsilon_effective: Effective epsilon of the PML. Match this to the material at the edge of your grid.
+        Default 1.
+    :param s_function: s_function created by prepare_s_function(...), allowing customization of pml parameters.
+        Default uses prepare_s_function() with no parameters.
+    :return: Complex cell widths (dx_lists)
+    """
+    if s_function is None:
+        s_function = prepare_s_function()
+
+    # Normalized distance to nearest boundary
+    def l(u, n, t):
+        return ((t - u).clip(0) + (u - (n - t)).clip(0)) / t
+
+    dx_a = [numpy.array(numpy.inf)] * 3
+    dx_b = [numpy.array(numpy.inf)] * 3
+
+    # divide by this to adjust for epsilon_effective and omega
+    s_correction = numpy.sqrt(epsilon_effective) * numpy.real(omega)
+
+    for k, th in enumerate(thicknesses):
+        s = shape[k]
+        if th > 0:
+            sr = numpy.arange(s)
+            dx_a[k] = 1 + 1j * s_function(l(sr, s, th)) / s_correction
+            dx_b[k] = 1 + 1j * s_function(l(sr+0.5, s, th)) / s_correction
+        else:
+            dx_a[k] = numpy.ones((s,))
+            dx_b[k] = numpy.ones((s,))
+    return [dx_a, dx_b]
+
+
+def stretch_with_scpml(dxes: dx_lists_t,
+                       axis: int,
+                       polarity: int,
+                       omega: float,
+                       epsilon_effective: float = 1.0,
+                       thickness: int = 10,
+                       s_function: s_function_type = None,
+                       ) -> dx_lists_t:
+    """
+        Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis.
+
+        :param dxes: dx_tuple with coordinates to stretch
+        :param axis: axis to stretch (0=x, 1=y, 2=z)
+        :param polarity: direction to stretch (-1 for -ve, +1 for +ve)
+        :param omega: Angular frequency for the simulation
+        :param epsilon_effective: Effective epsilon of the PML. Match this to the material at the edge of your grid.
+            Default 1.
+        :param thickness: number of cells to use for pml (default 10)
+        :param s_function: s_function created by prepare_s_function(...), allowing customization of pml parameters.
+            Default uses prepare_s_function() with no parameters.
+        :return: Complex cell widths
+    """
+    if s_function is None:
+        s_function = prepare_s_function()
+
+    dx_ai = dxes[0][axis].astype(complex)
+    dx_bi = dxes[1][axis].astype(complex)
+
+    pos = numpy.hstack((0, dx_ai.cumsum()))
+    pos_a = (pos[:-1] + pos[1:]) / 2
+    pos_b = pos[:-1]
+
+    # divide by this to adjust for epsilon_effective and omega
+    s_correction = numpy.sqrt(epsilon_effective) * numpy.real(omega)
+
+    if polarity > 0:
+        # front pml
+        bound = pos[thickness]
+        d = bound - pos[0]
+
+        def l_d(x):
+            return (bound - x) / (bound - pos[0])
+
+        slc = slice(thickness)
+
+    else:
+        # back pml
+        bound = pos[-thickness - 1]
+        d = pos[-1] - bound
+
+        def l_d(x):
+            return (x - bound) / (pos[-1] - bound)
+
+        if thickness == 0:
+            slc = slice(None)
+        else:
+            slc = slice(-thickness, None)
+
+    dx_ai[slc] *= 1 + 1j * s_function(l_d(pos_a[slc])) / d / s_correction
+    dx_bi[slc] *= 1 + 1j * s_function(l_d(pos_b[slc])) / d / s_correction
+
+    dxes[0][axis] = dx_ai
+    dxes[1][axis] = dx_bi
+
+    return dxes
+
+
+def generate_dx(pos: List[numpy.ndarray]) -> dx_lists_t:
+    """
+    Given a list of 3 ndarrays cell centers, creates the cell width parameters.
+
+    :param pos: List of 3 ndarrays of cell centers
+    :return: (dx_a, dx_b) cell widths (no pml)
+    """
+    if len(pos) != 3:
+        raise Exception('Must have len(pos) == 3')
+
+    dx_a = [numpy.array(numpy.inf)] * 3
+    dx_b = [numpy.array(numpy.inf)] * 3
+
+    for i, p_orig in enumerate(pos):
+        p = numpy.array(p_orig, dtype=float)
+        if p.size != 1:
+            p_shifted = numpy.hstack((p[1:], p[-1] + (p[1] - p[0])))
+            dx_a[i] = numpy.diff(p)
+            dx_b[i] = numpy.diff((p + p_shifted) / 2)
+    return dx_a, dx_b

fdfd_tools/operators.py

@@ -0,0 +1,397 @@
+"""
+Sparse matrix operators for use with electromagnetic wave equations.
+
+These functions return sparse-matrix (scipy.sparse.spmatrix) representations of
+ a variety of operators, intended for use with E and H fields vectorized using the
+ fdfd_tools.vec() and .unvec() functions (column-major/Fortran ordering).
+
+E- and H-field values are defined on a Yee cell; epsilon values should be calculated for
+ cells centered at each E component (mu at each H component).
+
+Many of these functions require a 'dxes' parameter, of type fdfd_tools.dx_lists_type,
+ which contains grid cell width information in the following format:
+ [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...], [dz_e_0, ...]],
+  [[dx_h_0, dx_h_1, ...], [dy_h_0, ...], [dz_h_0, ...]]]
+ where dx_e_0 is the x-width of the x=0 cells, as used when calculating dE/dx,
+ and dy_h_0 is  the y-width of the y=0 cells, as used when calculating dH/dy, etc.
+
+
+The following operators are included:
+- E-only wave operator
+- H-only wave operator
+- EH wave operator
+- Curl for use with E, H fields
+- E to H conversion
+- M to J conversion
+- Poynting cross products
+
+Also available:
+- Circular shifts
+- Discrete derivatives
+- Averaging operators
+- Cross product matrices
+"""
+
+from typing import List, Tuple
+import numpy
+import scipy.sparse as sparse
+
+from . import vec, dx_lists_t, vfield_t
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def e_full(omega: complex,
+           dxes: dx_lists_t,
+           epsilon: vfield_t,
+           mu: vfield_t = None
+           ) -> sparse.spmatrix:
+    """
+    Wave operator del x (1/mu * del x) - omega**2 * epsilon, for use with E-field,
+     with wave equation
+    (del x (1/mu * del x) - omega**2 * epsilon) E = -i * omega * J
+
+    To make this matrix symmetric, use the preconditions from e_full_preconditioners().
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param epsilon: Vectorized dielectric constant
+    :param mu: Vectorized magnetic permeability (default 1 everywhere)..
+    :return: Sparse matrix containing the wave operator
+    """
+    ce = curl_e(dxes)
+    ch = curl_h(dxes)
+
+    e = sparse.diags(epsilon)
+    if numpy.any(numpy.equal(mu, None)):
+        m_div = sparse.eye(epsilon.size)
+    else:
+        m_div = sparse.diags(1 / mu)
+
+    op = ch @ m_div @ ce - omega**2 * e
+    return op
+
+
+def e_full_preconditioners(dxes: dx_lists_t
+                           ) -> Tuple[sparse.spmatrix, sparse.spmatrix]:
+    """
+    Left and right preconditioners (Pl, Pr) for symmetrizing the e_full wave operator.
+
+    The preconditioned matrix A_symm = (Pl @ A @ Pr) is complex-symmetric
+     (non-Hermitian unless there is no loss or PMLs).
+
+    The preconditioner matrices are diagonal and complex, with Pr = 1 / Pl
+
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Preconditioner matrices (Pl, Pr)
+    """
+    p_squared = [dxes[0][0][:, None, None] * dxes[1][1][None, :, None] * dxes[1][2][None, None, :],
+                 dxes[1][0][:, None, None] * dxes[0][1][None, :, None] * dxes[1][2][None, None, :],
+                 dxes[1][0][:, None, None] * dxes[1][1][None, :, None] * dxes[0][2][None, None, :]]
+
+    p_vector = numpy.sqrt(vec(p_squared))
+    P_left = sparse.diags(p_vector)
+    P_right = sparse.diags(1 / p_vector)
+    return P_left, P_right
+
+
+def h_full(omega: complex,
+           dxes: dx_lists_t,
+           epsilon: vfield_t,
+           mu: vfield_t = None
+           ) -> sparse.spmatrix:
+    """
+    Wave operator del x (1/epsilon * del x) - omega**2 * mu, for use with H-field,
+     with wave equation
+    (del x (1/epsilon * del x) - omega**2 * mu) H = i * omega * M
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param epsilon: Vectorized dielectric constant
+    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+    :return: Sparse matrix containing the wave operator
+    """
+    ec = curl_e(dxes)
+    hc = curl_h(dxes)
+
+    e_div = sparse.diags(1 / epsilon)
+    if numpy.any(numpy.equal(mu, None)):
+        m = sparse.eye(epsilon.size)
+    else:
+        m = sparse.diags(mu)
+
+    A = ec @ e_div @ hc - omega**2 * m
+    return A
+
+
+def eh_full(omega, dxes, epsilon, mu=None):
+    """
+    Wave operator for [E, H] field representation. This operator implements Maxwell's
+     equations without cancelling out either E or H. The operator is
+     [[-i * omega * epsilon,  del x],
+      [del x, i * omega * mu]]
+
+    for use with a field vector of the form hstack(vec(E), vec(H)).
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param epsilon: Vectorized dielectric constant
+    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+    :return: Sparse matrix containing the wave operator
+    """
+    A2 = curl_e(dxes)
+    A1 = curl_h(dxes)
+
+    iwe = 1j * omega * sparse.diags(epsilon)
+    iwm = 1j * omega
+    if not numpy.any(numpy.equal(mu, None)):
+        iwm *= sparse.diags(mu)
+
+    A = sparse.bmat([[-iwe, A1],
+                     [A2, +iwm]])
+    return A
+
+
+def curl_h(dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Curl operator for use with the H field.
+
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Sparse matrix for taking the discretized curl of the H-field
+    """
+    return cross(deriv_back(dxes[1]))
+
+
+def curl_e(dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Curl operator for use with the E field.
+
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Sparse matrix for taking the discretized curl of the E-field
+    """
+    return cross(deriv_forward(dxes[0]))
+
+
+def e2h(omega: complex,
+        dxes: dx_lists_t,
+        mu: vfield_t = None,
+        ) -> sparse.spmatrix:
+    """
+    Utility operator for converting the E field into the H field.
+    For use with e_full -- assumes that there is no magnetic current M.
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+    :return: Sparse matrix for converting E to H
+    """
+    op = curl_e(dxes) / (-1j * omega)
+
+    if not numpy.any(numpy.equal(mu, None)):
+        op = sparse.diags(1 / mu) @ op
+
+    return op
+
+
+def m2j(omega: complex,
+        dxes: dx_lists_t,
+        mu: vfield_t = None):
+    """
+    Utility operator for converting M field into J.
+    Converts a magnetic current M into an electric current J.
+    For use with eg. e_full.
+
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param mu: Vectorized magnetic permeability (default 1 everywhere)
+    :return: Sparse matrix for converting E to H
+    """
+    op = curl_h(dxes) / (1j * omega)
+
+    if not numpy.any(numpy.equal(mu, None)):
+        op = op @ sparse.diags(1 / mu)
+
+    return op
+
+
+def rotation(axis: int, shape: List[int]) -> sparse.spmatrix:
+    """
+    Utility operator for performing a circular shift along a specified axis by 1 element.
+
+    :param axis: Axis to shift along. x=0, y=1, z=2
+    :param shape: Shape of the grid being shifted
+    :return: Sparse matrix for performing the circular shift
+    """
+    if len(shape) not in (2, 3):
+        raise Exception('Invalid shape: {}'.format(shape))
+    if axis not in range(len(shape)):
+        raise Exception('Invalid direction: {}, shape is {}'.format(axis, shape))
+
+    n = numpy.prod(shape)
+
+    shifts = [1 if k == axis else 0 for k in range(3)]
+    shifted_diags = [(numpy.arange(n) + s) % n for n, s in zip(shape, shifts)]
+    ijk = numpy.meshgrid(*shifted_diags, indexing='ij')
+
+    i_ind = numpy.arange(n)
+    j_ind = ijk[0] + ijk[1] * shape[0]
+    if len(shape) == 3:
+        j_ind += ijk[2] * shape[0] * shape[1]
+
+    vij = (numpy.ones(n), (i_ind, j_ind.flatten(order='F')))
+
+    D = sparse.csr_matrix(vij, shape=(n, n))
+    return D
+
+
+def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]:
+    """
+    Utility operators for taking discretized derivatives (forward variant).
+
+    :param dx_e: Lists of cell sizes for all axes [[dx_0, dx_1, ...], ...].
+    :return: List of operators for taking forward derivatives along each axis.
+    """
+    shape = [s.size for s in dx_e]
+    n = numpy.prod(shape)
+
+    dx_e_expanded = numpy.meshgrid(*dx_e, indexing='ij')
+
+    def deriv(axis):
+        return rotation(axis, shape) - sparse.eye(n)
+
+    Ds = [sparse.diags(+1 / dx.flatten(order='F')) @ deriv(a)
+          for a, dx in enumerate(dx_e_expanded)]
+
+    return Ds
+
+
+def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]:
+    """
+    Utility operators for taking discretized derivatives (backward variant).
+
+    :param dx_h: Lists of cell sizes for all axes [[dx_0, dx_1, ...], ...].
+    :return: List of operators for taking forward derivatives along each axis.
+    """
+    shape = [s.size for s in dx_h]
+    n = numpy.prod(shape)
+
+    dx_h_expanded = numpy.meshgrid(*dx_h, indexing='ij')
+
+    def deriv(axis):
+        return rotation(axis, shape) - sparse.eye(n)
+
+    Ds = [sparse.diags(-1 / dx.flatten(order='F')) @ deriv(a).T
+          for a, dx in enumerate(dx_h_expanded)]
+
+    return Ds
+
+
+def cross(B: List[sparse.spmatrix]) -> sparse.spmatrix:
+    """
+    Cross product operator
+
+    :param B: List [Bx, By, Bz] of sparse matrices corresponding to the x, y, z
+            portions of the operator on the left side of the cross product.
+    :return: Sparse matrix corresponding to (B x), where x is the cross product
+    """
+    n = B[0].shape[0]
+    zero = sparse.csr_matrix((n, n))
+    return sparse.bmat([[zero, -B[2], B[1]],
+                        [B[2], zero, -B[0]],
+                        [-B[1], B[0], zero]])
+
+
+def vec_cross(b: vfield_t) -> sparse.spmatrix:
+    """
+    Vector cross product operator
+
+    :param b: Vector on the left side of the cross product
+    :return: Sparse matrix corresponding to (b x), where x is the cross product
+    """
+    B = [sparse.diags(c) for c in numpy.split(b, 3)]
+    return cross(B)
+
+
+def avgf(axis: int, shape: List[int]) -> sparse.spmatrix:
+    """
+    Forward average operator (x4 = (x4 + x5) / 2)
+
+    :param axis: Axis to average along (x=0, y=1, z=2)
+    :param shape: Shape of the grid to average
+    :return: Sparse matrix for forward average operation
+    """
+    if len(shape) not in (2, 3):
+        raise Exception('Invalid shape: {}'.format(shape))
+
+    n = numpy.prod(shape)
+    return 0.5 * (sparse.eye(n) + rotation(axis, shape))
+
+
+def avgb(axis: int, shape: List[int]) -> sparse.spmatrix:
+    """
+    Backward average operator (x4 = (x4 + x3) / 2)
+
+    :param axis: Axis to average along (x=0, y=1, z=2)
+    :param shape: Shape of the grid to average
+    :return: Sparse matrix for backward average operation
+    """
+    return avgf(axis, shape).T
+
+
+def poynting_e_cross(e: vfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Operator for computing the Poynting vector, contining the (E x) portion of the Poynting vector.
+
+    :param e: Vectorized E-field for the ExH cross product
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Sparse matrix containing (E x) portion of Poynting cross product
+    """
+    shape = [len(dx) for dx in dxes[0]]
+
+    fx, fy, fz = [avgf(i, shape) for i in range(3)]
+    bx, by, bz = [avgb(i, shape) for i in range(3)]
+
+    dxag = [dx.flatten(order='F') for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
+    dbgx, dbgy, dbgz = [sparse.diags(dx.flatten(order='F'))
+                        for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
+
+    Ex, Ey, Ez = [sparse.diags(ei * da) for ei, da in zip(numpy.split(e, 3), dxag)]
+
+    n = numpy.prod(shape)
+    zero = sparse.csr_matrix((n, n))
+
+    P = sparse.bmat(
+        [[ zero,                -fx @ Ez @ bz @ dbgy,  fx @ Ey @ by @ dbgz],
+         [ fy @ Ez @ bz @ dbgx,  zero,                -fy @ Ex @ bx @ dbgz],
+         [-fz @ Ey @ by @ dbgx,  fz @ Ex @ bx @ dbgy,  zero]])
+    return P
+
+
+def poynting_h_cross(h: vfield_t, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Operator for computing the Poynting vector, containing the (H x) portion of the Poynting vector.
+
+    :param h: Vectorized H-field for the HxE cross product
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :return: Sparse matrix containing (H x) portion of Poynting cross product
+    """
+    shape = [len(dx) for dx in dxes[0]]
+
+    fx, fy, fz = [avgf(i, shape) for i in range(3)]
+    bx, by, bz = [avgb(i, shape) for i in range(3)]
+
+    dxbg = [dx.flatten(order='F') for dx in numpy.meshgrid(*dxes[1], indexing='ij')]
+    dagx, dagy, dagz = [sparse.diags(dx.flatten(order='F'))
+                        for dx in numpy.meshgrid(*dxes[0], indexing='ij')]
+
+    Hx, Hy, Hz = [sparse.diags(hi * db) for hi, db in zip(numpy.split(h, 3), dxbg)]
+
+    n = numpy.prod(shape)
+    zero = sparse.csr_matrix((n, n))
+
+    P = sparse.bmat(
+        [[ zero,                -by @ Hz @ fx @ dagy,  bz @ Hy @ fx @ dagz],
+         [ bx @ Hz @ fy @ dagx,  zero,                -bz @ Hx @ fy @ dagz],
+         [-bx @ Hy @ fz @ dagx,  by @ Hx @ fz @ dagy,  zero]])
+    return P

fdfd_tools/vectorization.py

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

fdfd_tools/waveguide.py

@@ -0,0 +1,309 @@
+"""
+Various operators and helper functions for solving for waveguide modes.
+
+Assuming a z-dependence of the from exp(-i * wavenumber * z), we can simplify Maxwell's
+ equations in the absence of sources to the form
+
+A @ [H_x, H_y] = wavenumber**2 * [H_x, H_y]
+
+with A =
+omega**2 * epsilon * mu +
+epsilon * [[-Dy], [Dx]] / epsilon * [-Dy, Dx] +
+[[Dx], [Dy]] / mu * [Dx, Dy] * mu
+
+which is the form used in this file.
+
+As the z-dependence is known, all the functions in this file assume a 2D grid
+ (ie. dxes = [[[dx_e_0, dx_e_1, ...], [dy_e_0, ...]], [[dx_h_0, ...], [dy_h_0, ...]]])
+ with propagation along the z axis.
+"""
+
+from typing import List, Tuple
+import numpy
+from numpy.linalg import norm
+import scipy.sparse as sparse
+
+from . import unvec, dx_lists_t, field_t, vfield_t
+from . import operators
+
+
+__author__ = 'Jan Petykiewicz'
+
+
+def operator(omega: complex,
+             dxes: dx_lists_t,
+             epsilon: vfield_t,
+             mu: vfield_t = None,
+             ) -> sparse.spmatrix:
+    """
+    Waveguide operator of the form
+
+    omega**2 * epsilon * mu +
+    epsilon * [[-Dy], [Dx]] / epsilon * [-Dy, Dx] +
+    [[Dx], [Dy]] / mu * [Dx, Dy] * mu
+
+    for use with a field vector of the form [H_x, H_y].
+
+    This operator can be used to form an eigenvalue problem of the form
+    A @ [H_x, H_y] = wavenumber**2 * [H_x, H_y]
+
+    which can then be solved for the eigenmodes of the system (an exp(-i * wavenumber * z)
+    z-dependence is assumed for the fields).
+
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Sparse matrix representation of the operator
+    """
+    if numpy.any(numpy.equal(mu, None)):
+        mu = numpy.ones_like(epsilon)
+
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    Dbx, Dby = operators.deriv_back(dxes[1])
+
+    eps_parts = numpy.split(epsilon, 3)
+    eps_yx = sparse.diags(numpy.hstack((eps_parts[1], eps_parts[0])))
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    mu_parts = numpy.split(mu, 3)
+    mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1])))
+    mu_z_inv = sparse.diags(1 / mu_parts[2])
+
+    op = omega ** 2 * eps_yx @ mu_xy + \
+        eps_yx @ sparse.vstack((-Dfy, Dfx)) @ eps_z_inv @ sparse.hstack((-Dby, Dbx)) + \
+        sparse.vstack((Dbx, Dby)) @ mu_z_inv @ sparse.hstack((Dfx, Dfy)) @ mu_xy
+
+    return op
+
+
+def normalized_fields(v: numpy.ndarray,
+                      wavenumber: complex,
+                      omega: complex,
+                      dxes: dx_lists_t,
+                      epsilon: vfield_t,
+                      mu: vfield_t = None
+                      ) -> Tuple[vfield_t, vfield_t]:
+    """
+    Given a vector v containing the vectorized H_x and H_y fields,
+     returns normalized, vectorized E and H fields for the system.
+
+    :param v: Vector containing H_x and H_y fields
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Normalized, vectorized (e, h) containing all vector components.
+    """
+    e = v2e(v, wavenumber, omega, dxes, epsilon, mu=mu)
+    h = v2h(v, wavenumber, dxes, mu=mu)
+
+    shape = [s.size for s in dxes[0]]
+    dxes_real = [[numpy.real(d) for d in numpy.meshgrid(*dxes[v], indexing='ij')] for v in (0, 1)]
+
+    E = unvec(e, shape)
+    H = unvec(h, shape)
+
+    S1 = E[0] * numpy.roll(numpy.conj(H[1]), 1, axis=0) * dxes_real[0][1] * dxes_real[1][0]
+    S2 = E[1] * numpy.roll(numpy.conj(H[0]), 1, axis=1) * dxes_real[0][0] * dxes_real[1][1]
+    S = 0.25 * ((S1 + numpy.roll(S1, 1, axis=0)) -
+                (S2 + numpy.roll(S2, 1, axis=1)))
+    P = 0.5 * numpy.real(S.sum())
+    assert P > 0, 'Found a mode propagating in the wrong direction! P={}'.format(P)
+
+    norm_amplitude = 1 / numpy.sqrt(P)
+    norm_angle = -numpy.angle(e[e.size//2])
+    norm_factor = norm_amplitude * numpy.exp(1j * norm_angle)
+
+    e *= norm_factor
+    h *= norm_factor
+
+    return e, h
+
+
+def v2h(v: numpy.ndarray,
+        wavenumber: complex,
+        dxes: dx_lists_t,
+        mu: vfield_t = None
+        ) -> vfield_t:
+    """
+    Given a vector v containing the vectorized H_x and H_y fields,
+     returns a vectorized H including all three H components.
+
+    :param v: Vector containing H_x and H_y fields
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Vectorized H field with all vector components
+    """
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    op = sparse.hstack((Dfx, Dfy))
+
+    if not numpy.any(numpy.equal(mu, None)):
+        mu_parts = numpy.split(mu, 3)
+        mu_xy = sparse.diags(numpy.hstack((mu_parts[0], mu_parts[1])))
+        mu_z_inv = sparse.diags(1 / mu_parts[2])
+
+        op = mu_z_inv @ op @ mu_xy
+
+    w = op @ v / (1j * wavenumber)
+    return numpy.hstack((v, w)).flatten()
+
+
+def v2e(v: numpy.ndarray,
+        wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfield_t,
+        mu: vfield_t = None
+        ) -> vfield_t:
+    """
+    Given a vector v containing the vectorized H_x and H_y fields,
+     returns a vectorized E containing all three E components
+
+    :param v: Vector containing H_x and H_y fields
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Vectorized E field with all vector components.
+    """
+    h2eop = h2e(wavenumber, omega, dxes, epsilon)
+    return h2eop @ v2h(v, wavenumber, dxes, mu)
+
+
+def e2h(wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        mu: vfield_t = None
+        ) -> sparse.spmatrix:
+    """
+    Returns an operator which, when applied to a vectorized E eigenfield, produces
+     the vectorized H eigenfield.
+
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Sparse matrix representation of the operator
+    """
+    op = curl_e(wavenumber, dxes) / (-1j * omega)
+    if not numpy.any(numpy.equal(mu, None)):
+        op = sparse.diags(1 / mu) @ op
+    return op
+
+
+def h2e(wavenumber: complex,
+        omega: complex,
+        dxes: dx_lists_t,
+        epsilon: vfield_t
+        ) -> sparse.spmatrix:
+    """
+    Returns an operator which, when applied to a vectorized H eigenfield, produces
+     the vectorized E eigenfield.
+
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :return: Sparse matrix representation of the operator
+    """
+    op = sparse.diags(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes)
+    return op
+
+
+def curl_e(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Discretized curl operator for use with the waveguide E field.
+
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :return: Sparse matrix representation of the operator
+    """
+    n = 1
+    for d in dxes[0]:
+        n *= len(d)
+
+    Bz = -1j * wavenumber * sparse.eye(n)
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    return operators.cross([Dfx, Dfy, Bz])
+
+
+def curl_h(wavenumber: complex, dxes: dx_lists_t) -> sparse.spmatrix:
+    """
+    Discretized curl operator for use with the waveguide H field.
+
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :return: Sparse matrix representation of the operator
+    """
+    n = 1
+    for d in dxes[1]:
+        n *= len(d)
+
+    Bz = -1j * wavenumber * sparse.eye(n)
+    Dbx, Dby = operators.deriv_back(dxes[1])
+    return operators.cross([Dbx, Dby, Bz])
+
+
+def h_err(h: vfield_t,
+          wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> float:
+    """
+    Calculates the relative error in the H field
+
+    :param h: Vectorized H field
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Relative error norm(OP @ h) / norm(h)
+    """
+    ce = curl_e(wavenumber, dxes)
+    ch = curl_h(wavenumber, dxes)
+
+    eps_inv = sparse.diags(1 / epsilon)
+
+    if numpy.any(numpy.equal(mu, None)):
+        op = ce @ eps_inv @ ch @ h - omega ** 2 * h
+    else:
+        op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h)
+
+    return norm(op) / norm(h)
+
+
+def e_err(e: vfield_t,
+          wavenumber: complex,
+          omega: complex,
+          dxes: dx_lists_t,
+          epsilon: vfield_t,
+          mu: vfield_t = None
+          ) -> float:
+    """
+    Calculates the relative error in the E field
+
+    :param e: Vectorized E field
+    :param wavenumber: Wavenumber satisfying A @ v == wavenumber**2 * v
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param mu: Vectorized magnetic permeability grid (default 1 everywhere)
+    :return: Relative error norm(OP @ e) / norm(e)
+    """
+    ce = curl_e(wavenumber, dxes)
+    ch = curl_h(wavenumber, dxes)
+
+    if numpy.any(numpy.equal(mu, None)):
+        op = ch @ ce @ e - omega ** 2 * (epsilon * e)
+    else:
+        mu_inv = sparse.diags(1 / mu)
+        op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
+
+    return norm(op) / norm(e)

fdfd_tools/waveguide_mode.py

@@ -0,0 +1,301 @@
+from typing import Dict, List
+import numpy
+import scipy.sparse as sparse
+import scipy.sparse.linalg as spalg
+
+from . import vec, unvec, dx_lists_t, vfield_t, field_t
+from . import operators, waveguide, functional
+
+
+def solve_waveguide_mode_2d(mode_number: int,
+                            omega: complex,
+                            dxes: dx_lists_t,
+                            epsilon: vfield_t,
+                            mu: vfield_t = None,
+                            wavenumber_correction: bool = True
+                            ) -> Dict[str, complex or field_t]:
+    """
+    Given a 2d region, attempts to solve for the eigenmode with the specified mode number.
+
+    :param mode_number: Number of the mode, 0-indexed
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param epsilon: Dielectric constant
+    :param mu: Magnetic permeability (default 1 everywhere)
+    :param wavenumber_correction: Whether to correct the wavenumber to
+        account for numerical dispersion (default True)
+    :return: {'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}
+    """
+
+    '''
+    Solve for the largest-magnitude eigenvalue of the real operator
+     by using power iteration.
+    '''
+    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))
+
+    # Use power iteration for 20 steps to estimate the dominant eigenvector
+    v = numpy.random.rand(A_r.shape[0])
+    for _ in range(20):
+        v = A_r @ v
+        v /= numpy.linalg.norm(v)
+
+    lm_eigval = v @ A_r @ v
+
+    '''
+    Shift by the absolute value of the largest eigenvalue, then find a few of the
+     largest (shifted) eigenvalues. The shift ensures that we find the largest
+     _positive_ eigenvalues, since any negative eigenvalues will be shifted to the range
+     0 >= neg_eigval + abs(lm_eigval) > abs(lm_eigval)
+    '''
+    shifted_A_r = A_r + abs(lm_eigval) * sparse.eye(A_r.shape[0])
+    eigvals, eigvecs = spalg.eigs(shifted_A_r, which='LM', k=mode_number + 3, ncv=50)
+
+    # Pick the eigenvalue we want from the few we found
+    k = eigvals.argsort()[-(mode_number+1)]
+    v = eigvecs[:, k]
+
+    '''
+    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 = None
+    for _ in range(40):
+        eigval = v @ A @ v
+        if numpy.linalg.norm(A @ v - eigval * v) < 1e-13:
+            break
+        w = spalg.spsolve(A - eigval * sparse.eye(A.shape[0]), v)
+        v = w / numpy.linalg.norm(w)
+
+    # Calculate the wave-vector (force the real part to be positive)
+    wavenumber = numpy.sqrt(eigval)
+    wavenumber *= numpy.sign(numpy.real(wavenumber))
+
+    e, h = waveguide.normalized_fields(v, wavenumber, omega, dxes, epsilon, mu)
+
+    '''
+    Perform correction on wavenumber to account for numerical dispersion.
+
+     See Numerical Dispersion in Taflove's FDTD book.
+     This correction term reduces the error in emitted power, but additional
+      error is introduced into the E_err and H_err terms. This effect becomes
+      more pronounced as beta increases.
+    '''
+    if wavenumber_correction:
+        wavenumber -= 2 * numpy.sin(numpy.real(wavenumber / 2)) - numpy.real(wavenumber)
+
+    shape = [d.size for d in dxes[0]]
+    fields = {
+        'wavenumber': wavenumber,
+        'E': unvec(e, shape),
+        'H': unvec(h, shape),
+    }
+
+    return fields
+
+
+def solve_waveguide_mode(mode_number: int,
+                         omega: complex,
+                         dxes: dx_lists_t,
+                         axis: int,
+                         polarity: int,
+                         slices: List[slice],
+                         epsilon: field_t,
+                         mu: field_t = None,
+                         wavenumber_correction: bool = True
+                         ) -> Dict[str, complex or numpy.ndarray]:
+    """
+    Given a 3D grid, selects a slice from the grid and attempts to
+     solve for an eigenmode propagating through that slice.
+
+    :param mode_number: Number of the mode, 0-indexed
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param axis: Propagation axis (0=x, 1=y, 2=z)
+    :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
+    :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
+        as the waveguide cross-section. slices[axis] should select only one
+    :param epsilon: Dielectric constant
+    :param mu: Magnetic permeability (default 1 everywhere)
+    :param wavenumber_correction: Whether to correct the wavenumber to
+        account for numerical dispersion (default True)
+    :return: {'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}
+    """
+    if mu is None:
+        mu = [numpy.ones_like(epsilon[0])] * 3
+
+    '''
+    Solve the 2D problem in the specified plane
+    '''
+    # Define rotation to set z as propagation direction
+    order = numpy.roll(range(3), 2 - axis)
+    reverse_order = numpy.roll(range(3), axis - 2)
+
+    # Reduce to 2D and solve the 2D problem
+    args_2d = {
+        'dxes': [[dx[i][slices[i]] for i in order[:2]] for dx in dxes],
+        'epsilon': vec([epsilon[i][slices].transpose(order) for i in order]),
+        'mu': vec([mu[i][slices].transpose(order) for i in order]),
+        'wavenumber_correction': wavenumber_correction,
+    }
+    fields_2d = solve_waveguide_mode_2d(mode_number, omega=omega, **args_2d)
+
+    '''
+    Apply corrections and expand to 3D
+    '''
+    # Scale based on dx in propagation direction
+    dxab_forward = numpy.array([dx[order[2]][slices[order[2]]] for dx in dxes])
+
+    # Adjust for propagation direction
+    fields_2d['E'][2] *= polarity
+    fields_2d['H'][2] *= polarity
+
+    # Apply phase shift to H-field
+    d_prop = 0.5 * sum(dxab_forward)
+    for a in range(3):
+        fields_2d['H'][a] *= numpy.exp(-polarity * 1j * 0.5 * fields_2d['wavenumber'] * d_prop)
+
+    # Expand E, H to full epsilon space we were given
+    E = [None]*3
+    H = [None]*3
+    for a, o in enumerate(reverse_order):
+        E[a] = numpy.zeros_like(epsilon[0], dtype=complex)
+        H[a] = numpy.zeros_like(epsilon[0], dtype=complex)
+
+        E[a][slices] = fields_2d['E'][o][:, :, None].transpose(reverse_order)
+        H[a][slices] = fields_2d['H'][o][:, :, None].transpose(reverse_order)
+
+    results = {
+        'wavenumber': fields_2d['wavenumber'],
+        'H': H,
+        'E': E,
+    }
+
+    return results
+
+
+def compute_source(E: field_t,
+                   H: field_t,
+                   wavenumber: complex,
+                   omega: complex,
+                   dxes: dx_lists_t,
+                   axis: int,
+                   polarity: int,
+                   slices: List[slice],
+                   mu: field_t = None,
+                   ) -> field_t:
+    """
+    Given an eigenmode obtained by solve_waveguide_mode, returns the current source distribution
+    necessary to position a unidirectional source at the slice location.
+
+    :param E: E-field of the mode
+    :param H: H-field of the mode (advanced by half of a Yee cell from E)
+    :param wavenumber: Wavenumber of the mode
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param axis: Propagation axis (0=x, 1=y, 2=z)
+    :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
+    :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
+        as the waveguide cross-section. slices[axis] should select only one
+    :param mu: Magnetic permeability (default 1 everywhere)
+    :return: J distribution for the unidirectional source
+    """
+    if mu is None:
+        mu = [1] * 3
+
+    J = [None]*3
+    M = [None]*3
+
+    src_order = numpy.roll(range(3), axis)
+    exp_iphi = numpy.exp(1j * polarity * wavenumber * dxes[1][axis][slices[axis]])
+    J[src_order[0]] = numpy.zeros_like(E[0])
+    J[src_order[1]] = +exp_iphi * H[src_order[2]] * polarity
+    J[src_order[2]] = -exp_iphi * H[src_order[1]] * polarity
+
+    M[src_order[0]] = numpy.zeros_like(E[0])
+    M[src_order[1]] = +numpy.roll(E[src_order[2]], -1, axis=axis)
+    M[src_order[2]] = -numpy.roll(E[src_order[1]], -1, axis=axis)
+
+    A1f = functional.curl_h(dxes)
+
+    Jm_iw = A1f([M[k] / mu[k] for k in range(3)])
+    for k in range(3):
+        J[k] += Jm_iw[k] / (-1j * omega)
+
+    return J
+
+
+def compute_overlap_e(E: field_t,
+                      H: field_t,
+                      wavenumber: complex,
+                      omega: complex,
+                      dxes: dx_lists_t,
+                      axis: int,
+                      polarity: int,
+                      slices: List[slice],
+                      mu: field_t = None,
+                      ) -> field_t:
+    """
+    Given an eigenmode obtained by solve_waveguide_mode, calculates overlap_e for the
+    mode orthogonality relation Integrate(((E x H_mode) + (E_mode x H)) dot dn)
+    [assumes reflection symmetry].
+
+    overlap_e makes use of the e2h operator to collapse the above expression into
+     (vec(E) @ vec(overlap_e)), allowing for simple calculation of the mode overlap.
+
+    :param E: E-field of the mode
+    :param H: H-field of the mode (advanced by half of a Yee cell from E)
+    :param wavenumber: Wavenumber of the mode
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header
+    :param axis: Propagation axis (0=x, 1=y, 2=z)
+    :param polarity: Propagation direction (+1 for +ve, -1 for -ve)
+    :param slices: epsilon[tuple(slices)] is used to select the portion of the grid to use
+        as the waveguide cross-section. slices[axis] should select only one
+    :param mu: Magnetic permeability (default 1 everywhere)
+    :return: overlap_e for calculating the mode overlap
+    """
+    cross_plane = [slice(None)] * 3
+    cross_plane[axis] = slices[axis]
+
+    # Determine phase factors for parallel slices
+    a_shape = numpy.roll([-1, 1, 1], axis)
+    a_E = numpy.real(dxes[0][axis]).cumsum()
+    a_H = numpy.real(dxes[1][axis]).cumsum()
+    iphi = -polarity * 1j * wavenumber
+    phase_E = numpy.exp(iphi * (a_E - a_E[slices[axis]])).reshape(a_shape)
+    phase_H = numpy.exp(iphi * (a_H - a_H[slices[axis]])).reshape(a_shape)
+
+    # Expand our slice to the entire grid using the calculated phase factors
+    Ee = [None]*3
+    He = [None]*3
+    for k in range(3):
+        Ee[k] = phase_E * E[k][tuple(cross_plane)]
+        He[k] = phase_H * H[k][tuple(cross_plane)]
+
+
+    # Write out the operator product for the mode orthogonality integral
+    domain = numpy.zeros_like(E[0], dtype=int)
+    domain[slices] = 1
+
+    npts = E[0].size
+    dn = numpy.zeros(npts * 3, dtype=int)
+    dn[0:npts] = 1
+    dn = numpy.roll(dn, npts * axis)
+
+    e2h = operators.e2h(omega, dxes, mu)
+    ds = sparse.diags(vec([domain]*3))
+    h_cross_ = operators.poynting_h_cross(vec(He), dxes)
+    e_cross_ = operators.poynting_e_cross(vec(Ee), dxes)
+
+    overlap_e = dn @ ds @ (-h_cross_ + e_cross_ @ e2h)
+
+    # Normalize
+    dx_forward = dxes[0][axis][slices[axis]]
+    norm_factor = numpy.abs(overlap_e @ vec(Ee))
+    overlap_e /= norm_factor * dx_forward
+
+    return unvec(overlap_e, E[0].shape)

float_raster.py

@@ -0,0 +1 @@
+../float_raster/float_raster.py

gridlock

@@ -0,0 +1 @@
+../gridlock/gridlock/

setup.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+
+from setuptools import setup, find_packages
+
+setup(name='fdfd_tools',
+      version='0.1',
+      description='FDFD Electromagnetic simulation tools',
+      author='Jan Petykiewicz',
+      author_email='anewusername@gmail.com',
+      url='https://mpxd.net/gogs/jan/fdfd_tools',
+      packages=find_packages(),
+      install_requires=[
+            'numpy',
+            'scipy',
+      ],
+      extras_require={
+      },
+      )