much faster custom medium, systematic tests

.gitignore

@@ -134,3 +134,6 @@ htmlcov/
 # Specific file and folder exclusions
 .idea
 .vscode
+
+# cProfile output
+*.prof

tests/test_components/test_autograd.py

@@ -3,14 +3,43 @@
 import pytest
 import matplotlib.pylab as plt
 import numpy as np
+import cProfile
+import typing
 
 import autograd as ag
 import autograd.numpy as npa
 
 import tidy3d as td
+import tidy3d.web as web
 
 from ..utils import run_emulated
 
+""" Test configuration """
+
+"""Test modes
+    pipeline: just run with emulated data, make sure gradient is not 0.0
+    adjoint: run pipeline with real data through web API
+    numerical: adjoint with an extra numerical derivative test after
+    speed: pipeline with cProfile to analyze performance
+"""
+
+TEST_MODES = ("pipeline", "adjoint", "numerical", "speed")
+TEST_MODE = TEST_MODES[-1]
+
+# number of elements in the parameters / input to the objective function
+N_PARAMS = 10
+
+# default starting args
+np.random.seed(1)
+params0 = np.random.random(N_PARAMS)
+
+# whether to plot the simulation within the objective function
+PLOT_SIM = False
+
+# whether to include a call to `objective(params)` in addition to gradient
+CALL_OBJECTIVE = False
+
+""" simulation configuration """
 
 WVL = 1.0
 FREQ0 = td.C_0 / WVL
@@ -26,13 +55,15 @@
 LX = 4 * WVL if IS_3D else 0.0
 PML_X = True if IS_3D else False
 
-PLOT_SIM = False
 
-DA_SHAPE_X = 12 if IS_3D else 1
-DA_SHAPE = (DA_SHAPE_X, 11, 10)
+# shape of the custom medium
+DA_SHAPE_X = 1 if IS_3D else 1
+DA_SHAPE = (DA_SHAPE_X, 500, 500)
 
-NUM_VERTICES = 42
+# number of vertices in the polyslab
+NUM_VERTICES = 12
 
+# sim that we add traced structures and monitors to
 SIM_BASE = td.Simulation(
     size=(LX, 3, LZ),
     run_time=1e-12,
@@ -74,44 +105,49 @@
 @pytest.fixture
 def use_emulated_run(monkeypatch):
     """If this fixture is used, the `tests.utils.run_emulated` function is used for simulation."""
-    import tidy3d.web.api.webapi as webapi
-
-    monkeypatch.setattr(webapi, "run", run_emulated)
-    _run_was_emulated[0] = True
 
+    if TEST_MODE in ("pipeline", "speed"):
+        import tidy3d.web.api.webapi as webapi
 
-# default starting args
-eps0 = 2.0
-center0 = (0.0, 0.0, 0.0)
-size_y0 = 1.0
-eps_arr0 = np.random.random(DA_SHAPE) + 1.0
-radius0 = 0.8
+        monkeypatch.setattr(webapi, "run", run_emulated)
+        _run_was_emulated[0] = True
 
-args0 = (eps0, center0, size_y0, eps_arr0, radius0)
-argnum = tuple(range(len(args0)))
 
-
-def make_structures(eps, center, size_y, eps_arr, radius) -> dict[str, td.Structure]:
+def make_structures(params: npa.ndarray) -> dict[str, td.Structure]:
     """Make a dictionary of the structures given the parameters."""
 
+    params_average = npa.mean(params)
+
+    # static components
     box = td.Box(center=(0, 0, 0), size=(1, 1, 1))
     med = td.Medium(permittivity=2.0)
 
+    # Structure with variable .medium
+    eps = 1 + np.random.random(N_PARAMS) @ params
+    conductivity = eps / 10.0
     medium = td.Structure(
         geometry=box,
-        medium=td.Medium(permittivity=eps, conductivity=eps / 10.0),
+        medium=td.Medium(permittivity=eps, conductivity=conductivity),
     )
 
+    # Structure with variable Box.center
+    center = (np.random.random((3, N_PARAMS)) - 0.5) @ params
+    x0, y0, z0 = center
     center_list = td.Structure(
-        geometry=td.Box(center=center, size=(1, 1, 1)),
+        geometry=td.Box(center=(x0, y0, z0), size=(1, 1, 1)),
         medium=med,
     )
 
+    # Structure with variable Box.center
+    size_y = 1 + np.random.random(N_PARAMS) @ params
     size_element = td.Structure(
         geometry=td.Box(center=(0, 0, 0), size=(1, size_y, 1)),
         medium=med,
     )
 
+    # custom medium with variable permittivity data
+    len_arr = np.prod(DA_SHAPE)
+    eps_arr = 1.0 + npa.dot(np.random.random((len_arr, N_PARAMS)), params).reshape(DA_SHAPE)
     nx, ny, nz = eps_arr.shape
     custom_med = td.Structure(
         geometry=box,
@@ -127,11 +163,12 @@ def make_structures(eps, center, size_y, eps_arr, radius) -> dict[str, td.Struct
         ),
     )
 
+    # Polyslab with variable radius about origin
+    radii = 1 + 0.1 * npa.dot(np.random.random((NUM_VERTICES, N_PARAMS)), params)
     phis = 2 * npa.pi * npa.linspace(0, 1, NUM_VERTICES + 1)[:NUM_VERTICES]
-    xs = radius * npa.cos(phis)
-    ys = radius * npa.sin(phis)
+    xs = radii * npa.cos(phis)
+    ys = radii * npa.sin(phis)
     vertices = npa.stack((xs, ys), axis=-1)
-
     polyslab = td.Structure(
         geometry=td.PolySlab(
             vertices=vertices,
@@ -150,36 +187,38 @@ def make_structures(eps, center, size_y, eps_arr, radius) -> dict[str, td.Struct
     )
 
 
-def make_monitors() -> dict[str, td.Monitor]:
+def make_monitors() -> dict[str, tuple[td.Monitor, typing.Callable[[td.SimulationData], float]]]:
     """Make a dictionary of all the possible monitors in the simulation."""
 
-    mode = td.ModeMonitor(
+    mode_mnt = td.ModeMonitor(
         size=(2, 2, 0),
         center=(0, 0, LZ / 2 - WVL),
         mode_spec=td.ModeSpec(),
         freqs=[FREQ0],
         name="mode",
     )
 
-    mode_pp = lambda mnt_data: npa.sum(abs(mnt_data.amps.values) ** 2)
+    mode_postprocess_fn = lambda mnt_data: npa.sum(abs(mnt_data.amps.values) ** 2)
 
-    diff = td.DiffractionMonitor(
+    diff_mnt = td.DiffractionMonitor(
         size=(td.inf, td.inf, 0),
         center=(0, 0, -LZ / 2 + WVL),
         freqs=[FREQ0],
         normal_dir="+",
         name="diff",
     )
 
-    diff_pp = lambda mnt_data: npa.sum(abs(mnt_data.amps.values) ** 2)
+    diff_postprocess_fn = lambda mnt_data: npa.sum(abs(mnt_data.amps.values) ** 2)
 
     return dict(
-        mode=(mode, mode_pp),
-        diff=(diff, diff_pp),
+        mode=(mode_mnt, mode_postprocess_fn),
+        diff=(diff_mnt, diff_postprocess_fn),
     )
 
 
 def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None:
+    """Plot the simulation."""
+
     plot_fn = sim.plot_eps if plot_eps else sim.plot
 
     f, (ax1, ax2, ax3) = plt.subplots(1, 3, tight_layout=True)
@@ -203,7 +242,9 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None:
     args.append((ALL_KEY, m))
 
 # or just set args manually to test certain things
-args = [("polyslab", "mode")]
+args = [("custom_med", "mode")]
+# args = [(ALL_KEY, "mode")]
+# args = [("medium", "mode")]
 
 
 @pytest.mark.parametrize("structure_key, monitor_key", args)
@@ -260,11 +301,52 @@ def objective(*args):
         value = postprocess(data)
         return value
 
-    val, grad = ag.value_and_grad(objective, argnum=argnum)(*args0)
+    # if speed test, get the profile
+    if TEST_MODE == "speed":
+        with cProfile.Profile() as pr:
+            val, grad = ag.value_and_grad(objective)(params0)
+            pr.print_stats(sort="cumtime")
+            pr.dump_stats("results.prof")
 
-    grad_are_0 = [np.all(np.array(grad_i) == 0) for grad_i in grad]
-
-    assert not npa.all(grad_are_0), "All gradient elements were 0, this combo was un-traced."
+    # otherwise, just test that it ran and the gradients are all non-zero
+    else:
+        if CALL_OBJECTIVE:
+            val = objective(params0)
+        val, grad = ag.value_and_grad(objective)(params0)
+        assert npa.all(grad != 0.0), "some gradients are 0"
+
+    # if 'numerical', we do a numerical gradient check
+    if TEST_MODE == "numerical":
+        delta = 1e-8
+        sims_numerical = {}
+
+        params_num = np.zeros((N_PARAMS, N_PARAMS))
+
+        for i in range(N_PARAMS):
+            for sign, pm_string in zip((-1, 1), "-+"):
+                task_name = f"{i}_{pm_string}"
+                params_i = np.copy(params0)
+                params_i[i] += sign * delta
+                params_num[i] = params_i.copy()
+                sim_i = make_sim(params_i)
+                sims_numerical[task_name] = sim_i
+
+        datas = web.run_async(sims_numerical, path_dir="data")
+
+        grad_num = np.zeros_like(grad)
+        objectives_num = np.zeros((len(params0), 2))
+        for i in range(N_PARAMS):
+            for sign, pm_string in zip((-1, 1), "-+"):
+                task_name = f"{i}_{pm_string}"
+                sim_data_i = datas[task_name]
+                obj_i = postprocess(sim_data_i)
+                objectives_num[i, (sign + 1) // 2] = obj_i
+                grad_num[i] += sign * obj_i / 2 / delta
+
+        print("adjoint: ", grad)
+        print("numerical: ", grad_num)
+
+        assert np.allclose(grad, grad_num), "gradients dont match"
 
     # for logging output
     td.config.logging_level = "WARNING"

tidy3d/components/base.py

@@ -940,7 +940,7 @@ def handle_value(x: Any, path: tuple[str, ...]) -> None:
                 if isinstance(x, Box):
                     field_mapping[path] = x
                 elif isinstance(x, DataArray):
-                    # NOTE: for some reason if I dont do `npa.array(x.tolist()) it wont trace
+                    # NOTE: here be dragons, if you dont copy it this way (tolist()) it will break
                     field_mapping[path] = npa.array(x.values.tolist())
 
             # for sequences, add (i,) to the path and handle each value
@@ -956,7 +956,7 @@ def handle_value(x: Any, path: tuple[str, ...]) -> None:
                     handle_value(val, path=sub_paths)
 
         # recursively parse the dictionary of this object
-        self_dict = self.dict()  # TODO: need copy here?
+        self_dict = self.dict()
         handle_value(self_dict, path=())
 
         # convert the resulting field_mapping to an autograd-traced dictionary
@@ -981,21 +981,15 @@ def insert_value(value, path: tuple[str, ...], sub_dict: dict):
 
             # only one element in path => leaf node. insert into the sub dict and don't recurse
             if len_sub_path == 0:
-                if isinstance(sub_dict[key], DataArray):
-                    # ensure you copy() because DataArrays are mutable and setting would affect self
-                    sub_dict[key] = sub_dict[key].copy()
+                sub_element = sub_dict[key]
+                if isinstance(sub_element, DataArray):
+                    # NOTE: if you don't copy, you'll mutate data in self and bad stuff will occur
+                    sub_dict[key] = sub_element.copy()
                     sub_dict[key].values = value
                 else:
                     sub_dict[key] = value
                 return
 
-            # one integer sub-path => index into tuple leaf node. insert value and don't recurse
-            # if len(sub_path) == 1:
-            #     (sub_key,) = sub_path
-            #     if isinstance(sub_key, int):
-            #         sub_dict[key][sub_key] = value
-            #         return
-
             # if 1 or more more elements in the path, and they aren't a tuple index (above), recurse
             sub_dict = sub_dict[key]
             insert_value(value=value, path=sub_path, sub_dict=sub_dict)

tidy3d/components/geometry/polyslab.py

@@ -1411,6 +1411,8 @@ def compute_derivative_vertices(
         E_der_slab = self.project_in_basis(E_der_at_edges, basis_vector=basis_vectors["slab"])
 
         # approximate permittivity in and out
+
+        # TODO: not a good approximation at all. better to use `medium.eps_model`
         eps_in = complex(
             np.sum(np.mean([np.mean(val) for _, val in eps_structure.field_components.items()]))
         )

tidy3d/components/medium.py

@@ -2013,14 +2013,10 @@ def _medium(self):
         either `CustomIsotropicMedium` or `CustomAnisotropicMedium`.
         """
 
-        _self = self.to_static()
-
-        self_dict = _self.dict(exclude={"type", "eps_dataset"})
+        self_dict = self.dict(exclude={"type", "eps_dataset"})
         # isotropic
-        if _self.eps_dataset is None:
-            self_dict.update(
-                {"permittivity": _self.permittivity, "conductivity": _self.conductivity}
-            )
+        if self.eps_dataset is None:
+            self_dict.update({"permittivity": self.permittivity, "conductivity": self.conductivity})
             return CustomIsotropicMedium.parse_obj(self_dict)
 
         def get_eps_sigma(eps_complex: SpatialDataArray, freq: float) -> tuple: