Support for Nonuniform grid spec in grid_size

tests/test_grid.py

@@ -20,7 +20,6 @@ def test_field_grid():
     c = Coords(x=x, y=y, z=z)
     f = FieldGrid(x=c, y=c, z=c)
 
-
 def test_grid():
 
     boundaries_x = np.arange(-1, 2, 1)
@@ -41,6 +40,87 @@ def test_grid():
     assert np.all(g.yee.E.x.z == np.array([-3, -2, -1, 0, 1, 2]))
 
 
+def test_sim_nonuniform_small():
+    # tests when the nonuniform grid doesnt cover the simulation size
+
+    size_x = 18
+    num_layers_pml_x = 2
+    grid_size_x = [2, 1, 3]
+    sim = td.Simulation(
+        center=(1, 0, 0),
+        size=(size_x, 4, 4), 
+        grid_size=(grid_size_x, 1, 1),
+        pml_layers=[td.PML(num_layers=num_layers_pml_x), None, None]
+    )
+
+    bound_coords = sim.grid.boundaries.x
+    dls = np.diff(bound_coords)
+
+    dl_min = grid_size_x[0]
+    dl_max = grid_size_x[-1]
+
+    # checks the bounds were adjusted correctly
+    # (smaller than sim size as is, but larger than sim size with one dl added on each edge)
+    assert np.sum(dls) <= size_x + num_layers_pml_x*dl_min + num_layers_pml_x*dl_max
+    assert np.sum(dls)+dl_min+dl_max >= size_x + num_layers_pml_x*dl_min + num_layers_pml_x*dl_max
+
+    # tests that PMLs were added correctly
+    for i in range(num_layers_pml_x):
+        assert np.diff(bound_coords[i:i+2]) == dl_min
+        assert np.diff(bound_coords[-2-i:len(bound_coords)-i]) == dl_max
+
+    # tests that all the grid sizes are in there
+    for size in grid_size_x:
+        assert size in dls
+
+    # tests that nothing but the grid sizes are in there
+    for dl in dls:
+        assert dl in grid_size_x
+
+    # tests that it gives exactly what we expect
+    assert np.all(bound_coords == np.array([-12,-10,-8,-6,-4,-2,0,1,4,7,10,13,16]))
+
+def test_sim_nonuniform_large():
+    # tests when the nonuniform grid extends beyond the simulation size
+
+    size_x = 18
+    num_layers_pml_x = 2
+    grid_size_x = [2, 3, 4, 1, 2, 1, 3, 1, 2, 3, 4]
+    sim = td.Simulation(
+        center=(1, 0, 0),
+        size=(size_x, 4, 4), 
+        grid_size=(grid_size_x, 1, 1),
+        pml_layers=[td.PML(num_layers=num_layers_pml_x), None, None]
+    )
+
+    bound_coords = sim.grid.boundaries.x
+    dls = np.diff(bound_coords)
+
+    dl_min = grid_size_x[0]
+    dl_max = grid_size_x[-1]
+
+    # checks the bounds were adjusted correctly
+    # (smaller than sim size as is, but larger than sim size with one dl added on each edge)
+    assert np.sum(dls) <= size_x + num_layers_pml_x*dl_min + num_layers_pml_x*dl_max
+    assert np.sum(dls)+dl_min+dl_max >= size_x + num_layers_pml_x*dl_min + num_layers_pml_x*dl_max
+
+    # tests that PMLs were added correctly
+    for i in range(num_layers_pml_x):
+        assert np.diff(bound_coords[i:i+2]) == grid_size_x[0]
+        assert np.diff(bound_coords[-2-i:len(bound_coords)-i]) == grid_size_x[-1]
+
+    # tests that all the grid sizes are in there
+    for size in grid_size_x:
+        assert size in dls
+
+    # tests that nothing but the grid sizes are in there
+    for dl in dls:
+        assert dl in grid_size_x
+
+    # tests that it gives exactly what we expect
+    # assert np.all(bound_coords == np.array([-12,-10,-8,-6,-4,-2,0,1,4,7,10,13,16]))
+
+
 def test_sim_grid():
 
     sim = td.Simulation(

tidy3d/components/simulation.py

@@ -11,7 +11,7 @@
 
 from .validators import assert_unique_names, assert_objects_in_sim_bounds, set_names
 from .geometry import Box
-from .types import Symmetry, Ax, Shapely, FreqBound
+from .types import Symmetry, Ax, Shapely, FreqBound, GridSize
 from .grid import Coords1D, Grid, Coords
 from .medium import Medium, MediumType, AbstractMedium
 from .structure import Structure
@@ -42,8 +42,10 @@ class Simulation(Box):  # pylint:disable=too-many-public-methods
     size : Tuple[float, float, float]
         (microns) Size of simulation domain in x, y, and z.
         Each element must be non-negative.
-    grid_size : Tuple[float, float, float]
-        (microns) Grid size along x, y, and z.
+    grid_size : Tuple[Union[float, List[float]], Union[float, List[float]], Union[float, List[float]]]
+        (microns) If components are float, uniform grid size along x, y, and z.
+        If components are array like, defines an array of nonuniform grid sizes centered at ``simulation.center``.
+        Note: if supplied sizes do not cover ``simulation.size``, the first and last sizes are repeated to cover size.
         Each element must be non-negative.
     run_time : float = 0.0
         Total electromagnetic evolution time in seconds.
@@ -140,7 +142,7 @@ class Simulation(Box):  # pylint:disable=too-many-public-methods
     """
     # pylint:enable=line-too-long
 
-    grid_size: Tuple[pydantic.PositiveFloat, pydantic.PositiveFloat, pydantic.PositiveFloat]
+    grid_size: Tuple[GridSize, GridSize, GridSize]
     medium: MediumType = Medium()
     run_time: pydantic.NonNegativeFloat = 0.0
     structures: List[Structure] = []
@@ -800,6 +802,52 @@ def tmesh(self) -> Coords1D:
         dt = self.dt
         return np.arange(0.0, self.run_time + dt, dt)
 
+    def _make_bound_coords_uniform(self, dl, center, size, num_layers):
+        """creates coordinate boundaries with uniform mesh (dl is float)"""
+
+        num_cells = int(np.floor(size / dl))
+
+        # Make sure there's at least one cell
+        num_cells = max(num_cells, 1)
+
+        # snap to grid, recenter, and add PML
+        size_snapped = dl * num_cells
+        bound_coords = center + np.linspace(-size_snapped / 2, size_snapped / 2, num_cells + 1)
+        bound_coords = self._add_pml_to_bounds(num_layers, bound_coords)
+        return bound_coords
+
+    @staticmethod
+    def _make_bound_coords_nonuniform(dl, center, size, num_layers):
+        """creates coordinate boundaries with non-uniform mesh (dl is arraylike)"""
+
+        # get bounding coordinates
+        dl = np.array(dl)
+        bound_coords = np.array([np.sum(dl[:i]) for i in range(len(dl) + 1)])
+
+        # shift coords to center at center of simulation along dimension
+        bound_coords = bound_coords - np.sum(dl)/2 + center
+
+        # chop off any coords outside of simulation bounds
+        bound_min = center - size/2
+        bound_max = center + size/2
+        bound_coords = bound_coords[bound_coords <= bound_max]
+        bound_coords = bound_coords[bound_coords >= bound_min]
+
+        # if not extending to simulation bounds, repeat beginning and end
+        dl_min = dl[0]
+        dl_max = dl[-1]
+        while bound_coords[0] - dl_min >= bound_min:
+            bound_coords = np.insert(bound_coords, 0, bound_coords[0] - dl_min)
+        while bound_coords[-1] + dl_max <= bound_max:
+            bound_coords = np.append(bound_coords, bound_coords[-1] + dl_max)
+
+        # add PML layers in using dl on edges
+        for _ in range(num_layers[0]):
+            bound_coords = np.insert(bound_coords, 0, bound_coords[0] - dl_min)
+        for _ in range(num_layers[1]):
+            bound_coords = np.append(bound_coords, bound_coords[-1] + dl_max)
+        return bound_coords
+
     @property
     def grid(self) -> Grid:
         """FDTD grid spatial locations and information.
@@ -812,14 +860,10 @@ def grid(self) -> Grid:
         cell_boundary_dict = {}
         zipped_vals = zip("xyz", self.grid_size, self.center, self.size, self.num_pml_layers)
         for key, dl, center, size, num_layers in zipped_vals:
-            num_cells = int(np.floor(size / dl))
-            # Make sure there's at least one cell
-            num_cells = max(num_cells, 1)
-            size_snapped = dl * num_cells
-            # if size_snapped != size:
-            #     log.warning(f"dl = {dl} not commensurate with simulation size = {size}")
-            bound_coords = center + np.linspace(-size_snapped / 2, size_snapped / 2, num_cells + 1)
-            bound_coords = self._add_pml_to_bounds(num_layers, bound_coords)
+            if isinstance(dl, float):
+                bound_coords = self._make_bound_coords_uniform(dl, center, size, num_layers)
+            else:
+                bound_coords = self._make_bound_coords_nonuniform(dl, center, size, num_layers)
             cell_boundary_dict[key] = bound_coords
         boundaries = Coords(**cell_boundary_dict)
         return Grid(boundaries=boundaries)