update fdtd and add some fdtd tests

fdfd_tools/fdtd.py

@@ -3,6 +3,8 @@ import numpy
 
 from . import dx_lists_t, field_t
 
+#TODO fix pmls
+
 __author__ = 'Jan Petykiewicz'
 
 
@@ -71,7 +73,7 @@ def maxwell_e(dt: float, dxes: dx_lists_t = None) -> functional_matrix:
     curl_h_fun = curl_h(dxes)
 
     def me_fun(e: field_t, h: field_t, epsilon: field_t):
-        e += dt * curl_h_fun(h)/ epsilon
+        e += dt * curl_h_fun(h) / epsilon
         return e
 
     return me_fun
@@ -150,7 +152,12 @@ def cpml(direction:int,
          dt: float,
          epsilon: field_t,
          thickness: int = 8,
+         ln_R_per_layer: float = -1.6,
          epsilon_eff: float = 1,
+         mu_eff: float = 1,
+         m: float = 3.5,
+         ma: float = 1,
+         cfs_alpha: float = 0,
          dtype: numpy.dtype = numpy.float32,
          ) -> Tuple[Callable, Callable, Dict[str, field_t]]:
 
@@ -166,9 +173,9 @@ def cpml(direction:int,
     if epsilon_eff <= 0:
         raise Exception('epsilon_eff must be positive')
 
-    m = (3.5, 1)
-    sigma_max = 0.8 * (m[0] + 1) / numpy.sqrt(epsilon_eff)
-    alpha_max = 0  # TODO: Decide what to do about non-zero alpha
+    sigma_max = -ln_R_per_layer / 2 * (m + 1)
+    kappa_max = numpy.sqrt(epsilon_eff * mu_eff)
+    alpha_max = cfs_alpha
     transverse = numpy.delete(range(3), direction)
     u, v = transverse
 
@@ -187,14 +194,17 @@ def cpml(direction:int,
     expand_slice[direction] = slice(None)
 
     def par(x):
-        sigma = ((x / thickness) ** m[0]) * sigma_max
-        alpha = ((1 - x / thickness) ** m[1]) * alpha_max
-        p0 = numpy.exp(-(sigma + alpha) * dt)
-        p1 = sigma / (sigma + alpha) * (p0 - 1)
-        return p0[expand_slice], p1[expand_slice]
+        scaling = (x / thickness) ** m
+        sigma = scaling * sigma_max
+        kappa = 1 + scaling * (kappa_max - 1)
+        alpha = ((1 - x / thickness) ** ma) * alpha_max
+        p0 = numpy.exp(-(sigma / kappa + alpha) * dt)
+        p1 = sigma / (sigma + kappa * alpha) * (p0 - 1)
+        p2 = 1 / kappa
+        return p0[expand_slice], p1[expand_slice], p2[expand_slice]
 
-    p0e, p1e = par(xe)
-    p0h, p1h = par(xh)
+    p0e, p1e, p2e = par(xe)
+    p0h, p1h, p2h = par(xh)
 
     region = [slice(None)] * 3
     if polarity < 0:
@@ -204,12 +214,9 @@ def cpml(direction:int,
     else:
         raise Exception('Bad polarity!')
 
-    if direction == 1:
-        se = 1
-    else:
-        se = -1
+    se = 1 if direction == 1 else -1
 
-    # TODO check if epsilon is uniform?
+    # TODO check if epsilon is uniform in pml region?
     shape = list(epsilon[0].shape)
     shape[direction] = thickness
     psi_e = [numpy.zeros(shape, dtype=dtype), numpy.zeros(shape, dtype=dtype)]
@@ -222,37 +229,107 @@ def cpml(direction:int,
         'psi_h_v': psi_h[1],
     }
 
+    # Note that this is kinda slow -- would be faster to reuse dHv*p2h for the original
+    #  H update, but then you have multiple arrays and a monolithic (field + pml) update operation
     def pml_e(e: field_t, h: field_t, epsilon: field_t) -> Tuple[field_t, field_t]:
+        dHv = h[v][region] - numpy.roll(h[v], 1, axis=direction)[region]
+        dHu = h[u][region] - numpy.roll(h[u], 1, axis=direction)[region]
         psi_e[0] *= p0e
-        psi_e[0] += p1e * (h[v][region] - numpy.roll(h[v], 1, axis=direction)[region])
+        psi_e[0] += p1e * dHv * p2e
         psi_e[1] *= p0e
-        psi_e[1] += p1e * (h[u][region] - numpy.roll(h[u], 1, axis=direction)[region])
-        e[u][region] += se * dt * psi_e[0] / epsilon[u][region]
-        e[v][region] -= se * dt * psi_e[1] / epsilon[v][region]
+        psi_e[1] += p1e * dHu * p2e
+        e[u][region] += se * dt / epsilon[u][region] * (psi_e[0] + (p2e - 1) * dHv)
+        e[v][region] -= se * dt / epsilon[v][region] * (psi_e[1] + (p2e - 1) * dHu)
         return e, h
 
     def pml_h(e: field_t, h: field_t) -> Tuple[field_t, field_t]:
+        dEv = (numpy.roll(e[v], -1, axis=direction)[region] - e[v][region])
+        dEu = (numpy.roll(e[u], -1, axis=direction)[region] - e[u][region])
         psi_h[0] *= p0h
-        psi_h[0] += p1h * (numpy.roll(e[v], -1, axis=direction)[region] - e[v][region])
+        psi_h[0] += p1h * dEv * p2h
         psi_h[1] *= p0h
-        psi_h[1] += p1h * (numpy.roll(e[u], -1, axis=direction)[region] - e[u][region])
-        h[u][region] -= se * dt * psi_h[0]
-        h[v][region] += se * dt * psi_h[1]
+        psi_h[1] += p1h * dEu * p2h
+        h[u][region] -= se * dt * (psi_h[0] + (p2h - 1) * dEv)
+        h[v][region] += se * dt * (psi_h[1] + (p2h - 1) * dEu)
         return e, h
 
     return pml_e, pml_h, fields
 
 
 def poynting(e, h):
-    s = [numpy.roll(e[1], -1, axis=0) * h[2] - numpy.roll(e[2], -1, axis=0) * h[1],
+    s = (numpy.roll(e[1], -1, axis=0) * h[2] - numpy.roll(e[2], -1, axis=0) * h[1],
          numpy.roll(e[2], -1, axis=1) * h[0] - numpy.roll(e[0], -1, axis=1) * h[2],
-         numpy.roll(e[0], -1, axis=2) * h[1] - numpy.roll(e[1], -1, axis=2) * h[0]]
+         numpy.roll(e[0], -1, axis=2) * h[1] - numpy.roll(e[1], -1, axis=2) * h[0])
     return numpy.array(s)
 
 
-def div_poyting(dt, dxes, e, h):
-    s = poynting(e, h)
-    ds = (s[0] - numpy.roll(s[0], 1, axis=0) +
-          s[1] - numpy.roll(s[1], 1, axis=1) +
-          s[2] - numpy.roll(s[2], 1, axis=2))
+def poynting_divergence(dt, dxes, s=None, *, e=None, h=None): # TODO dxes
+    if s is None:
+        s = poynting(e, h)
+    ds = ((s[0] - numpy.roll(s[0], 1, axis=0)) / numpy.sqrt(dxes[0][0] * dxes[1][0])[:, None, None] +
+          (s[1] - numpy.roll(s[1], 1, axis=1)) / numpy.sqrt(dxes[0][1] * dxes[1][1])[None, :, None] +
+          (s[2] - numpy.roll(s[2], 1, axis=2)) / numpy.sqrt(dxes[0][2] * dxes[1][2])[None, None, :] )
     return ds
+
+
+def energy_hstep(e0, h1, e2, epsilon=None, mu=None, dxes=None):
+    u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes)
+    return u
+
+
+def energy_estep(h0, e1, h2, epsilon=None, mu=None, dxes=None):
+    u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes)
+    return u
+
+
+def delta_energy_h2e(dt, e0, h1, e2, h3, epsilon=None, mu=None, dxes=None):
+    """
+    This is just from (e2 * e2 + h3 * h1) - (h1 * h1 + e0 * e2)
+    """
+    de = e2 * (e2 - e0) / dt
+    dh = h1 * (h3 - h1) / dt
+    du = dt * dxmul(de, dh, epsilon, mu, dxes)
+    return du
+
+
+def delta_energy_e2h(dt, h0, e1, h2, e3, epsilon=None, mu=None, dxes=None):
+    """
+    This is just from (h2 * h2 + e3 * e1) - (e1 * e1 + h0 * h2)
+    """
+    de = e1 * (e3 - e1) / dt
+    dh = h2 * (h2 - h0) / dt
+    du = dxmul(de, dh, epsilon, mu, dxes)
+    return du
+
+
+def delta_energy_j(j0, e1, dxes=None):
+    if dxes is None:
+        dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
+
+    du = ((j0 * e1).sum(axis=0) *
+          dxes[0][0][:, None, None] *
+          dxes[0][1][None, :, None] *
+          dxes[0][2][None, None, :])
+    return du
+
+
+def dxmul(ee, hh, epsilon=None, mu=None, dxes=None):
+    if epsilon is None:
+        epsilon = 1
+    if mu is None:
+        mu = 1
+    if dxes is None:
+        dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
+
+    result = ((ee * epsilon).sum(axis=0) *
+              dxes[0][0][:, None, None] *
+              dxes[0][1][None, :, None] *
+              dxes[0][2][None, None, :] +
+              (hh * mu).sum(axis=0) *
+              dxes[1][0][:, None, None] *
+              dxes[1][1][None, :, None] *
+              dxes[1][2][None, None, :])
+    return result
+
+
+

fdfd_tools/test_fdtd.py

@@ -0,0 +1,68 @@
+import unittest
+import numpy
+
+from fdfd_tools import fdtd
+
+class TestBasic2D(unittest.TestCase):
+    def setUp(self):
+        shape = [3, 5, 5, 1]
+        dt = 0.5
+        epsilon = numpy.ones(shape, dtype=float)
+
+        src_mask = numpy.zeros_like(epsilon, dtype=bool)
+        src_mask[1, 2, 2, 0] = True
+
+        e = numpy.zeros_like(epsilon)
+        h = numpy.zeros_like(epsilon)
+        e[src_mask] = 32
+        es = [e]
+        hs = [h]
+
+        eh2h = fdtd.maxwell_h(dt=dt)
+        eh2e = fdtd.maxwell_e(dt=dt)
+        for _ in range(9):
+            e = e.copy()
+            h = h.copy()
+            eh2h(e, h)
+            eh2e(e, h, epsilon)
+            es.append(e)
+            hs.append(h)
+
+        self.es = es
+        self.hs = hs
+        self.dt = dt
+        self.epsilon = epsilon
+        self.src_mask = src_mask
+
+    def test_initial_fields(self):
+        # Make sure initial fields didn't change
+        e0 = self.es[0]
+        h0 = self.hs[0]
+        self.assertEqual(e0[1, 2, 2, 0], 32)
+
+        self.assertFalse(e0[~self.src_mask].any())
+        self.assertFalse(h0.any())
+
+
+    def test_initial_energy(self):
+        e0 = self.es[0]
+        h0 = self.hs[0]
+        h1 = self.hs[1]
+        mask = self.src_mask[1]
+
+        # Make sure initial energy and E dot J are correct
+        energy0 = fdtd.energy_estep(h0=h0, e1=e0, h2=self.hs[1])
+        e_dot_j_0 = fdtd.delta_energy_j(j0=e0 - 0, e1=e0)
+        self.assertEqual(energy0[mask], 32 * 32)
+        self.assertFalse(energy0[~mask].any())
+        self.assertEqual(e_dot_j_0[mask], 32 * 32)
+        self.assertFalse(e_dot_j_0[~mask].any())
+
+
+    def test_energy_conservation(self):
+        for ii in range(1, 8):
+            with self.subTest(i=ii):
+                u_estep = fdtd.energy_estep(h0=self.hs[ii], e1=self.es[ii], h2=self.hs[ii + 1])
+                u_hstep = fdtd.energy_hstep(e0=self.es[ii-1], h1=self.hs[ii], e2=self.es[ii])
+                self.assertTrue(numpy.allclose(u_estep.sum(), 32 * 32))
+                self.assertTrue(numpy.allclose(u_hstep.sum(), 32 * 32))