diff --git a/python/adjoint/basis.py b/python/adjoint/basis.py
index 3e138ee42..25370f272 100644
--- a/python/adjoint/basis.py
+++ b/python/adjoint/basis.py
@@ -14,13 +14,11 @@
 class Basis(ABC):
     """
     """
-    def __init__(
-            self,
-            rho_vector=None,
-            volume=None,
-            size=None,
-            center=mp.Vector3(),
-    ):
+    def __init__(self,
+                 rho_vector=None,
+                 volume=None,
+                 size=None,
+                 center=mp.Vector3()):
         self.volume = volume if volume else mp.Volume(center=center, size=size)
         self.rho_vector = rho_vector
 
@@ -54,6 +52,9 @@ class BilinearInterpolationBasis(Basis):
     Simple bilinear interpolation basis set.
     '''
     def __init__(self, resolution, symmetry=None, **kwargs):
+        ''' 
+        
+        '''
         self.dim = 2
 
         super(BilinearInterpolationBasis, self).__init__(**kwargs)
@@ -194,14 +195,8 @@ def get_bilinear_row(self, rx, ry, rho_x, rho_y):
 
         return weights, interp_idx
 
-    def gen_interpolation_matrix(
-        self,
-        rho_x,
-        rho_y,
-        rho_x_interp,
-        rho_y_interp,
-        rho_z_interp,
-    ):
+    def gen_interpolation_matrix(self, rho_x, rho_y, rho_x_interp,
+                                 rho_y_interp, rho_z_interp):
         '''
         Generates a bilinear interpolation matrix.
         
diff --git a/python/adjoint/filter_source.py b/python/adjoint/filter_source.py
index cce5fc881..b69ef2d80 100644
--- a/python/adjoint/filter_source.py
+++ b/python/adjoint/filter_source.py
@@ -4,14 +4,12 @@
 
 
 class FilteredSource(CustomSource):
-    def __init__(
-        self,
-        center_frequency,
-        frequencies,
-        frequency_response,
-        dt,
-        time_src=None,
-    ):
+    def __init__(self,
+                 center_frequency,
+                 frequencies,
+                 frequency_response,
+                 dt,
+                 time_src=None):
         dt = dt / 2  # divide by two to compensate for staggered E,H time interval
         self.dt = dt
         self.frequencies = frequencies
diff --git a/python/adjoint/filters.py b/python/adjoint/filters.py
index b2605f740..833cf2df5 100644
--- a/python/adjoint/filters.py
+++ b/python/adjoint/filters.py
@@ -3,14 +3,15 @@
 """
 
 import numpy as np
-from autograd import numpy as npa
+import jax
+from jax import numpy as npj
 import meep as mp
 from scipy import special
+import jax.numpy.fft as fft
 
 
-def _centered(arr, newshape):
+def centered(arr, newshape):
     '''Helper function that reformats the padded array of the fft filter operation.
-
     Borrowed from scipy:
     https://github.com/scipy/scipy/blob/v1.4.1/scipy/signal/signaltools.py#L263-L270
     '''
@@ -23,282 +24,139 @@ def _centered(arr, newshape):
     return arr[tuple(myslice)]
 
 
-def _edge_pad(arr, pad):
-
-    # fill sides
-    left = npa.tile(arr[0, :], (pad[0][0], 1))  # left side
-    right = npa.tile(arr[-1, :], (pad[0][1], 1))  # right side
-    top = npa.tile(arr[:, 0], (pad[1][0], 1)).transpose()  # top side
-    bottom = npa.tile(arr[:, -1], (pad[1][1], 1)).transpose()  # bottom side)
-
-    # fill corners
-    top_left = npa.tile(arr[0, 0], (pad[0][0], pad[1][0]))  # top left
-    top_right = npa.tile(arr[-1, 0], (pad[0][1], pad[1][0]))  # top right
-    bottom_left = npa.tile(arr[0, -1], (pad[0][0], pad[1][1]))  # bottom left
-    bottom_right = npa.tile(arr[-1, -1],
-                            (pad[0][1], pad[1][1]))  # bottom right
-
-    out = npa.concatenate((npa.concatenate(
-        (top_left, top, top_right)), npa.concatenate((left, arr, right)),
-                           npa.concatenate(
-                               (bottom_left, bottom, bottom_right))),
-                          axis=1)
-
-    return out
+def atleast_3d(ary):
+    if ary.ndim == 0:
+        result = npj.expand_dims(ary, axis=(0, 1, 2))  #ary.reshape(1, 1, 1)
+    elif ary.ndim == 1:
+        result = npj.expand_dims(ary, axis=(1, 2))
+    elif ary.ndim == 2:
+        result = npj.expand_dims(ary, axis=(2))
+    else:
+        result = ary
+    return result
 
 
-def _zero_pad(arr, pad):
+def _centered(arr, newshape):
+    '''Helper function that reformats the padded array of the fft filter operation.
 
-    # fill sides
-    left = npa.tile(0, (pad[0][0], arr.shape[1]))  # left side
-    right = npa.tile(0, (pad[0][1], arr.shape[1]))  # right side
-    top = npa.tile(0, (arr.shape[0], pad[1][0]))  # top side
-    bottom = npa.tile(0, (arr.shape[0], pad[1][1]))  # bottom side
+    Borrowed from scipy:
+    https://github.com/scipy/scipy/blob/v1.4.1/scipy/signal/signaltools.py#L263-L270
+    '''
+    # Return the center newshape portion of the array.
+    newshape = npj.array(newshape)
+    currshape = npj.array(arr.shape)
+    startind = (currshape - newshape) // 2
+    endind = startind + newshape
+    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
+    return arr[tuple(myslice)]
 
-    # fill corners
-    top_left = npa.tile(0, (pad[0][0], pad[1][0]))  # top left
-    top_right = npa.tile(0, (pad[0][1], pad[1][0]))  # top right
-    bottom_left = npa.tile(0, (pad[0][0], pad[1][1]))  # bottom left
-    bottom_right = npa.tile(0, (pad[0][1], pad[1][1]))  # bottom right
 
-    out = npa.concatenate((npa.concatenate(
-        (top_left, top, top_right)), npa.concatenate((left, arr, right)),
-                           npa.concatenate(
-                               (bottom_left, bottom, bottom_right))),
-                          axis=1)
+def meep_filter(x, kernel):
+    '''
+    General convolution for 1D, 2D, and 3D design spaces.
 
-    return out
+    Currently, Jax's XLA convolution kernel *always* uses 
+    the overlap-add method, which is very costly for 
+    the large arrays we typically convolve.
 
+    In the meantime, we can do a manual n-dimensional
+    FFT (also using jax). Note, however, that Jax
+    doesn't support automatic padding (like numpy does)
+    so we need to do the padding ourselves (in addition
+    to the edge padding we use for boundary conditions).
 
-def simple_2d_filter(x, kernel, Lx, Ly, resolution, symmetries=[]):
-    """A simple 2d filter algorithm that is differentiable with autograd.
-    Uses a 2D fft approach since it is typically faster and preserves the shape
-    of the input and output arrays.
-    
-    The ffts pad the operation to prevent any circular convolution garbage.
+    Note that we need to fftshift the output since we 
+    padded *both* sides of our signals.
 
     Parameters
     ----------
-    x : array_like (2D)
-        Input array to be filtered. Must be 2D.
-    kernel : array_like (2D)
+    x : array_like (1D, 2D, or 3D)
+        Input array to be filtered.
+    kernel : array_like (1D, 2D, or 3D)
         Filter kernel (before the DFT). Must be same size as `x`
-    Lx : float
-        Length of design region in X direction (in "meep units")
-    Ly : float
-        Length of design region in Y direction (in "meep units")
-    resolution : int
-        Resolution of the design grid (not the meep simulation resolution)
-    symmetries : list
-        Symmetries to impose on the parameter field (either mp.X or mp.Y)
-    
-    Returns
-    -------
-    array_like (2D)
-        The output of the 2d convolution.
-    """
-    # Get 2d parameter space shape
-    Nx = int(Lx * resolution)
-    Ny = int(Ly * resolution)
-    (kx, ky) = kernel.shape
-
-    # Adjust parameter space shape for symmetries
-    if mp.X in symmetries:
-        Nx = int(Nx / 2)
-    if mp.Y in symmetries:
-        Ny = int(Ny / 2)
-
-    # Ensure the input is 2D
-    x = x.reshape(Nx, Ny)
-
-    # Perform the required reflections for symmetries
-    if mp.X in symmetries:
-        if kx % 2 == 1:
-            x = npa.concatenate((x, x[-1, :][None, :], x[::-1, :]), axis=0)
-        else:
-            x = npa.concatenate((x, x[::-1, :]), axis=0)
-    if mp.Y in symmetries:
-        if ky % 2 == 1:
-            x = npa.concatenate((x[:, ::-1], x[:, -1][:, None], x), axis=1)
-        else:
-            x = npa.concatenate((x[:, ::-1], x), axis=1)
-
-    # pad the kernel and input to avoid circular convolution and
-    # to ensure boundary conditions are met.
-    kernel = _zero_pad(kernel, ((kx, kx), (ky, ky)))
-    x = _edge_pad(x, ((kx, kx), (ky, ky)))
-
-    # Transform to frequency domain for fast convolution
-    H = npa.fft.fft2(kernel)
-    X = npa.fft.fft2(x)
-
-    # Convolution (multiplication in frequency domain)
-    Y = H * X
-
-    # We need to fftshift since we padded both sides if each dimension of our input and kernel.
-    y = npa.fft.fftshift(npa.real(npa.fft.ifft2(Y)))
-
-    # Remove all the extra padding
-    y = _centered(y, (kx, ky))
-
-    # Remove the added symmetry domains
-    if mp.X in symmetries:
-        y = y[0:Nx, :]
-    if mp.Y in symmetries:
-        y = y[:, -Ny:]
-
-    return y
-
-
-def cylindrical_filter(x, radius, Lx, Ly, resolution, symmetries=[]):
-    '''A uniform cylindrical filter [1]. Typically allows for sharper transitions. 
     
-    Parameters
-    ----------
-    x : array_like (2D)
-        Design parameters
-    radius : float
-        Filter radius (in "meep units")
-    Lx : float
-        Length of design region in X direction (in "meep units")
-    Ly : float
-        Length of design region in Y direction (in "meep units")
-    resolution : int
-        Resolution of the design grid (not the meep simulation resolution)
-    symmetries : list
-        Symmetries to impose on the parameter field (either mp.X or mp.Y)
-
     Returns
     -------
-    array_like (2D)
-        Filtered design parameters.
-    
-    References
-    ----------
-    [1] Lazarov, B. S., Wang, F., & Sigmund, O. (2016). Length scale and manufacturability in 
-    density-based topology optimization. Archive of Applied Mechanics, 86(1-2), 189-218.
+    array_like (1D, 2D, or 3D)
+        The output of the N-D convolution.
     '''
-    # Get 2d parameter space shape
-    Nx = int(Lx * resolution)
-    Ny = int(Ly * resolution)
-
-    # Formulate grid over entire design region
-    xv, yv = np.meshgrid(np.linspace(-Lx / 2, Lx / 2, Nx),
-                         np.linspace(-Ly / 2, Ly / 2, Ny),
-                         sparse=True,
-                         indexing='ij')
-
-    # Calculate kernel
-    kernel = np.where(np.abs(xv**2 + yv**2) <= radius**2, 1, 0).T
+    # store shapes
+    x_shape = x.shape
+    k_shape = kernel.shape
 
-    # Normalize kernel
-    kernel = kernel / np.sum(kernel.flatten())  # Normalize the filter
+    # edge pad
+    npad = *((s, s) for s in x_shape),
+    kernelp = npj.pad(kernel, npad, mode='edge')
+    xp = npj.pad(x, npad, mode='edge')
 
-    # Filter the response
-    y = simple_2d_filter(x, kernel, Lx, Ly, resolution, symmetries)
+    # convolve
+    '''
+    As is the case with any convolution, the filter will
+    introduce a delay in every specified dimension. To 
+    compensate for this delay, we can perform a zero-phase
+    filter operation. We first convolve normally, and then
+    "flip" the signal and convolve again, effectively eliminating
+    the phase delay. This effectively applies the filter kernel
+    twice, however. We sqrt the kernel in the frequency domain
+    to somewhat compensate for this.
+    '''
 
-    return y
+    yp = fft.fftshift(fft.ifftn(fft.fftn(xp) * (fft.fftn(kernelp)))).real
+    yp = npj.flip(fft.fftshift(
+        fft.ifftn(fft.fftn(npj.flip(yp, axis=None)) * fft.fftn(kernelp))).real,
+                  axis=None)
 
+    # remove paddings
+    return _centered(yp, x_shape)
 
-def conic_filter(x, radius, Lx, Ly, resolution, symmetries=[]):
-    '''A linear conic filter, also known as a "Hat" filter in the literature [1].
-    
-    Parameters
-    ----------
-    x : array_like (2D)
-        Design parameters
-    radius : float
-        Filter radius (in "meep units")
-    Lx : float
-        Length of design region in X direction (in "meep units")
-    Ly : float
-        Length of design region in Y direction (in "meep units")
-    resolution : int
-        Resolution of the design grid (not the meep simulation resolution)
-    symmetries : list
-        Symmetries to impose on the parameter field (either mp.X or mp.Y)
 
-    Returns
-    -------
-    array_like (2D)
-        Filtered design parameters.
-    
-    References
-    ----------
-    [1] Lazarov, B. S., Wang, F., & Sigmund, O. (2016). Length scale and manufacturability in 
-    density-based topology optimization. Archive of Applied Mechanics, 86(1-2), 189-218.
-    '''
-    # Get 2d parameter space shape
-    Nx = int(Lx * resolution)
-    Ny = int(Ly * resolution)
+def cylindrical_filter(x, radius):
+    # filter radius is in pixels, so make sure radius >= 1
+    if radius < 1:
+        raise ValueError("The filter radius must be at least 1 pixel")
 
-    # Formulate grid over entire design region
-    xv, yv = np.meshgrid(np.linspace(-Lx / 2, Lx / 2, Nx),
-                         np.linspace(-Ly / 2, Ly / 2, Ny),
-                         sparse=True,
-                         indexing='ij')
+    x = atleast_3d(x)
+    xl = npj.linspace(-x.shape[0] / 2, x.shape[0] / 2, x.shape[0])
+    yl = npj.linspace(-x.shape[1] / 2, x.shape[1] / 2, x.shape[1])
+    zl = npj.linspace(-x.shape[2] / 2, x.shape[2] / 2, x.shape[2])
 
-    # Calculate kernel
-    kernel = np.where(
-        np.abs(xv**2 + yv**2) <= radius**2,
-        (1 - np.sqrt(abs(xv**2 + yv**2)) / radius), 0)
+    X, Y, Z = npj.meshgrid(xl, yl, zl, sparse=True, indexing='ij')
+    kernel = X**2 + Y**2 + Z**2 <= radius**2
+    kernel = kernel / npj.sum(kernel.flatten())  # normalize
 
-    # Normalize kernel
-    kernel = kernel / np.sum(kernel.flatten())  # Normalize the filter
+    return meep_filter(npj.squeeze(x), npj.squeeze(kernel))
 
-    # Filter the response
-    y = simple_2d_filter(x, kernel, Lx, Ly, resolution, symmetries)
 
-    return y
+def conic_filter(x, radius):
+    x = atleast_3d(x)
+    xl = npj.linspace(-x.shape[0] / 2, x.shape[0] / 2, x.shape[0])
+    yl = npj.linspace(-x.shape[1] / 2, x.shape[1] / 2, x.shape[1])
+    zl = npj.linspace(-x.shape[2] / 2, x.shape[2] / 2, x.shape[2])
 
+    X, Y, Z = npj.meshgrid(xl, yl, zl, sparse=True, indexing='ij')
+    kernel = np.where(
+        np.abs(X**2 + Y**2 + Z**2) <= radius**2,
+        (1 - np.sqrt(npj.abs(X**2 + Y**2 + Z**2)) / radius), 0)
+    kernel = kernel / npj.sum(kernel.flatten())  # normalize
 
-def gaussian_filter(x, sigma, Lx, Ly, resolution, symmetries=[]):
-    '''A simple gaussian filter of the form exp(-x **2 / sigma ** 2) [1].
-    
-    Parameters
-    ----------
-    x : array_like (2D)
-        Design parameters
-    sigma : float
-        Filter radius (in "meep units")
-    Lx : float
-        Length of design region in X direction (in "meep units")
-    Ly : float
-        Length of design region in Y direction (in "meep units")
-    resolution : int
-        Resolution of the design grid (not the meep simulation resolution)
+    return meep_filter(npj.squeeze(x), npj.squeeze(kernel))
 
-    Returns
-    -------
-    array_like (2D)
-        Filtered design parameters.
-    
-    References
-    ----------
-    [1] Wang, E. W., Sell, D., Phan, T., & Fan, J. A. (2019). Robust design of 
-    topology-optimized metasurfaces. Optical Materials Express, 9(2), 469-482.
-    '''
-    # Get 2d parameter space shape
-    Nx = int(Lx * resolution)
-    Ny = int(Ly * resolution)
 
+def gaussian_filter(x, sigma):
     gaussian = lambda x, sigma: np.exp(-x**2 / sigma**2)
+    x = atleast_3d(x)
+    xl = npj.linspace(-x.shape[0] / 2, x.shape[0] / 2, x.shape[0])
+    yl = npj.linspace(-x.shape[1] / 2, x.shape[1] / 2, x.shape[1])
+    zl = npj.linspace(-x.shape[2] / 2, x.shape[2] / 2, x.shape[2])
 
-    # Formulate grid over entire design region
-    xv = np.linspace(-Lx / 2, Lx / 2, Nx)
-    yv = np.linspace(-Ly / 2, Ly / 2, Ny)
+    X, Y, Z = npj.meshgrid(xl, yl, zl, sparse=True, indexing='ij')
+    kernel = np.multiply.outer(np.outer(gaussian(X, sigma), gaussian(Y,
+                                                                     sigma)),
+                               gaussian(Z, sigma))  # Gaussian filter kernel
+    kernel = kernel / npj.sum(kernel.flatten())  # normalize
 
-    # Calculate kernel
-    kernel = np.outer(gaussian(xv, sigma),
-                      gaussian(yv, sigma))  # Gaussian filter kernel
+    return meep_filter(npj.squeeze(x), npj.squeeze(kernel))
 
-    # Normalize kernel
-    kernel = kernel / np.sum(kernel.flatten())  # Normalize the filter
-
-    # Filter the response
-    y = simple_2d_filter(x, kernel, Lx, Ly, resolution, symmetries)
-
-    return y
 
 
 '''
@@ -307,7 +165,7 @@ def gaussian_filter(x, sigma, Lx, Ly, resolution, symmetries=[]):
 '''
 
 
-def exponential_erosion(x, radius, beta, Lx, Ly, resolution):
+def exponential_erosion(x, radius, beta):
     ''' Performs and exponential erosion operation.
     
     Parameters
@@ -339,12 +197,11 @@ def exponential_erosion(x, radius, beta, Lx, Ly, resolution):
     and Multidisciplinary Optimization, 54(1), 15-21.
     '''
 
-    x_hat = npa.exp(beta * (1 - x))
-    return 1 - npa.log(
-        cylindrical_filter(x_hat, radius, Lx, Ly, resolution).flatten()) / beta
+    x_hat = npj.exp(beta * (1 - x))
+    return 1 - npj.log(cylindrical_filter(x_hat, radius)) / beta
 
 
-def exponential_dilation(x, radius, beta, Lx, Ly, resolution):
+def exponential_dilation(x, radius, beta):
     ''' Performs a exponential dilation operation.
     
     Parameters
@@ -376,12 +233,11 @@ def exponential_dilation(x, radius, beta, Lx, Ly, resolution):
     and Multidisciplinary Optimization, 54(1), 15-21.
     '''
 
-    x_hat = npa.exp(beta * x)
-    return npa.log(
-        cylindrical_filter(x_hat, radius, Lx, Ly, resolution).flatten()) / beta
+    x_hat = npj.exp(beta * x)
+    return npj.log(cylindrical_filter(x_hat, radius)) / beta
 
 
-def heaviside_erosion(x, radius, beta, Lx, Ly, resolution):
+def heaviside_erosion(x, radius, beta):
     ''' Performs a heaviside erosion operation.
     
     Parameters
@@ -411,11 +267,11 @@ def heaviside_erosion(x, radius, beta, Lx, Ly, resolution):
     numerical methods in engineering, 61(2), 238-254.
     '''
 
-    x_hat = cylindrical_filter(x, radius, Lx, Ly, resolution).flatten()
-    return npa.exp(-beta * (1 - x_hat)) + npa.exp(-beta) * (1 - x_hat)
+    x_hat = cylindrical_filter(x, radius)
+    return npj.exp(-beta * (1 - x_hat)) + npj.exp(-beta) * (1 - x_hat)
 
 
-def heaviside_dilation(x, radius, beta, Lx, Ly, resolution):
+def heaviside_dilation(x, radius, beta):
     ''' Performs a heaviside dilation operation.
     
     Parameters
@@ -445,11 +301,11 @@ def heaviside_dilation(x, radius, beta, Lx, Ly, resolution):
     numerical methods in engineering, 61(2), 238-254.
     '''
 
-    x_hat = cylindrical_filter(x, radius, Lx, Ly, resolution).flatten()
-    return 1 - npa.exp(-beta * x_hat) + npa.exp(-beta) * x_hat
+    x_hat = cylindrical_filter(x, radius)
+    return 1 - npj.exp(-beta * x_hat) + npj.exp(-beta) * x_hat
 
 
-def geometric_erosion(x, radius, alpha, Lx, Ly, resolution):
+def geometric_erosion(x, radius, alpha):
     ''' Performs a geometric erosion operation.
     
     Parameters
@@ -477,12 +333,11 @@ def geometric_erosion(x, radius, alpha, Lx, Ly, resolution):
     [1] Svanberg, K., & Svärd, H. (2013). Density filters for topology optimization based on the 
     Pythagorean means. Structural and Multidisciplinary Optimization, 48(5), 859-875.
     '''
-    x_hat = npa.log(x + alpha)
-    return npa.exp(cylindrical_filter(x_hat, radius, Lx, Ly,
-                                      resolution)).flatten() - alpha
+    x_hat = npj.log(x + alpha)
+    return npj.exp(cylindrical_filter(x_hat, radius)) - alpha
 
 
-def geometric_dilation(x, radius, alpha, Lx, Ly, resolution):
+def geometric_dilation(x, radius, alpha):
     ''' Performs a geometric dilation operation.
     
     Parameters
@@ -511,12 +366,11 @@ def geometric_dilation(x, radius, alpha, Lx, Ly, resolution):
     Pythagorean means. Structural and Multidisciplinary Optimization, 48(5), 859-875.
     '''
 
-    x_hat = npa.log(1 - x + alpha)
-    return -npa.exp(cylindrical_filter(x_hat, radius, Lx, Ly,
-                                       resolution)).flatten() + alpha + 1
+    x_hat = npj.log(1 - x + alpha)
+    return -npj.exp(cylindrical_filter(x_hat, radius)) + alpha + 1
 
 
-def harmonic_erosion(x, radius, alpha, Lx, Ly, resolution):
+def harmonic_erosion(x, radius, alpha):
     ''' Performs a harmonic erosion operation.
     
     Parameters
@@ -546,11 +400,10 @@ def harmonic_erosion(x, radius, alpha, Lx, Ly, resolution):
     '''
 
     x_hat = 1 / (x + alpha)
-    return 1 / cylindrical_filter(x_hat, radius, Lx, Ly,
-                                  resolution).flatten() - alpha
+    return 1 / cylindrical_filter(x_hat, radius) - alpha
 
 
-def harmonic_dilation(x, radius, alpha, Lx, Ly, resolution):
+def harmonic_dilation(x, radius, alpha):
     ''' Performs a harmonic dilation operation.
     
     Parameters
@@ -580,8 +433,7 @@ def harmonic_dilation(x, radius, alpha, Lx, Ly, resolution):
     '''
 
     x_hat = 1 / (1 - x + alpha)
-    return 1 - 1 / cylindrical_filter(x_hat, radius, Lx, Ly,
-                                      resolution).flatten() + alpha
+    return 1 - 1 / cylindrical_filter(x_hat, radius) + alpha
 
 
 '''
@@ -613,9 +465,9 @@ def tanh_projection(x, beta, eta):
     formulations in topology optimization. Structural and Multidisciplinary Optimization, 43(6), 767-784.
     '''
 
-    return (npa.tanh(beta * eta) +
-            npa.tanh(beta *
-                     (x - eta))) / (npa.tanh(beta * eta) + npa.tanh(beta *
+    return (npj.tanh(beta * eta) +
+            npj.tanh(beta *
+                     (x - eta))) / (npj.tanh(beta * eta) + npj.tanh(beta *
                                                                     (1 - eta)))
 
 
@@ -642,10 +494,10 @@ def heaviside_projection(x, beta, eta):
     density-based topology optimization. Archive of Applied Mechanics, 86(1-2), 189-218.
     '''
 
-    case1 = eta * npa.exp(-beta * (eta - x) / eta) - (eta - x) * npa.exp(-beta)
-    case2 = 1 - (1 - eta) * npa.exp(-beta * (x - eta) /
-                                    (1 - eta)) - (eta - x) * npa.exp(-beta)
-    return npa.where(x < eta, case1, case2)
+    case1 = eta * npj.exp(-beta * (eta - x) / eta) - (eta - x) * npj.exp(-beta)
+    case2 = 1 - (1 - eta) * npj.exp(-beta * (x - eta) /
+                                    (1 - eta)) - (eta - x) * npj.exp(-beta)
+    return npj.where(x < eta, case1, case2)
 
 
 '''
@@ -791,18 +643,13 @@ def indicator_solid(x, c, filter_f, threshold_f, resolution):
 
     filtered_field = filter_f(x)
     design_field = threshold_f(filtered_field)
-    gradient_filtered_field = npa.gradient(filtered_field)
-    grad_mag = (gradient_filtered_field[0] *
-                resolution)**2 + (gradient_filtered_field[1] * resolution)**2
-    if grad_mag.ndim != 2:
-        raise ValueError(
-            "The gradient fields must be 2 dimensional. Check input array and filter functions."
-        )
-    I_s = design_field * npa.exp(-c * grad_mag)
-    return I_s
+    gradient_filtered_field = npj.gradient(filtered_field, resolution)
+    #gradient_filtered_field = gradient_filtered_field if isinstance(gradient_filtered_field,list) else [gradient_filtered_field]
+    grad_mag = npj.sum(npj.array(gradient_filtered_field)**2, axis=0)
+    return design_field * npj.exp(-c * grad_mag)
 
 
-def constraint_solid(x, c, eta_e, filter_f, threshold_f, resolution):
+def constraint_solid(x, c, eta_e, filter_f, threshold_f, resolution=1):
     '''Calculates the constraint function of the solid phase needed for minimum length optimization [1].
     
     Parameters
@@ -835,9 +682,9 @@ def constraint_solid(x, c, eta_e, filter_f, threshold_f, resolution):
     '''
 
     filtered_field = filter_f(x)
-    I_s = indicator_solid(x.reshape(filtered_field.shape), c, filter_f,
-                          threshold_f, resolution).flatten()
-    return npa.mean(I_s * npa.minimum(filtered_field.flatten() - eta_e, 0)**2)
+    I_s = indicator_solid(x, c, filter_f, threshold_f, resolution)
+    return npj.mean(I_s.flatten() *
+                    npj.minimum(filtered_field.flatten() - eta_e, 0)**2)
 
 
 def indicator_void(x, c, filter_f, threshold_f, resolution):
@@ -867,19 +714,15 @@ def indicator_void(x, c, filter_f, threshold_f, resolution):
     geometric constraints. Computer Methods in Applied Mechanics and Engineering, 293, 266-282.
     '''
 
-    filtered_field = filter_f(x).reshape(x.shape)
+    filtered_field = filter_f(x)
     design_field = threshold_f(filtered_field)
-    gradient_filtered_field = npa.gradient(filtered_field)
-    grad_mag = (gradient_filtered_field[0] *
-                resolution)**2 + (gradient_filtered_field[1] * resolution)**2
-    if grad_mag.ndim != 2:
-        raise ValueError(
-            "The gradient fields must be 2 dimensional. Check input array and filter functions."
-        )
-    return (1 - design_field) * npa.exp(-c * grad_mag)
+    gradient_filtered_field = npj.gradient(filtered_field, resolution)
+    #gradient_filtered_field = gradient_filtered_field if isinstance(list, gradient_filtered_field) else [gradient_filtered_field]
+    grad_mag = npj.sum(npj.array(gradient_filtered_field)**2, axis=0)
+    return (1 - design_field) * npj.exp(-c * grad_mag)
 
 
-def constraint_void(x, c, eta_d, filter_f, threshold_f, resolution):
+def constraint_void(x, c, eta_d, filter_f, threshold_f, resolution=1):
     '''Calculates the constraint function of the void phase needed for minimum length optimization [1].
     
     Parameters
@@ -912,9 +755,9 @@ def constraint_void(x, c, eta_d, filter_f, threshold_f, resolution):
     '''
 
     filtered_field = filter_f(x)
-    I_v = indicator_void(x.reshape(filtered_field.shape), c, filter_f,
-                         threshold_f, resolution).flatten()
-    return npa.mean(I_v * npa.minimum(eta_d - filtered_field.flatten(), 0)**2)
+    I_v = indicator_void(x, c, filter_f, threshold_f, resolution)
+    return npj.mean(I_v.flatten() *
+                    npj.minimum(eta_d - filtered_field.flatten(), 0)**2)
 
 
 def gray_indicator(x):
@@ -937,4 +780,34 @@ def gray_indicator(x):
     [1] Lazarov, B. S., Wang, F., & Sigmund, O. (2016). Length scale and manufacturability in 
     density-based topology optimization. Archive of Applied Mechanics, 86(1-2), 189-218.
     '''
-    return npa.mean(4 * x.flatten() * (1 - x.flatten())) * 100
+    return npj.mean(4 * x.flatten() * (1 - x.flatten()))
+
+
+def F_diff_open_close(x, f_open, f_close, beta=64, p=2):
+    '''
+
+    '''
+    return npj.mean(
+        tanh_projection(npj.abs(f_close(x) - f_open(x)).flatten()**2), beta,
+        0.5) * 100
+
+
+def M_diff_open_close(x, f_open, f_close, p=2):
+    '''
+
+    '''
+    return npj.mean(npj.abs(f_close(x) - f_open(x)).flatten()**p) * 100
+
+
+def M_diff_identity_open(x, f_open, p=2):
+    '''
+
+    '''
+    return npj.mean(npj.abs(x - f_open(x)).flatten()**p) * 100
+
+
+def M_diff_identity_close(x, f_close, p=2):
+    '''
+
+    '''
+    return npj.mean(npj.abs(x - f_close(x)).flatten()**p) * 100
diff --git a/python/adjoint/objective.py b/python/adjoint/objective.py
index 802b5571a..3a90692dc 100644
--- a/python/adjoint/objective.py
+++ b/python/adjoint/objective.py
@@ -43,7 +43,7 @@ def place_adjoint_source(self, dJ):
 
     def get_evaluation(self):
         """Evaluates the objective quantity."""
-        if self._eval:
+        if self._eval is not None:
             return self._eval
         else:
             raise RuntimeError(
@@ -330,11 +330,8 @@ def place_adjoint_source(self, dJ):
         dJ = dJ.flatten()
         farpt_list = np.array([list(pi) for pi in self.far_pts]).flatten()
         far_pt0 = self.far_pts[0]
-        far_pt_vec = py_v3_to_vec(
-            self.sim.dimensions,
-            far_pt0,
-            self.sim.is_cylindrical,
-        )
+        far_pt_vec = py_v3_to_vec(self.sim.dimensions, far_pt0,
+                                  self.sim.is_cylindrical)
 
         all_nearsrcdata = self._monitor.swigobj.near_sourcedata(
             far_pt_vec, farpt_list, self._nfar_pts, dJ)
diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py
index 244103116..fe5e20da0 100644
--- a/python/adjoint/optimization_problem.py
+++ b/python/adjoint/optimization_problem.py
@@ -1,6 +1,6 @@
 import meep as mp
 import numpy as np
-from autograd import grad, jacobian
+from jax import grad, jacobian
 from collections import namedtuple
 
 Grid = namedtuple('Grid', ['x', 'y', 'z', 'w'])
@@ -8,14 +8,12 @@
 
 
 class DesignRegion(object):
-    def __init__(
-            self,
-            design_parameters,
-            volume=None,
-            size=None,
-            center=mp.Vector3(),
-            MaterialGrid=None,
-    ):
+    def __init__(self,
+                 design_parameters,
+                 volume=None,
+                 size=None,
+                 center=mp.Vector3(),
+                 MaterialGrid=None):
         self.volume = volume if volume else mp.Volume(center=center, size=size)
         self.size = self.volume.size
         self.center = self.volume.center
@@ -108,12 +106,10 @@ def __init__(
                 fmax = fcen + 0.5 * df
                 fmin = fcen - 0.5 * df
                 dfreq = (fmax - fmin) / (nf - 1)
-                self.frequencies = np.linspace(
-                    fmin,
-                    fmin + dfreq * nf,
-                    num=nf,
-                    endpoint=False,
-                )
+                self.frequencies = np.linspace(fmin,
+                                               fmin + dfreq * nf,
+                                               num=nf,
+                                               endpoint=False)
                 self.nf = nf
 
         if self.nf == 1:
@@ -236,16 +232,9 @@ def forward_run(self):
 
         # Forward run
         self.sim.run(until_after_sources=stop_when_dft_decayed(
-            self.sim,
-            self.design_region_monitors,
-            self.decay_dt,
-            self.decay_fields,
-            self.fcen_idx,
-            self.decay_by,
-            True,
-            self.minimum_run_time,
-            self.maximum_run_time,
-        ))
+            self.sim, self.design_region_monitors, self.decay_dt,
+            self.decay_fields, self.fcen_idx, self.decay_by, True,
+            self.minimum_run_time, self.maximum_run_time))
 
         # record objective quantities from user specified monitors
         self.results_list = []
@@ -281,8 +270,8 @@ def prepare_adjoint_run(self):
         for ar in range(len(self.objective_functions)):
             for mi, m in enumerate(self.objective_arguments):
                 dJ = jacobian(self.objective_functions[ar],
-                              mi)(*self.results_list)
-                # get gradient of objective w.r.t. monitor
+                              mi)(*self.results_list
+                                  )  # get gradient of objective w.r.t. monitor
                 if np.any(dJ):
                     self.adjoint_sources[ar] += m.place_adjoint_source(
                         dJ)  # place the appropriate adjoint sources
@@ -311,16 +300,9 @@ def adjoint_run(self):
 
             # Adjoint run
             self.sim.run(until_after_sources=stop_when_dft_decayed(
-                self.sim,
-                self.design_region_monitors,
-                self.decay_dt,
-                self.decay_fields,
-                self.fcen_idx,
-                self.decay_by,
-                True,
-                self.minimum_run_time,
-                self.maximum_run_time,
-            ))
+                self.sim, self.design_region_monitors, self.decay_dt,
+                self.decay_fields, self.fcen_idx, self.decay_by, True,
+                self.minimum_run_time, self.maximum_run_time))
 
             # Store adjoint fields for each design set of design variables in array (x,y,z,field_components,frequencies)
             self.a_E.append([[
@@ -340,12 +322,9 @@ def adjoint_run(self):
     def calculate_gradient(self):
         # Iterate through all design regions and calculate gradient
         self.gradient = [[
-            dr.get_gradient(
-                self.sim,
-                self.a_E[ar][dri],
-                self.d_E[dri],
-                self.frequencies,
-            ) for dri, dr in enumerate(self.design_regions)
+            dr.get_gradient(self.sim, self.a_E[ar][dri], self.d_E[dri],
+                            self.frequencies)
+            for dri, dr in enumerate(self.design_regions)
         ] for ar in range(len(self.objective_functions))]
 
         # Cleanup list of lists
@@ -362,13 +341,11 @@ def calculate_gradient(self):
         # Return optimizer's state to initialization
         self.current_state = "INIT"
 
-    def calculate_fd_gradient(
-        self,
-        num_gradients=1,
-        db=1e-4,
-        design_variables_idx=0,
-        filter=None,
-    ):
+    def calculate_fd_gradient(self,
+                              num_gradients=1,
+                              db=1e-4,
+                              design_variables_idx=0,
+                              filter=None):
         '''
         Estimate central difference gradients.
 
@@ -405,8 +382,7 @@ def calculate_fd_gradient(
         fd_gradient_idx = np.random.choice(
             self.num_design_params[design_variables_idx],
             num_gradients,
-            replace=False,
-        )
+            replace=False)
 
         for k in fd_gradient_idx:
 
@@ -430,16 +406,9 @@ def calculate_fd_gradient(
                     m.register_monitors(self.frequencies))
 
             self.sim.run(until_after_sources=stop_when_dft_decayed(
-                self.sim,
-                self.forward_monitors,
-                self.decay_dt,
-                self.decay_fields,
-                self.fcen_idx,
-                self.decay_by,
-                True,
-                self.minimum_run_time,
-                self.maximum_run_time,
-            ))
+                self.sim, self.forward_monitors, self.decay_dt,
+                self.decay_fields, self.fcen_idx, self.decay_by, True,
+                self.minimum_run_time, self.maximum_run_time))
 
             # record final objective function value
             results_list = []
@@ -465,16 +434,9 @@ def calculate_fd_gradient(
 
             # add monitor used to track dft convergence
             self.sim.run(until_after_sources=stop_when_dft_decayed(
-                self.sim,
-                self.forward_monitors,
-                self.decay_dt,
-                self.decay_fields,
-                self.fcen_idx,
-                self.decay_by,
-                True,
-                self.minimum_run_time,
-                self.maximum_run_time,
-            ))
+                self.sim, self.forward_monitors, self.decay_dt,
+                self.decay_fields, self.fcen_idx, self.decay_by, True,
+                self.minimum_run_time, self.maximum_run_time))
 
             # record final objective function value
             results_list = []
@@ -531,17 +493,15 @@ def plot2D(self, init_opt=False, **kwargs):
         self.sim.plot2D(**kwargs)
 
 
-def stop_when_dft_decayed(
-    simob,
-    mon,
-    dt,
-    c,
-    fcen_idx,
-    decay_by,
-    yee_grid=False,
-    minimum_run_time=0,
-    maximum_run_time=None,
-):
+def stop_when_dft_decayed(simob,
+                          mon,
+                          dt,
+                          c,
+                          fcen_idx,
+                          decay_by,
+                          yee_grid=False,
+                          minimum_run_time=0,
+                          maximum_run_time=None):
     '''Step function that monitors the relative change in DFT fields for a list of monitors.
 
     mon ............. a list of monitors
diff --git a/python/adjoint/utils.py b/python/adjoint/utils.py
index 110a4c4c2..600ca8fa8 100644
--- a/python/adjoint/utils.py
+++ b/python/adjoint/utils.py
@@ -80,11 +80,11 @@ def gather_monitor_values(monitors: List[ObjectiveQuantity]) -> onp.ndarray:
     Args:
       monitors: the mode monitors.
 
-    Returns:
-      a rank-2 ndarray, where the dimensions are (monitor, frequency), of dtype
-      complex128.  Note that these values refer to the mode as oriented (i.e. they
-      are unidirectional).
-    """
+  Returns:
+    a rank-2 ndarray, where the dimensions are (monitor, frequency), of dtype
+    complex128.  Note that these values refer to the mode as oriented (i.e. they
+    are unidirectional).
+  """
     monitor_values = []
     for monitor in monitors:
         monitor_values.append(monitor())
@@ -101,21 +101,21 @@ def gather_design_region_fields(
 ) -> List[List[onp.ndarray]]:
     """Collects the design region DFT fields from the simulation.
 
-    Args:
-     simulation: the simulation object.
-     design_region_monitors: the installed design region monitors.
-     frequencies: the frequencies to monitor.
-
-    Returns:
-      A list of lists.  Each entry (list) in the overall list corresponds one-to-
-      one with a declared design region.  For each such contained list, the
-      entries correspond to the field components that are monitored.  The entries
-      are ndarrays of rank 4 with dimensions (freq, x, y, (z-or-pad)).
-
-      The design region fields are sampled on the *Yee grid*.  This makes them
-      fairly awkward to inspect directly.  Their primary use case is supporting
-      gradient calculations.
-    """
+  Args:
+   simulation: the simulation object.
+   design_region_monitors: the installed design region monitors.
+   frequencies: the frequencies to monitor.
+
+  Returns:
+    A list of lists.  Each entry (list) in the overall list corresponds one-to-
+    one with a declared design region.  For each such contained list, the
+    entries correspond to the field components that are monitored.  The entries
+    are ndarrays of rank 4 with dimensions (freq, x, y, (z-or-pad)).
+
+    The design region fields are sampled on the *Yee grid*.  This makes them
+    fairly awkward to inspect directly.  Their primary use case is supporting
+    gradient calculations.
+  """
     fwd_fields = []
     for monitor in design_region_monitors:
         fields_by_component = []
@@ -134,17 +134,17 @@ def validate_and_update_design(
         design_variables: Iterable[onp.ndarray]) -> None:
     """Validate the design regions and variables.
 
-    In particular the design variable should be 1,2,3-D and the design region
-    shape should match the design variable shape after dimension expansion.
-    The arguments are modified in place.
+  In particular the design variable should be 1,2,3-D and the design region
+  shape should match the design variable shape after dimension expansion.
+  The arguments are modified in place.
 
     Args:
       design_regions: List of mpa.DesignRegion,
       design_variables: Iterable with numpy arrays representing design variables.
 
-    Raises:
-      ValueError if the validation of dimensions fails.
-    """
+  Raises:
+    ValueError if the validation of dimensions fails.
+  """
     for i, (design_region,
             design_variable) in enumerate(zip(design_regions,
                                               design_variables)):
diff --git a/python/adjoint/wrapper.py b/python/adjoint/wrapper.py
index d36d1b89d..26d37e3df 100644
--- a/python/adjoint/wrapper.py
+++ b/python/adjoint/wrapper.py
@@ -64,48 +64,46 @@ def loss(x):
 class MeepJaxWrapper:
     """Wraps a Meep simulation object into a JAX-differentiable callable.
 
-    Attributes:
-        simulation: the pre-configured Meep simulation object.
-        sources: a list of Meep sources for the forward simulation.
-        monitors: a list of eigenmode coefficient monitors from the `meep.adjoint`
-          module.
-        design_regions: a list of design regions from the `meep.adjoint` module.
-        frequencies: the list of frequencies, in normalized Meep units.
-        measurement_interval: the time interval between DFT field convergence
-          measurements, in Meep time units. The default value is 50.
-        dft_field_components: a list of Meep field components, such as `mp.Ex`,
-          `mp.Hy`, etc, whose DFT will be monitored for convergence to stop the
-          simulation. The default is `mp.Ez`.
-        dft_threshold: the threshold for DFT field convergence. Once the norm of the
-          change in the fields (the maximum over all design regions and field
-          components) is less than this value, the simulation will be stopped. The
-          default value is 1e-6.
-        minimum_run_time: the minimum run time of the simulation, in Meep time
-          units. The default value is 0.
-        maximum_run_time: the maximum run time of the simulation, in Meep time
-          units. The default value is infinity.
-        until_after_sources: whether `maximum_run_time` should be ignored until the
-          sources have turned off. This parameter specifies whether `until` or
-          `until_after_sources` is used. See
-          https://meep.readthedocs.io/en/latest/Python_User_Interface/#Simulation
+  Attributes:
+      simulation: the pre-configured Meep simulation object.
+      sources: a list of Meep sources for the forward simulation.
+      monitors: a list of eigenmode coefficient monitors from the `meep.adjoint`
+        module.
+      design_regions: a list of design regions from the `meep.adjoint` module.
+      frequencies: the list of frequencies, in normalized Meep units.
+      measurement_interval: the time interval between DFT field convergence
+        measurements, in Meep time units. The default value is 50.
+      dft_field_components: a list of Meep field components, such as `mp.Ex`,
+        `mp.Hy`, etc, whose DFT will be monitored for convergence to stop the
+        simulation. The default is `mp.Ez`.
+      dft_threshold: the threshold for DFT field convergence. Once the norm of the
+        change in the fields (the maximum over all design regions and field
+        components) is less than this value, the simulation will be stopped. The
+        default value is 1e-6.
+      minimum_run_time: the minimum run time of the simulation, in Meep time
+        units. The default value is 0.
+      maximum_run_time: the maximum run time of the simulation, in Meep time
+        units. The default value is infinity.
+      until_after_sources: whether `maximum_run_time` should be ignored until the
+        sources have turned off. This parameter specifies whether `until` or
+        `until_after_sources` is used. See
+        https://meep.readthedocs.io/en/latest/Python_User_Interface/#Simulation
           for more information. The default is true.
-    """
+  """
     _log_fn = print
 
-    def __init__(
-        self,
-        simulation: mp.Simulation,
-        sources: List[mp.Source],
-        monitors: List[EigenmodeCoefficient],
-        design_regions: List[DesignRegion],
-        frequencies: List[float],
-        measurement_interval: float = 50.0,
-        dft_field_components: Tuple[int, ...] = (mp.Ez, ),
-        dft_threshold: float = 1e-6,
-        minimum_run_time: float = 0,
-        maximum_run_time: float = onp.inf,
-        until_after_sources: bool = True,
-    ):
+    def __init__(self,
+                 simulation: mp.Simulation,
+                 sources: List[mp.Source],
+                 monitors: List[EigenmodeCoefficient],
+                 design_regions: List[DesignRegion],
+                 frequencies: List[float],
+                 measurement_interval: float = 50.0,
+                 dft_field_components: Tuple[int, ...] = (mp.Ez, ),
+                 dft_threshold: float = 1e-6,
+                 minimum_run_time: float = 0,
+                 maximum_run_time: float = onp.inf,
+                 until_after_sources: bool = True):
         self.simulation = simulation
         self.sources = sources
         self.monitors = monitors
@@ -124,20 +122,20 @@ def __init__(
     def __call__(self, designs: List[jnp.ndarray]) -> jnp.ndarray:
         """Performs a Meep simulation, taking a list of designs and returning mode overlaps.
 
-        Args:
-          designs: a list of design variables as 1D, 2D, or 3D JAX arrays. Valid shapes for
-          design variables are (Nx, Ny, Nz) where Nx{y,z} match the elements of the
-          `grid_size` constructor argument of Meep's `MaterialGrid` used for the
-          corresponding design region. Singleton dimensions of the `grid_size` may be
-          omitted from the corresponding design variable. For example, a design variable
-          with a shape of either (10, 20) or (10, 20, 1) would be compatible with a
-          `grid_size` of (10, 20, 1). Similarly, a design variable with shapes of (25,),
-          (25, 1), or (25, 1, 1) would be compatible with a `grid_size` of (25, 1, 1).
-
-        Returns:
-          a complex-valued JAX ndarray of differentiable mode monitor overlaps values with
-          a shape of (num monitors, num frequencies).
-        """
+    Args:
+      designs: a list of design variables as 1D, 2D, or 3D JAX arrays. Valid shapes for
+      design variables are (Nx, Ny, Nz) where Nx{y,z} match the elements of the
+      `grid_size` constructor argument of Meep's `MaterialGrid` used for the
+      corresponding design region. Singleton dimensions of the `grid_size` may be
+      omitted from the corresponding design variable. For example, a design variable
+      with a shape of either (10, 20) or (10, 20, 1) would be compatible with a
+      `grid_size` of (10, 20, 1). Similarly, a design variable with shapes of (25,),
+      (25, 1), or (25, 1, 1) would be compatible with a `grid_size` of (25, 1, 1).
+
+    Returns:
+      a complex-valued JAX ndarray of differentiable mode monitor overlaps values with
+      a shape of (num monitors, num frequencies).
+    """
         return self._simulate_fn(designs)
 
     def _reset_convergence_measurement(self,
@@ -202,19 +200,15 @@ def _run_adjoint_simulation(self, monitor_values_grad):
         self._reset_convergence_measurement(design_region_monitors)
         self.simulation.run(**sim_run_args)
 
-        return utils.gather_design_region_fields(
-            self.simulation,
-            design_region_monitors,
-            self.frequencies,
-        )
+        return utils.gather_design_region_fields(self.simulation,
+                                                 design_region_monitors,
+                                                 self.frequencies)
 
-    def _calculate_vjps(
-        self,
-        fwd_fields,
-        adj_fields,
-        design_variable_shapes,
-        sum_freq_partials=True,
-    ):
+    def _calculate_vjps(self,
+                        fwd_fields,
+                        adj_fields,
+                        design_variable_shapes,
+                        sum_freq_partials=True):
         """Calculates the VJP for a given set of forward and adjoint fields."""
         return utils.calculate_vjps(
             self.simulation,
@@ -298,19 +292,19 @@ def _are_dfts_converged(self, sim: mp.Simulation) -> bool:
     def _simulation_run_callback(self, sim: mp.Simulation) -> bool:
         """A callback function returning `True` when the simulation should stop.
 
-        This is a step function that gets called at each time step of the Meep
-        simulation, taking a Meep simulation object as an input and returning `True`
-        when the simulation should be terminated and returning `False` when the
-        simulation should continue. The resulting step function is designed to be
-        used with the `until` and `until_after_sources` arguments of the Meep
-        simulation `run()` routine.
+    This is a step function that gets called at each time step of the Meep
+    simulation, taking a Meep simulation object as an input and returning `True`
+    when the simulation should be terminated and returning `False` when the
+    simulation should continue. The resulting step function is designed to be
+    used with the `until` and `until_after_sources` arguments of the Meep
+    simulation `run()` routine.
 
-        Args:
-          sim: a Meep simulation object.
+    Args:
+      sim: a Meep simulation object.
 
-        Returns:
-          a boolean indicating whether the simulation should be terminated.
-        """
+    Returns:
+      a boolean indicating whether the simulation should be terminated.
+    """
         current_meep_time = sim.round_time()
         if current_meep_time <= self._last_measurement_meep_time + self.measurement_interval:
             return False