diff --git a/fdtd.py b/fdtd.py
index e586ee3..8e30fd0 100644
--- a/fdtd.py
+++ b/fdtd.py
@@ -75,7 +75,7 @@ def perturbed_l3(a: float, radius: float, **kwargs) -> Pattern:
 def main():
     max_t = 8000            # number of timesteps
 
-    dx = 40                 # discretization (nm/cell)
+    dx = 15                 # discretization (nm/cell)
     pml_thickness = 8       # (number of cells)
 
     wl = 1550               # Excitation wavelength and fwhm
@@ -109,13 +109,16 @@ def main():
                    eps=n_slab**2)
     mask = perturbed_l3(a, r)
 
-    grid.draw_polygons(surface_normal=gridlock.Direction.z,
+    for i in range(5000):
+        print(i)
+        grid.draw_polygons(surface_normal=gridlock.Direction.z,
                        center=[0, 0, 0],
                        thickness=2 * th,
                        eps=n_air**2,
                        polygons=mask.as_polygons())
+    return
 
-    print(grid.shape)
+    print(grid.shape, numpy.prod(grid.shape))
     # #### Create the simulation grid ####
     sim = Simulation(grid.grids)
 
diff --git a/fdtd/base.py b/fdtd/base.py
index 6285fb9..3717a65 100644
--- a/fdtd/base.py
+++ b/fdtd/base.py
@@ -30,7 +30,8 @@ const int diy = sz;
 const int diz = 1;
 """
 
-# Given a linear index i and shape sx, sy, sz, defines x, y, and z as the 3D indices of the current element (i).
+# Given a linear index i and shape sx, sy, sz, defines x, y, and z
+#  as the 3D indices of the current element (i).
 xyz_source = """
 // Convert linear index to field index (xyz)
 const int x = i / (sz * sy);