autograd geometry group

CHANGELOG.md

@@ -8,9 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [2.8.0rc1]
 
 ### Added
+- Support for differentiation with respect to `GeometryGroup.geometries` elements.
 - Users can now export `SimulationData` to MATLAB `.mat` files with the `to_mat_file` method.
 - Introduce RF material library. Users can now export `rf_material_library` from `tidy3d.plugins.microwave`.
 
+### Changed
+
 ### Fixed
 - Bug where boundary layers would be plotted too small in 2D simulations.
 

tests/test_components/test_autograd.py

@@ -243,13 +243,26 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
         medium=med,
     )
 
+    # geometry group
+    geo_group = td.Structure(
+        geometry=td.GeometryGroup(
+            geometries=[
+                medium.geometry,
+                center_list.geometry,
+                size_element.geometry,
+            ],
+        ),
+        medium=td.Medium(permittivity=eps, conductivity=conductivity),
+    )
+
     return dict(
         medium=medium,
         center_list=center_list,
         size_element=size_element,
         custom_med=custom_med,
         custom_med_vec=custom_med_vec,
         polyslab=polyslab,
+        geo_group=geo_group,
     )
 
 
@@ -336,6 +349,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None:
     "custom_med",
     "custom_med_vec",
     "polyslab",
+    "geo_group",
 )
 monitor_keys_ = ("mode", "diff", "field_vol", "field_point")
 
@@ -356,7 +370,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None:
     args = [("polyslab", "mode")]
 
 
-# args = [("custom_med_vec", "mode")]
+# args = [("geo_group", "mode")]
 
 
 def get_functions(structure_key: str, monitor_key: str) -> typing.Callable:
@@ -716,42 +730,62 @@ def f3(x):
 
 
 # @pytest.mark.timeout(18.0)
-# def test_many_structures_timeout():
-#     """Test that a metalens-like simulation with many structures can be initialized fast enough."""
-
-#     with cProfile.Profile() as pr:
-#         import time
-
-#         t = time.time()
-
-#         Nx, Ny = 200, 200
-#         sim_size = [Nx, Ny, 5]
-
-#         geoms = []
-#         for ix in range(Nx):
-#             for iy in range(Ny):
-#                 verts = ((ix, iy), (ix + 0.5, iy), (ix + 0.5, iy + 0.5), (ix, iy + 0.5))
-#                 geom = td.PolySlab(slab_bounds=(0, 1), vertices=verts)
-#                 geoms.append(geom)
-
-#         metalens = td.Structure(
-#             geometry=td.GeometryGroup(geometries=geoms),
-#             medium=td.material_library["Si3N4"]["Horiba"],
-#         )
-
-#         src = td.PlaneWave(
-#             source_time=td.GaussianPulse(freq0=2.5e14, fwidth=1e13),
-#             center=(0, 0, -1),
-#             size=(td.inf, td.inf, 0),
-#             direction="+",
-#         )
-
-#         sim = td.Simulation(size=sim_size, structures=[metalens], sources=[src], run_time=1e-12)
-
-#         t2 = time.time() - t
-#         pr.print_stats(sort="cumtime")
-#         pr.dump_stats("sim_test.prof")
-#         print(f"structures took {t2} seconds")
+def _test_many_structures():
+    """Test that a metalens-like simulation with many structures can be initialized fast enough."""
+
+    with cProfile.Profile() as pr:
+        import time
+
+        t = time.time()
+
+        N_length = 200
+        Nx, Ny = N_length, N_length
+        sim_size = [Nx, Ny, 5]
+
+        def f(x):
+            monitor, postprocess = make_monitors()["field_point"]
+            monitor = monitor.updated_copy(center=(0, 0, 0))
+
+            geoms = []
+            for ix in range(Nx):
+                for iy in range(Ny):
+                    ix = ix + x
+                    iy = iy + x
+                    verts = ((ix, iy), (ix + 0.5, iy), (ix + 0.5, iy + 0.5), (ix, iy + 0.5))
+                    geom = td.PolySlab(slab_bounds=(0, 1), vertices=verts)
+                    geoms.append(geom)
+
+            metalens = td.Structure(
+                geometry=td.GeometryGroup(geometries=geoms),
+                medium=td.material_library["Si3N4"]["Horiba"],
+            )
+
+            src = td.PlaneWave(
+                source_time=td.GaussianPulse(freq0=2.5e14, fwidth=1e13),
+                center=(0, 0, -1),
+                size=(td.inf, td.inf, 0),
+                direction="+",
+            )
+
+            sim = td.Simulation(
+                size=sim_size,
+                structures=[metalens],
+                sources=[src],
+                monitors=[monitor],
+                run_time=1e-12,
+            )
+
+            data = run_emulated(sim, task_name="test")
+            return postprocess(data, data[monitor.name])
+
+        x0 = 0.0
+        ag.grad(f)(x0)
+
+        t2 = time.time() - t
+        pr.print_stats(sort="cumtime")
+        pr.dump_stats("sim_test.prof")
+        print(f"structures took {t2} seconds")
+
 
 """ times (tyler's system)
 * original : 35 sec

tidy3d/components/autograd/__init__.py

@@ -0,0 +1,23 @@
+from .types import (
+    AutogradFieldMap,
+    AutogradTraced,
+    TracedCoordinate,
+    TracedFloat,
+    TracedSize,
+    TracedSize1D,
+    TracedVertices,
+)
+from .utils import get_static
+
+__all__ = [
+    "TracedFloat",
+    "TracedSize1D",
+    "TracedSize",
+    "TracedCoordinate",
+    "TracedVertices",
+    "AutogradTraced",
+    "AutogradFieldMap",
+    "get_static",
+    "integrate_within_bounds",
+    "DerivativeInfo",
+]

tidy3d/components/autograd/derivative_utils.py

@@ -0,0 +1,121 @@
+# utilities for autograd derivative passing
+from __future__ import annotations
+
+import numpy as np
+import pydantic.v1 as pd
+import xarray as xr
+
+from ..base import Tidy3dBaseModel
+from ..data.data_array import ScalarFieldDataArray
+from ..types import Bound, tidycomplex
+from .types import PathType
+from .utils import get_static
+
+# we do this because importing these creates circular imports
+FieldData = dict[str, ScalarFieldDataArray]
+PermittivityData = dict[str, ScalarFieldDataArray]
+
+
+class DerivativeInfo(Tidy3dBaseModel):
+    """Stores derivative information passed to the ``.compute_derivatives`` methods."""
+
+    paths: list[PathType] = pd.Field(
+        ...,
+        title="Paths to Traced Fields",
+        description="List of paths to the traced fields that need derivatives calculated.",
+    )
+
+    E_der_map: FieldData = pd.Field(
+        ...,
+        title="Electric Field Gradient Map",
+        description='Dataset where the field components ``("Ex", "Ey", "Ez")`` store the '
+        "multiplication of the forward and adjoint electric fields. The tangential components "
+        "of this dataset is used when computing adjoint gradients for shifting boundaries. "
+        "All components are used when computing volume-based gradients.",
+    )
+
+    D_der_map: FieldData = pd.Field(
+        ...,
+        title="Displacement Field Gradient Map",
+        description='Dataset where the field components ``("Ex", "Ey", "Ez")`` store the '
+        "multiplication of the forward and adjoint displacement fields. The normal component "
+        "of this dataset is used when computing adjoint gradients for shifting boundaries.",
+    )
+
+    eps_data: PermittivityData = pd.Field(
+        ...,
+        title="Permittivity Dataset",
+        description="Dataset of relative permittivity values along all three dimensions. "
+        "Used for automatically computing permittivity inside or outside of a simple geometry.",
+    )
+
+    eps_in: tidycomplex = pd.Field(
+        title="Permittivity Inside",
+        description="Permittivity inside of the ``Structure``. "
+        "Typically computed from ``Structure.medium.eps_model``."
+        "Used when it can not be computed from ``eps_data`` or when ``eps_approx==True``.",
+    )
+
+    eps_out: tidycomplex = pd.Field(
+        ...,
+        title="Permittivity Outside",
+        description="Permittivity outside of the ``Structure``. "
+        "Typically computed from ``Simulation.medium.eps_model``."
+        "Used when it can not be computed from ``eps_data`` or when ``eps_approx==True``.",
+    )
+
+    bounds: Bound = pd.Field(
+        ...,
+        title="Geometry Bounds",
+        description="Bounds corresponding to the structure, used in ``Medium`` calculations.",
+    )
+
+    eps_approx: bool = pd.Field(
+        False,
+        title="Use Permittivity Approximation",
+        description="If ``True``, approximates outside permittivity using ``Simulation.medium``"
+        "and the inside permittivity using ``Structure.medium``. "
+        "Only set ``True`` for ``GeometryGroup`` handling where it is difficult to automatically "
+        "evaluate the inside and outside relative permittivity for each geometry.",
+    )
+
+    def updated_paths(self, paths: list[PathType]) -> DerivativeInfo:
+        """Update this ``DerivativeInfo`` with new set of paths."""
+        return self.updated_copy(paths=paths)
+
+
+# TODO: could we move this into a DataArray method?
+def integrate_within_bounds(arr: xr.DataArray, dims: list[str], bounds: Bound) -> xr.DataArray:
+    """integrate a data array within bounds, assumes bounds are [2, N] for N dims."""
+
+    _arr = arr.copy()
+
+    # order bounds with dimension first (N, 2)
+    bounds = np.array(bounds).T
+
+    all_coords = {}
+
+    # loop over all dimensions
+    for dim, (bmin, bmax) in zip(dims, bounds):
+        bmin = get_static(bmin)
+        bmax = get_static(bmax)
+
+        coord_values = _arr.coords[dim].values
+
+        # reset all coordinates outside of bounds to the bounds, so that dL = 0 in integral
+        coord_values[coord_values < bmin] = bmin
+        coord_values[coord_values > bmax] = bmax
+
+        all_coords[dim] = coord_values
+
+    _arr = _arr.assign_coords(**all_coords)
+
+    # uses trapezoidal rule
+    # https://docs.xarray.dev/en/stable/generated/xarray.DataArray.integrate.html
+    return _arr.integrate(coord=dims)
+
+
+__all__ = [
+    "integrate_within_bounds",
+    "DerivativeInfo",
+]