diff --git a/CHANGELOG.md b/CHANGELOG.md index 61652088ca..67b51bdbcb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- Gradient computation for rotated boxes in Transformed. +- Gradient computation for rotated angle of Transformed box. +- Gradient computation for chained transformation of box. + ## [2.8.4] - 2025-05-15 ### Added diff --git a/tests/test_components/test_autograd.py b/tests/test_components/test_autograd.py index a4d582ced5..8b9ffeaae3 100644 --- a/tests/test_components/test_autograd.py +++ b/tests/test_components/test_autograd.py @@ -2208,3 +2208,123 @@ def objective(x): with pytest.raises(ValueError): g = ag.grad(objective)(1.0) + + +def make_sim_rotation(center: tuple, size: tuple, angle: float, axis: int): + wavelength = 1.5 + L = 10 * wavelength + freq0 = td.C_0 / wavelength + buffer = 1.0 * wavelength + + # Source + src = td.PointDipole( + center=(-L / 2 + buffer, 0, 0), + source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10.0), + polarization="Ez", + ) + # Monitor + mnt = td.FieldMonitor( + center=( + +L / 2 - buffer, + 0.5 * buffer, + 0.5 * buffer, + ), + size=(0.0, 0.0, 0.0), + freqs=[freq0], + name="point", + ) + # The box geometry + base_box = td.Box(center=center, size=size) + if angle is not None: + base_box = base_box.rotated(angle, axis) + + scatterer = td.Structure( + geometry=base_box, + medium=td.Medium(permittivity=2.0), + ) + + sim = td.Simulation( + size=(L, L, L), + grid_spec=td.GridSpec.auto(min_steps_per_wvl=50), + structures=[scatterer], + sources=[src], + monitors=[mnt], + run_time=120 / freq0, + ) + return sim + + +def objective_fn(center, size, angle, axis): + sim = make_sim_rotation(center, size, angle, axis) + sim_data = web.run(sim, task_name="emulated_rot_test", local_gradient=True, verbose=False) + return anp.sum(sim_data.get_intensity("point").values) + + +def get_grad(center, size, angle, axis): + def wrapped(c, s): + return objective_fn(c, s, angle, axis) + + val, (grad_c, grad_s) = ag.value_and_grad(wrapped, argnum=(0, 1))(center, size) + return val, grad_c, grad_s + + +@pytest.mark.numerical +@pytest.mark.parametrize( + "angle_deg, axis", + [ + (0.0, 1), + (180.0, 1), + (90.0, 1), + (270.0, 1), + ], +) +def test_box_rotation_gradients(use_emulated_run, angle_deg, axis): + center0 = (0.0, 0.0, 0.0) + size0 = (2.0, 2.0, 2.0) + + angle_rad = np.deg2rad(angle_deg) + val, grad_c, grad_s = get_grad(center0, size0, angle=None, axis=None) + npx, npy, npz = grad_c + sSx, sSy, sSz = grad_s + + assert not np.allclose(grad_c, 0.0), "center gradient is all zero." + assert not np.allclose(grad_s, 0.0), "size gradient is all zero." + + if angle_deg == 180.0: + # rotating 180° about y => (x,z) become negated, y stays same + _, grad_c_ref, grad_s_ref = get_grad(center0, size0, angle_rad, axis) + rSx, rSy, rSz = grad_s_ref + rx, ry, rz = grad_c_ref + + assert np.allclose(grad_c[0], -grad_c_ref[0], atol=1e-6), "center_x sign mismatch" + assert np.allclose(grad_c[1], grad_c_ref[1], atol=1e-6), "center_y mismatch" + assert np.allclose(grad_c[2], -grad_c_ref[2], atol=1e-6), "center_z sign mismatch" + assert np.allclose(grad_s, grad_s_ref, atol=1e-6), "size grads changed unexpectedly" + + elif angle_deg == 90.0: + # rotating 90° about y => new x= old z, new z=- old x, y stays same + _, grad_c_ref, grad_s_ref = get_grad(center0, size0, angle_rad, axis) + rSx, rSy, rSz = grad_s_ref + rx, ry, rz = grad_c_ref + + assert np.allclose(npx, rz, atol=1e-6), "center_x != old center_z" + assert np.allclose(npy, ry, atol=1e-6), "center_y changed unexpectedly" + assert np.allclose(npz, -rx, atol=1e-6), "center_z != - old center_x" + + assert np.allclose(sSx, rSz, atol=1e-6), "size_x != old size_z" + assert np.allclose(sSy, rSy, atol=1e-6), "size_y changed unexpectedly" + assert np.allclose(sSz, rSx, atol=1e-6), "size_z != old size_x" + + elif angle_deg == 270.0: + # rotating 270° about y => new x= - old z, new z= old x, y stays same + _, grad_c_ref, grad_s_ref = get_grad(center0, size0, angle_rad, axis) + rSx, rSy, rSz = grad_s_ref + rx, ry, rz = grad_c_ref + + assert np.allclose(npx, -rz, atol=1e-6), "center_x != - old center_z" + assert np.allclose(npy, ry, atol=1e-6), "center_y changed unexpectedly" + assert np.allclose(npz, rx, atol=1e-6), "center_z != old center_x" + + assert np.allclose(sSx, rSz, atol=1e-6), "size_x != old size_z" + assert np.allclose(sSy, rSy, atol=1e-6), "size_y changed unexpectedly" + assert np.allclose(sSz, rSx, atol=1e-6), "size_z != old size_x" diff --git a/tests/test_components/test_box_chained_derivatives.py b/tests/test_components/test_box_chained_derivatives.py new file mode 100644 index 0000000000..da49699b67 --- /dev/null +++ b/tests/test_components/test_box_chained_derivatives.py @@ -0,0 +1,253 @@ +""" +FD vs AD checks for every affine parameter of a dielectric box: + • centre c = (cx,cy,cz) + • size a = (ax,ay,az) + • rotation θ about each axis + • scale s = (sx,sy,sz) + • translation t = (tx,ty,tz) +""" + +import atexit +import os +from collections import defaultdict + +import autograd +import autograd.numpy as anp +import matplotlib.pyplot as plt +import numpy as np +import pytest +import tidy3d as td +import tidy3d.web as web +from autograd import tuple + +# ───────── switches ─────────────────────────────────────────────────────── +SAVE = False # save raw .npz +PLOT = True # make png plots +OUT = "./fd_ad_all_results" + +# ───────── physical / geometric constants ───────────────────────────────── +λ = 1.5 +f0 = td.C_0 / λ +Lbox = 10 * λ +buffer = 1.0 * λ +Tsim = 120 / f0 + +# baseline shape parameters ( **plain Python lists / tuples** ) +center0 = [0.0, 0.0, 0.0] +size0 = [2.0, 2.0, 2.0] +eps_box = 2.0 + +theta0 = np.pi / 4 # baseline rotation (x-axis) +axis0 = 0 +scale0 = [1.2, 1.3, 0.9] +trans0 = [0.45, 0.20, 0.30] + +# finite-difference steps +Δθ = 0.015 +Δxyz = 0.03 + +AXES = (0, 1, 2) +LAB = {0: "x", 1: "y", 2: "z"} + +plots = defaultdict(list) + + +# ───────── simulation builder ───────────────────────────────────────────── +def make_simulation(center, size, scale, trans, theta, axis): + """Return a Tidy3D Simulation with the requested affine parameters.""" + src = td.PointDipole( + center=(-Lbox / 2 + 0.5 * buffer, 0, 0), + source_time=td.GaussianPulse(freq0=f0, fwidth=f0 / 10), + polarization="Ez", + ) + + mon = td.FieldMonitor( + center=(+Lbox / 2 - 0.5 * buffer, 2.0 * buffer, 2.0 * buffer), + size=(0, 0, 0), + freqs=[f0], + name="m", + ) + + geom = ( + td.Box(center=tuple(center), size=tuple(size)) + .rotated(theta, axis=axis) + .scaled(x=scale[0], y=scale[1], z=scale[2]) + .translated(x=trans[0], y=trans[1], z=trans[2]) + ) + struct = td.Structure(geometry=geom, medium=td.Medium(permittivity=eps_box)) + + return td.Simulation( + size=(Lbox, Lbox, Lbox), + run_time=Tsim, + grid_spec=td.GridSpec.auto(min_steps_per_wvl=50), + sources=[src], + monitors=[mon], + structures=[struct], + ) + + +def objective(c, a, s, t, θ, ax): + sim = make_simulation(c, a, s, t, θ, ax) + data = web.run(sim, task_name="fd_ad_all", verbose=False, local_gradient=True) + return anp.sum(data.get_intensity("m").values) + + +def finite_diff(fun, x0, δ): + return (fun(x0 + δ) - fun(x0 - δ)) / (2 * δ) + + +def _assert_close(fd_val, ad_val, tag, axis=None, extra="", tol=0.35): + """ + Raise AssertionError if |FD-AD|/max(|FD|,1e-12) >= tol. + """ + rel = abs(fd_val - ad_val) / max(abs(fd_val), 1e-12) + if rel >= tol: + ax_lbl = {0: "x", 1: "y", 2: "z"}.get(axis, "") + axis_str = f"_{ax_lbl}" if ax_lbl else "" + raise AssertionError( + f"{tag}{axis_str}{extra}: FD–AD mismatch " + f"(rel diff {rel:.2%}) | FD={fd_val:.4e}, AD={ad_val:.4e}" + ) + + +# 1. ROTATION θ_k (k = 0,1,2) +θ_vals = np.array([0, np.pi / 4, np.pi / 2]) + + +@pytest.mark.numerical +@pytest.mark.parametrize("k", AXES, ids=[f"θ_{LAB[a]}" for a in AXES]) +@pytest.mark.parametrize("θ", θ_vals) +def test_fd_vs_ad_rotation(k, θ): + f = lambda th: objective(center0, size0, scale0, trans0, th, k) + g_ad = autograd.grad(f)(θ) + g_fd = finite_diff(f, θ, Δθ) + + plots[("θ", k, "fd")].append((θ, g_fd)) + plots[("θ", k, "ad")].append((θ, g_ad)) + + assert np.isfinite(g_fd) and np.isfinite(g_ad) + _assert_close(g_fd, g_ad, tag="θ", axis=k) + + +# 2. SCALE s_k +@pytest.mark.numerical +@pytest.mark.parametrize("k", AXES, ids=[f"s_{LAB[a]}" for a in AXES]) +def test_fd_vs_ad_scale(k): + def f(sk): + s = list(scale0) + s[k] = sk + return objective(center0, size0, s, trans0, theta0, axis0) + + g_ad = autograd.grad(f)(scale0[k]) + g_fd = finite_diff(f, scale0[k], Δxyz) + + plots[("s", k, "fd")].append(g_fd) + plots[("s", k, "ad")].append(g_ad) + + assert np.isfinite(g_fd) and np.isfinite(g_ad), "NaN/Inf in FD or AD result" + _assert_close(g_fd, g_ad, tag="scale", axis=k) + + +# 3. TRANSLATION t_k +@pytest.mark.numerical +@pytest.mark.parametrize("k", AXES, ids=[f"t_{LAB[a]}" for a in AXES]) +def test_fd_vs_ad_translation(k): + def f(tk): + t = list(trans0) + t[k] = tk + return objective(center0, size0, scale0, t, theta0, axis0) + + g_ad = autograd.grad(f)(trans0[k]) + g_fd = finite_diff(f, trans0[k], Δxyz) + + plots[("t", k, "fd")].append(g_fd) + plots[("t", k, "ad")].append(g_ad) + + assert np.isfinite(g_fd) and np.isfinite(g_ad), "NaN/Inf in FD or AD result" + _assert_close(g_fd, g_ad, tag="trans", axis=k) + + +# # 4. SIZE a_k +@pytest.mark.numerical +@pytest.mark.parametrize("k", AXES, ids=[f"a_{LAB[a]}" for a in AXES]) +def test_fd_vs_ad_size(k): + def f(ak): + a = list(size0) + a[k] = ak + return objective(center0, a, scale0, trans0, theta0, axis0) + + g_ad = autograd.grad(f)(size0[k]) + g_fd = finite_diff(f, size0[k], Δxyz) + + plots[("a", k, "fd")].append(g_fd) + plots[("a", k, "ad")].append(g_ad) + + assert np.isfinite(g_fd) and np.isfinite(g_ad), "NaN/Inf in FD or AD result" + _assert_close(g_fd, g_ad, tag="size", axis=k) + + +# # 5. CENTRE c_k +@pytest.mark.numerical +@pytest.mark.parametrize("k", AXES, ids=[f"c_{LAB[a]}" for a in AXES]) +def test_fd_vs_ad_center(k): + def f(ck): + c = list(center0) + c[k] = ck + return objective(c, size0, scale0, trans0, theta0, axis0) + + g_ad = autograd.grad(f)(center0[k]) + g_fd = finite_diff(f, center0[k], Δxyz) + + plots[("c", k, "fd")].append(g_fd) + plots[("c", k, "ad")].append(g_ad) + + assert np.isfinite(g_fd) and np.isfinite(g_ad), "NaN/Inf in FD or AD result" + _assert_close(g_fd, g_ad, tag="center", axis=k) + + +# ───────── save / plot after test run ──────────────────────────── +def _save_and_plot(): + if not (SAVE or PLOT): + return + + os.makedirs(OUT, exist_ok=True) + + if SAVE: + np.savez_compressed(os.path.join(OUT, "gradients.npz"), **plots) + + labels, errs = [], [] + + for k in AXES: + fd_pairs = plots.get(("θ", k, "fd"), []) + ad_pairs = plots.get(("θ", k, "ad"), []) + for (theta, fd), (_, ad) in zip(sorted(fd_pairs), sorted(ad_pairs)): + lbl = f"θ_{LAB[k]} {theta:.2f}" + rel = abs(fd - ad) / max(abs(fd), 1e-12) + labels.append(lbl) + errs.append(rel) + + for tag, pretty in zip(("s", "t", "a", "c"), ("scale", "trans", "size", "center")): + for k in AXES: + fd_vals = plots.get((tag, k, "fd"), []) + ad_vals = plots.get((tag, k, "ad"), []) + if fd_vals: + rel = abs(fd_vals[0] - ad_vals[0]) / max(abs(fd_vals[0]), 1e-12) + labels.append(f"{pretty}_{LAB[k]}") + errs.append(rel) + + if not (PLOT and errs): + return + + plt.figure(figsize=(12, 5)) + plt.bar(range(len(errs)), errs) + plt.xticks(range(len(errs)), labels, rotation=75, ha="right", fontsize=8) + plt.ylabel("relative |FD – AD| / max(|FD|)") + plt.title("FD vs AD relative errors for all parameters") + plt.tight_layout() + fname = os.path.join(OUT, "bar_rel_error_all_params.png") + plt.savefig(fname, dpi=150) + plt.close() + print(f"[plot] saved {fname}") + + +atexit.register(_save_and_plot) diff --git a/tests/test_components/test_box_compute_derivatives.py b/tests/test_components/test_box_compute_derivatives.py new file mode 100644 index 0000000000..4060d212f4 --- /dev/null +++ b/tests/test_components/test_box_compute_derivatives.py @@ -0,0 +1,246 @@ +import atexit +import os +from collections import defaultdict + +import autograd +import autograd.numpy as anp +import matplotlib.pyplot as plt +import numpy as np +import pytest +import tidy3d as td +import tidy3d.web as web + +SAVE_RESULTS = False +PLOT_RESULTS = False +RESULTS_DIR = "./fd_ad_results" +results_collector = defaultdict(list) + +wavelength = 1.5 +freq0 = td.C_0 / wavelength +L = 10 * wavelength +buffer = 1.0 * wavelength +run_time = 120 / freq0 + + +SCENARIOS = [ + { + "name": "(1) normal", + "has_background": False, + "background_eps": 3.0, + "box_eps": 2.0, + "rotation_deg": None, + "rotation_axis": None, + }, + { + "name": "(2) perm=1.5", + "has_background": True, + "background_eps": 1.5, + "box_eps": 2.0, + "rotation_deg": None, + "rotation_axis": None, + }, + { + "name": "(3) rotation=0 deg about z", + "has_background": False, + "background_eps": 1.5, + "box_eps": 2.0, + "rotation_deg": 0.0, + "rotation_axis": 2, + }, + { + "name": "(4) rotation=90 deg about z", + "has_background": False, + "background_eps": 1.5, + "box_eps": 2.0, + "rotation_deg": 90.0, + "rotation_axis": 2, + }, + { + "name": "(5) rotation=45 deg about y", + "has_background": False, + "background_eps": 1.5, + "box_eps": 2.0, + "rotation_deg": 45.0, + "rotation_axis": 1, + }, + { + "name": "(6) rotation=45 deg about x", + "has_background": False, + "background_eps": 1.5, + "box_eps": 2.0, + "rotation_deg": 45.0, + "rotation_axis": 0, + }, + { + "name": "(7) rotation=45 deg about z", + "has_background": False, + "background_eps": 1.5, + "box_eps": 2.0, + "rotation_deg": 45.0, + "rotation_axis": 2, + }, +] + +PARAM_LABELS = ["center_x", "center_x", "center_y", "center_z", "size_x", "size_y", "size_z"] + + +def make_sim(center: tuple, size: tuple, scenario: dict): + source = td.PointDipole( + center=(-L / 2 + buffer, 0.0, 0.0), + source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10.0), + polarization="Ez", + ) + + monitor = td.FieldMonitor( + center=(+L / 2 - buffer, 0.5 * buffer, 0.5 * buffer), + size=(0, 0, 0), + freqs=[freq0], + name="point_out", + ) + + structures = [] + if scenario["has_background"]: + back_box = td.Box(center=(0.0, 0.0, 0.0), size=(4.0, 1.6, 1.6)) + background_box = td.Structure( + geometry=back_box, + medium=td.Medium(permittivity=scenario["background_eps"]), + ) + structures.append(background_box) + + scatter_box = td.Box(center=center, size=size) + + if scenario["rotation_deg"] is not None: + angle_rad = np.deg2rad(scenario["rotation_deg"]) + rotated_geom = scatter_box.rotated(angle_rad, scenario["rotation_axis"]) + else: + rotated_geom = scatter_box + + scatter_struct = td.Structure( + geometry=rotated_geom, + medium=td.Medium(permittivity=scenario["box_eps"]), + ) + structures.append(scatter_struct) + + sim = td.Simulation( + size=(L, L, L), + run_time=run_time, + grid_spec=td.GridSpec.auto(min_steps_per_wvl=50), + sources=[source], + monitors=[monitor], + structures=structures, + ) + return sim + + +def objective_fn(center, size, scenario): + sim = make_sim(center, size, scenario) + sim_data = web.run(sim, task_name="autograd_vs_fd_scenario", local_gradient=True, verbose=False) + return anp.sum(sim_data.get_intensity("point_out").values) + + +def fd_vs_ad_param(center, size, scenario, param_label, delta=1e-3): + val_and_grad_fn = autograd.value_and_grad( + lambda c, s: objective_fn(c, s, scenario), argnum=(0, 1) + ) + _, (grad_center, grad_size) = val_and_grad_fn(center, size) + + param_map = { + "center_x": (0, "center"), + "center_y": (1, "center"), + "center_z": (2, "center"), + "size_x": (0, "size"), + "size_y": (1, "size"), + "size_z": (2, "size"), + } + idx, which = param_map[param_label] + if which == "center": + ad_val = grad_center[idx] + else: + ad_val = grad_size[idx] + + center_arr = np.array(center, dtype=float) + size_arr = np.array(size, dtype=float) + + if which == "center": + cplus = center_arr.copy() + cminus = center_arr.copy() + cplus[idx] += delta + cminus[idx] -= delta + p_plus = objective_fn(tuple(cplus), tuple(size_arr), scenario) + p_minus = objective_fn(tuple(cminus), tuple(size_arr), scenario) + else: + splus = size_arr.copy() + sminus = size_arr.copy() + splus[idx] += delta + sminus[idx] -= delta + p_plus = objective_fn(tuple(center_arr), tuple(splus), scenario) + p_minus = objective_fn(tuple(center_arr), tuple(sminus), scenario) + + fd_val = (p_plus - p_minus) / (2.0 * delta) + return fd_val, ad_val, p_plus, p_minus + + +@pytest.mark.numerical +@pytest.mark.parametrize("scenario", SCENARIOS, ids=[s["name"] for s in SCENARIOS]) +@pytest.mark.parametrize( + "param_label", ["center_x", "center_y", "center_z", "size_x", "size_y", "size_z"] +) +def test_autograd_vs_fd_scenarios(scenario, param_label): + center0 = (0.0, 0.0, 0.0) + size0 = (2.0, 2.0, 2.0) + delta = 0.03 + + fd_val, ad_val, p_plus, p_minus = fd_vs_ad_param(center0, size0, scenario, param_label, delta) + + assert np.isfinite(fd_val), f"FD derivative is not finite for param={param_label}" + assert np.isfinite(ad_val), f"AD derivative is not finite for param={param_label}" + + denom = max(abs(fd_val), 1e-12) + rel_diff = abs(fd_val - ad_val) / denom + assert rel_diff < 0.3, f"Autograd vs FD mismatch: param={param_label}, diff={rel_diff:.1%}" + + results_collector[param_label].append((scenario["name"], rel_diff)) + + if SAVE_RESULTS: + os.makedirs(RESULTS_DIR, exist_ok=True) + results_data = { + "scenario_name": scenario["name"], + "param_label": param_label, + "delta": float(delta), + "fd_val": float(fd_val), + "ad_val": float(ad_val), + "p_plus": float(p_plus), + "p_minus": float(p_minus), + "rel_diff": float(rel_diff), + } + filename_npy = f"fd_ad_{scenario['name'].replace(' ', '_')}_{param_label}.npy" + np.save(os.path.join(RESULTS_DIR, filename_npy), results_data) + + +def finalize_plotting(): + if not PLOT_RESULTS: + return + + os.makedirs(RESULTS_DIR, exist_ok=True) + + for param_label in PARAM_LABELS: + scenario_data = results_collector[param_label] + if not scenario_data: + continue + scenario_names = [sd[0] for sd in scenario_data] + rel_diffs = [sd[1] for sd in scenario_data] + + plt.figure(figsize=(6, 4)) + plt.bar(scenario_names, rel_diffs, color="blue") + plt.xticks(rotation=45, ha="right") + plt.title(f"Relative Error for param='{param_label}'\n(FD vs AD)") + plt.ylabel("Relative Error") + plt.tight_layout() + + filename_png = f"rel_error_{param_label.replace('_', '-')}.png" + plt.savefig(os.path.join(RESULTS_DIR, filename_png)) + plt.close() + print(f"Saved bar chart => {filename_png}") + + +atexit.register(finalize_plotting) diff --git a/tests/test_components/test_box_rotation_derivatives.py b/tests/test_components/test_box_rotation_derivatives.py new file mode 100644 index 0000000000..71b9d04654 --- /dev/null +++ b/tests/test_components/test_box_rotation_derivatives.py @@ -0,0 +1,152 @@ +""" +Finite‑difference (FD) vs Autograd (AD) gradients with respect to the +rotation angle θ of a dielectric box, tested for rotations about all +three coordinate axes (x, y, z). A separate line plot is produced for +each axis. +""" + +import atexit +import os +from collections import defaultdict + +import autograd +import autograd.numpy as anp +import matplotlib.pyplot as plt +import numpy as np +import pytest +import tidy3d as td +import tidy3d.web as web + +# ───── global switches ──────────────────────────────────────────────────── +SAVE_RESULTS = False +PLOT_RESULTS = False +RESULTS_DIR = "./fd_ad_theta_results" + +# ───── simulation constants ─────────────────────────────────────────────── +wavelength = 1.5 +freq0 = td.C_0 / wavelength +L = 10 * wavelength +buffer = 1.0 * wavelength +run_time = 120 / freq0 + +# baseline geometry / material +center0 = (0.0, 0.0, 0.0) +size0 = (2.0, 2.0, 2.0) +eps0 = 2.0 + +# sweep parameters +delta = 0.016 # FD step +theta_sweep = np.array( + [0.0, np.pi / 4.0, np.pi / 2.0, 3 * np.pi / 4.0, np.pi, 3 * np.pi / 2.0, 2 * np.pi] +) +AXES = [0, 1, 2] # x, y, z +AXIS_LABEL = {0: "x", 1: "y", 2: "z"} + +plots = {a: defaultdict(list) for a in AXES} + + +# ───── helper: build a simulation ───────────────────────────────────────── +def make_simulation(center, size, eps, theta, axis): + """Tidy3D simulation with the box rotated by θ about the given axis.""" + source = td.PointDipole( + center=(-L / 2 + 0.5 * buffer, 0.0, 0.0), + source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10.0), + polarization="Ez", + ) + monitor = td.FieldMonitor( + center=(+L / 2 - 0.5 * buffer, 2 * buffer, 2 * buffer), + size=(0, 0, 0), + freqs=[freq0], + name="m", + ) + + box = td.Box(center=center, size=size) + geom = box.rotated(theta, axis) + struct = td.Structure(geometry=geom, medium=td.Medium(permittivity=eps)) + + return td.Simulation( + size=(L, L, L), + run_time=run_time, + grid_spec=td.GridSpec.auto(min_steps_per_wvl=50), + sources=[source], + monitors=[monitor], + structures=[struct], + ) + + +# ───── objective & FD helper ────────────────────────────────────────────── +def objective_fn(center, size, eps, theta, axis): + sim = make_simulation(center, size, eps, theta, axis) + data = web.run(sim, task_name="fd_ad_theta", verbose=False, local_gradient=True) + return anp.sum(data.get_intensity("m").values) + + +def finite_diff_theta(center, size, eps, theta, axis, delta=1e-3): + p_plus = objective_fn(center, size, eps, theta + delta, axis) + p_minus = objective_fn(center, size, eps, theta - delta, axis) + return (p_plus - p_minus) / (2.0 * delta) + + +# ───── pytest parametrised check ───────────────────────────────────────── +@pytest.mark.numerical +@pytest.mark.parametrize("axis", AXES, ids=[f"axis_{AXIS_LABEL[a]}" for a in AXES]) +@pytest.mark.parametrize("theta", theta_sweep) +def test_fd_vs_ad_theta(theta, axis): + grad_theta_fn = autograd.grad(lambda th: objective_fn(center0, size0, eps0, th, axis)) + + fd_val = finite_diff_theta(center0, size0, eps0, theta, axis, delta) + ad_val = grad_theta_fn(theta) + + # keep for plotting + plots[axis]["θ"].append(theta) + plots[axis]["fd"].append(fd_val) + plots[axis]["ad"].append(ad_val) + + # sanity & agreement checks + assert np.isfinite(fd_val) and np.isfinite(ad_val) + rel_diff = abs(fd_val - ad_val) / max(abs(fd_val), 1e-12) + assert rel_diff < 0.3, f"axis={axis} θ={theta:.3f}: FD={fd_val:.4e}, AD={ad_val:.4e}" + + if SAVE_RESULTS: + os.makedirs(RESULTS_DIR, exist_ok=True) + np.savez( + os.path.join(RESULTS_DIR, f"axis_{AXIS_LABEL[axis]}_θ_{theta:.4f}.npz"), + theta=float(theta), + axis=int(axis), + fd=float(fd_val), + ad=float(ad_val), + ) + + +# ───── plot after pytest run ────────────────────────────────────────────── +def finalize_plotting(): + if not PLOT_RESULTS: + return + + os.makedirs(RESULTS_DIR, exist_ok=True) + + for axis in AXES: + if not plots[axis]["θ"]: + continue + + idx = np.argsort(plots[axis]["θ"]) + θ = np.array(plots[axis]["θ"])[idx] + fd = np.array(plots[axis]["fd"])[idx] + ad = np.array(plots[axis]["ad"])[idx] + + plt.figure(figsize=(6, 4)) + plt.plot(θ, fd, label="Finite diff", marker="o") + plt.plot(θ, ad, label="Autograd", marker="s", linestyle="--") + plt.xlabel("θ [rad]") + plt.ylabel("∂(intensity)/∂θ") + plt.title(f"FD vs AD gradient (rotation about {AXIS_LABEL[axis]}‑axis)") + plt.legend() + plt.tight_layout() + + fname = os.path.join(RESULTS_DIR, f"grad_theta_fd_vs_ad_axis_{AXIS_LABEL[axis]}.png") + plt.savefig(fname, dpi=150) + plt.close() + print(f"[plot] saved ⇒ {fname}") + + +atexit.register(finalize_plotting) diff --git a/tidy3d/components/geometry/base.py b/tidy3d/components/geometry/base.py index 9fddfddc72..d5a9dbec62 100644 --- a/tidy3d/components/geometry/base.py +++ b/tidy3d/components/geometry/base.py @@ -10,7 +10,6 @@ import autograd.numpy as np import pydantic.v1 as pydantic import shapely -import xarray as xr try: from matplotlib import patches @@ -28,7 +27,10 @@ from ...log import log from ...packaging import check_import, verify_packages_import from ..autograd import AutogradFieldMap, TracedCoordinate, TracedSize, get_static -from ..autograd.derivative_utils import DerivativeInfo, integrate_within_bounds +from ..autograd.derivative_utils import ( + DerivativeInfo, + DerivativeSurfaceMesh, +) from ..base import Tidy3dBaseModel, cached_property from ..transformation import ReflectionFromPlane, RotationAroundAxis from ..types import ( @@ -61,6 +63,7 @@ ) POLY_GRID_SIZE = 1e-12 +_BOX_FACE_GRID_SIZE = 1e-3 _shapely_operations = { @@ -973,7 +976,9 @@ def rotated(self, angle: float, axis: Union[Axis, Coordinate]) -> Geometry: :class:`Geometry` Rotated copy of this geometry. """ - return Transformed(geometry=self, transform=Transformed.rotation(angle, axis)) + T = Transformed.rotation(angle, axis) + rotated_geom = Transformed(geometry=self, transform=T) + return rotated_geom def reflected(self, normal: Coordinate) -> Geometry: """Return a reflected copy of this geometry. @@ -1653,6 +1658,133 @@ def __invert__(self): operation="difference", geometry_a=Box(size=(inf, inf, inf)), geometry_b=self ) + def build_box_face_mesh( + self, + center: np.ndarray, + size: np.ndarray, + axis_normal: int, # 0,1,2 → x,y,z faces + min_max_index: int, # 0 = − side, 1 = + side + linear_matrix: np.ndarray | None = None, # 3×3 rotation/scale/shear + translation: np.ndarray | None = None, + ) -> tuple[DerivativeSurfaceMesh, np.ndarray]: + """Build a mesh for the face of a box, given the center, size, axis normal, and min/max index. + The mesh is built in the local coordinate system of the box, and then transformed to the global + coordinate system using the transformation matrix.""" + + if axis_normal == 0: + canonical_normal = np.array([1.0, 0.0, 0.0]) + elif axis_normal == 1: + canonical_normal = np.array([0.0, 1.0, 0.0]) + elif axis_normal == 2: + canonical_normal = np.array([0.0, 0.0, 1.0]) + else: + raise ValueError("Invalid axis_normal") + + if min_max_index == 0: + canonical_normal *= -1.0 + + if linear_matrix is None: + linear_matrix = np.eye(3) + if translation is None: + translation = np.zeros(3) + + n_local = np.linalg.inv(linear_matrix).T @ canonical_normal + n_local = n_local / np.linalg.norm(n_local) + + def compute_tangential_vectors( + normal: np.ndarray, eps: float = 1e-8 + ) -> tuple[np.ndarray, np.ndarray]: + """Compute any two perpendicular tangential vectors t1, t2, given a normal.""" + if abs(normal[0]) > abs(normal[2]): + t1 = np.array([-normal[1], normal[0], 0.0]) + else: + t1 = np.array([0.0, -normal[2], normal[1]]) + t1_norm = np.linalg.norm(t1) + if t1_norm < eps: + raise ValueError("Degenerate normal vector.") + t1 = t1 / t1_norm + t2 = np.cross(normal, t1) + t2 /= np.linalg.norm(t2) + return t1, t2 + + t1_local, t2_local = compute_tangential_vectors(n_local) + + min_bound = np.array(center) - np.array(size) / 2.0 + max_bound = np.array(center) + np.array(size) / 2.0 + bounds_old = np.column_stack((min_bound, max_bound)) + + corners = np.array( + [ + [bounds_old[0, i], bounds_old[1, j], bounds_old[2, k]] + for i in (0, 1) + for j in (0, 1) + for k in (0, 1) + ] + ) + + connectivity = { + 0: { # Faces perpendicular to x-axis + 0: [0, 1, 3, 2], + 1: [4, 5, 7, 6], + }, + 1: { # Faces perpendicular to y-axis + 0: [0, 1, 5, 4], + 1: [2, 3, 7, 6], + }, + 2: { # Faces perpendicular to z-axis + 0: [0, 4, 6, 2], + 1: [1, 5, 7, 3], + }, + } + + face_indices = connectivity[axis_normal][min_max_index] + face_corners = corners[face_indices, :] + + transformed_corners = (linear_matrix @ face_corners.T).T + translation + p1, p2, p3, p4 = transformed_corners + + # Determine number of points along each face edge based on _BOX_FACE_GRID_SIZE + edge1 = transformed_corners[1] - transformed_corners[0] + edge2 = transformed_corners[2] - transformed_corners[0] + len_s = np.linalg.norm(edge1) + len_t = np.linalg.norm(edge2) + num_s = max(1, int(np.ceil(len_s / _BOX_FACE_GRID_SIZE))) + num_t = max(1, int(np.ceil(len_t / _BOX_FACE_GRID_SIZE))) + s_vals = np.linspace(0, 1, 2 * num_s + 1)[1::2] + t_vals = np.linspace(0, 1, 2 * num_t + 1)[1::2] + S, T = np.meshgrid(s_vals, t_vals, indexing="ij") + + X = (1 - S) * (1 - T) * p1[0] + S * (1 - T) * p2[0] + S * T * p3[0] + (1 - S) * T * p4[0] + Y = (1 - S) * (1 - T) * p1[1] + S * (1 - T) * p2[1] + S * T * p3[1] + (1 - S) * T * p4[1] + Z = (1 - S) * (1 - T) * p1[2] + S * (1 - T) * p2[2] + S * T * p3[2] + (1 - S) * T * p4[2] + + centers = np.column_stack([X.ravel(), Y.ravel(), Z.ravel()]) + + tri1_area = 0.5 * np.linalg.norm(np.cross((p2 - p1), (p3 - p1))) + tri2_area = 0.5 * np.linalg.norm(np.cross((p4 - p1), (p3 - p1))) + face_area = tri1_area + tri2_area + + num_cells = (num_s) * (num_t) + if num_cells > 0: + cell_area = face_area / num_cells + else: + cell_area = face_area + + areas = cell_area * np.ones(centers.shape[0]) + + normals = np.tile(n_local, (centers.shape[0], 1)) + perps1 = np.tile(t1_local, (centers.shape[0], 1)) + perps2 = np.tile(t2_local, (centers.shape[0], 1)) + + surface_mesh = DerivativeSurfaceMesh( + centers=centers, + areas=areas, + normals=normals, + perps1=perps1, + perps2=perps2, + ) + return surface_mesh, n_local + """ Abstract subclasses """ @@ -1993,8 +2125,7 @@ def _normal_axis(self) -> Axis: """Axis normal to the Box. Errors if box is not planar.""" if self.size.count(0.0) != 1: raise ValidationError( - "Tried to get 'normal_axis' of 'Box' that is not planar. " - f"Given 'size={self.size}.'" + f"Tried to get 'normal_axis' of 'Box' that is not planar. Given 'size={self.size}.'" ) return self.size.index(0.0) @@ -2496,11 +2627,15 @@ def _surface_area(self, bounds: Bound) -> float: """ Autograd code """ - def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: - """Compute the adjoint derivatives for this object.""" + def compute_derivatives( + self, derivative_info: DerivativeInfo, transform_matrix: np.ndarray = None + ) -> AutogradFieldMap: + """Compute the adjoint derivatives for this object with optional rotation matrix.""" - # get gradients w.r.t. each of the 6 faces (in normal direction) - vjps_faces = self.derivative_faces(derivative_info=derivative_info) + # Compute gradients w.r.t. each of the 6 faces using the rotation matrix (if any) + vjps_faces = self.derivative_faces( + derivative_info=derivative_info, transform_matrix=transform_matrix + ) # post-process these values to give the gradients w.r.t. center and size vjps_center_size = self.derivatives_center_size(vjps_faces=vjps_faces) @@ -2538,8 +2673,10 @@ def derivatives_center_size(vjps_faces: Bound) -> dict[str, Coordinate]: size=tuple(vjp_size.tolist()), ) - def derivative_faces(self, derivative_info: DerivativeInfo) -> Bound: - """Derivative with respect to normal position of 6 faces of ``Box``.""" + def derivative_faces( + self, derivative_info: DerivativeInfo, transform_matrix: np.ndarray = None + ) -> Bound: + """Compute derivatives with respect to the normal position of 6 faces of `Box`, using rotation matrix if provided.""" # change in permittivity between inside and outside vjp_faces = np.zeros((2, 3)) @@ -2550,6 +2687,7 @@ def derivative_faces(self, derivative_info: DerivativeInfo) -> Bound: min_max_index=min_max_index, axis_normal=axis, derivative_info=derivative_info, + transform_matrix=transform_matrix, ) # record vjp for this face @@ -2560,105 +2698,28 @@ def derivative_faces(self, derivative_info: DerivativeInfo) -> Bound: def derivative_face( self, min_max_index: int, - axis_normal: Axis, + axis_normal: int, derivative_info: DerivativeInfo, + transform_matrix: np.ndarray = None, ) -> float: - """Compute the derivative w.r.t. shifting a face in the normal direction.""" - - # normal and tangential dims - dim_normal, dims_perp = self.pop_axis("xyz", axis=axis_normal) - fld_normal, flds_perp = self.pop_axis(("Ex", "Ey", "Ez"), axis=axis_normal) + """ + Compute the derivative (VJP) with respect to shifting a face of a rotated box, + using full integration over that face. + """ + if transform_matrix is None: + L_f = np.eye(3) # linear part (no rotation/scale/shear) + t_f = np.zeros(3) # translation - # normal and tangential fields - D_normal = derivative_info.D_der_map[fld_normal].sel(f=derivative_info.frequency) - Es_perp = tuple( - derivative_info.E_der_map[key].sel(f=derivative_info.frequency) for key in flds_perp + mesh, _ = self.build_box_face_mesh( + center=np.asarray(self.center, float), + size=np.asarray(self.size, float), + axis_normal=axis_normal, + min_max_index=min_max_index, + linear_matrix=L_f, + translation=t_f, ) - - # normal and tangential bounds - bounds_T = np.array(derivative_info.bounds).T # put (xyz) first dimension - bounds_normal, bounds_perp = self.pop_axis(bounds_T, axis=axis_normal) - - # define the integration plane - coord_normal_face = bounds_normal[min_max_index] - bounds_perp = np.array(bounds_perp).T # put (min / max) first dimension for integrator - - # normal field data coordinates - fld_coords_normal = D_normal.coords[dim_normal] - - # condition: a face is entirely outside of the domain, skip! - sign = (-1, 1)[min_max_index] - normal_coord_positive = sign * coord_normal_face - fld_coords_positive = sign * fld_coords_normal - if all(fld_coords_positive < normal_coord_positive): - log.info( - f"skipping VJP for 'Box' face '{dim_normal}{'-+'[min_max_index]}' " - "as it is entirely outside of the simulation domain." - ) - return 0.0 - - # grab permittivity data inside and outside edge in normal direction - eps_xyz = [ - derivative_info.eps_data[f"eps_{dim}{dim}"].sel(f=derivative_info.frequency) - for dim in "xyz" - ] - - # number of cells from the edge of data to register "inside" (index = num_cells_in - 1) - num_cells_in = 4 - - # if not enough data, just use best guess using eps in medium and simulation - needs_eps_approx = any(len(eps.coords[dim_normal]) <= num_cells_in for eps in eps_xyz) - - if derivative_info.eps_approx or needs_eps_approx: - eps_xyz_inside = 3 * [derivative_info.eps_in] - eps_xyz_outside = 3 * [derivative_info.eps_out] - # TODO: not tested... - - # otherwise, try to grab the data at the edges - else: - if min_max_index == 0: - index_out, index_in = (0, num_cells_in - 1) - else: - index_out, index_in = (-1, -num_cells_in) - eps_xyz_inside = [eps.isel(**{dim_normal: index_in}) for eps in eps_xyz] - eps_xyz_outside = [eps.isel(**{dim_normal: index_out}) for eps in eps_xyz] - - # put in normal / tangential basis - eps_in_normal, eps_in_perps = self.pop_axis(eps_xyz_inside, axis=axis_normal) - eps_out_normal, eps_out_perps = self.pop_axis(eps_xyz_outside, axis=axis_normal) - - # compute integration pre-factors - delta_eps_perps = [eps_in - eps_out for eps_in, eps_out in zip(eps_in_perps, eps_out_perps)] - delta_eps_inv_normal = 1.0 / eps_in_normal - 1.0 / eps_out_normal - - def integrate_face(arr: xr.DataArray) -> complex: - """Interpolate and integrate a scalar field data over the face using bounds.""" - - arr_at_face = arr.interp(**{dim_normal: float(coord_normal_face)}, assume_sorted=True) - - integral_result = integrate_within_bounds( - arr=arr_at_face, - dims=dims_perp, - bounds=bounds_perp, - ) - - return complex(integral_result) - - # put together VJP using D_normal and E_perp integration - vjp_value = 0.0 - - # perform D-normal integral - integrand_D = -delta_eps_inv_normal * D_normal - integral_D = integrate_face(integrand_D) - vjp_value += integral_D - - # perform E-perpendicular integrals - for E_perp, delta_eps_perp in zip(Es_perp, delta_eps_perps): - integrand_E = E_perp * delta_eps_perp - integral_E = integrate_face(integrand_E) - vjp_value += integral_E - - return np.real(vjp_value) + vjp = derivative_info.grad_surfaces(surface_mesh=mesh) + return float(np.real(np.sum(vjp))) """Compound subclasses""" @@ -2685,7 +2746,15 @@ def _transform_is_invertible(cls, val): @pydantic.validator("geometry") def _geometry_is_finite(cls, val): - if not np.isfinite(val.bounds).all(): + def preprocess(value): + return value._value if isinstance(value, np.numpy_boxes.ArrayBox) else value + + processed_bounds = tuple( + tuple(preprocess(coord) for coord in bound) for bound in val.bounds + ) + + # Ensure all values are finite + if not np.isfinite(processed_bounds).all(): raise ValidationError( "Transformations are only supported on geometries with finite dimensions. " "Try using a large value instead of 'inf' when creating geometries that undergo " @@ -2889,9 +2958,17 @@ def rotation(angle: float, axis: Union[Axis, Coordinate]) -> MatrixReal4x4: numpy.ndarray Transform matrix with shape (4, 4). """ - transform = np.eye(4) - transform[:3, :3] = RotationAroundAxis(angle=angle, axis=axis).matrix - return transform + + R = RotationAroundAxis(angle=angle, axis=axis).matrix + # Avoid inplace ops + zeros_col = np.zeros((3, 1), dtype=R.dtype) + top = np.concatenate([R, zeros_col], axis=1) # (3, 4) + + bottom = np.concatenate( + [np.zeros((1, 3), dtype=R.dtype), np.ones((1, 1), dtype=R.dtype)], axis=1 + ) # (1, 4) + + return np.concatenate([top, bottom], axis=0) @staticmethod def reflection(normal: Coordinate) -> MatrixReal4x4: @@ -2958,6 +3035,85 @@ def _update_from_bounds(self, bounds: Tuple[float, float], axis: Axis) -> Transf new_geometry = self.geometry._update_from_bounds(bounds=new_bounds, axis=axis) return self.updated_copy(geometry=new_geometry) + def _compute_transform_derivative(self, derivative_info: DerivativeInfo): + """Return the 3×3 tensor ∂J/∂R accumulated over all six faces. + + The rotation matrix **R** is reconstructed from the first entry in + ``self.transform_history`` (which must be a *rotate* operation). + """ + # --------- pull the final affine map once -------------------------------- + L_f = self.transform[:3, :3] # no assumption on orthogonality + t_f = self.transform[:3, 3] + L_f_inv = np.linalg.inv(L_f) + + center = np.asarray(self.geometry.center, float) # local box frame + size = np.asarray(self.geometry.size, float) + + dJ_dL_f = np.zeros((3, 3)) + dJ_dt_f = np.zeros(3) + + for ax in (0, 1, 2): # x-, y-, z-faces + for side in (0, 1): # − / + + mesh, n_face = self.build_box_face_mesh( + center=center, + size=size, + axis_normal=ax, + min_max_index=side, + linear_matrix=L_f, + translation=t_f, + ) + + g = derivative_info.grad_surfaces(mesh).values.real.ravel() # (N,) + + # ---- linear part ------------------------------------------------ + X_ref = (L_f_inv @ (mesh.centers - t_f).T).T # (N,3) + dJ_dL_f += n_face[:, None] * (g[:, None] * X_ref).sum(axis=0) + + # ---- translation part ------------------------------------------ + dJ_dt_f += n_face * g.sum() # g_i n_i + + dJ_dT = np.zeros((4, 4)) + dJ_dT[:3, :3] = dJ_dL_f + dJ_dT[:3, 3] = dJ_dt_f + return dJ_dT + + def compute_derivatives(self, derivative_info: DerivativeInfo): + """Compute adjoint derivatives for the transformed geometry.""" + derivative_map = {} + + # ---------- split requested paths ------------------------------------ + transform_paths = [p for p in derivative_info.paths if p[0] == "transform"] + if derivative_info.paths == [("transform",)]: + derivative_info = derivative_info.updated_copy( + paths=[("geometry", "center"), ("geometry", "size"), ("transform",)], + deep=False, + ) + geometry_paths = [p for p in derivative_info.paths if p[0] == "geometry"] + + if "transform" in [p[0] for p in transform_paths]: + transform_paths = [("transform", i, j) for i in range(4) for j in range(4)] + + # ---------- geometry-level derivatives ------------------------------ + if geometry_paths: + geo_info = derivative_info.updated_copy( + paths=[p[1:] for p in geometry_paths], deep=False + ) + geo_derivs = self.geometry.compute_derivatives(geo_info, self.transform) + + derivative_map[("geometry", "center")] = np.asarray( + geo_derivs.get(("center",), (0.0, 0.0, 0.0)), float + ) + derivative_map[("geometry", "size")] = np.asarray( + geo_derivs.get(("size",), (0.0, 0.0, 0.0)), float + ) + + # ---------- transform-level derivatives ----------------------------- + if any(p[0] == "transform" for p in derivative_info.paths): + dJ_dT = self._compute_transform_derivative(derivative_info) + derivative_map[("transform",)] = dJ_dT + + return derivative_map + class ClipOperation(Geometry): """Class representing the result of a set operation between geometries.""" diff --git a/tidy3d/components/transformation.py b/tidy3d/components/transformation.py index ea7c1ee494..5b3ca3c095 100644 --- a/tidy3d/components/transformation.py +++ b/tidy3d/components/transformation.py @@ -5,8 +5,10 @@ from abc import ABC, abstractmethod from typing import Union +import autograd.numpy as anp import numpy as np import pydantic.v1 as pd +from autograd.numpy.numpy_boxes import ArrayBox from ..constants import RADIAN from ..exceptions import ValidationError @@ -44,7 +46,6 @@ def rotate_vector(self, vector: ArrayFloat2D) -> ArrayFloat2D: if self.isidentity: return vector - if len(vector.shape) == 1: return self.matrix @ vector @@ -66,7 +67,6 @@ def rotate_tensor(self, tensor: TensorReal) -> TensorReal: if self.isidentity: return tensor - return np.matmul(self.matrix, np.matmul(tensor, self.matrix.T)) @@ -114,15 +114,14 @@ def isidentity(self) -> bool: def matrix(self) -> TensorReal: """Rotation matrix.""" - if self.isidentity: + if self.isidentity and not isinstance(self.angle, ArrayBox): return np.eye(3) - - norm = np.linalg.norm(self.axis) + norm = anp.linalg.norm(self.axis) n = self.axis / norm - c = np.cos(self.angle) - s = np.sin(self.angle) - K = np.array([[0, -n[2], n[1]], [n[2], 0, -n[0]], [-n[1], n[0], 0]]) - R = np.eye(3) + s * K + (1 - c) * K @ K + c = anp.cos(self.angle) + s = anp.sin(self.angle) + K = anp.array([[0, -n[2], n[1]], [n[2], 0, -n[0]], [-n[1], n[0], 0]]) + R = anp.eye(3) + s * K + (1 - c) * K @ K return R diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index c5e1f506b2..1ff9ccb78a 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -1013,32 +1013,28 @@ def postprocess_adj( eps_background = None # manually override simulation medium as the background structure - if not isinstance(structure.geometry, td.Box): - # auto permittivity detection - sim_orig = sim_data_orig.simulation - plane_eps = eps_fwd.monitor.geometry - - # get permittivity without this structure - structs_no_struct = list(sim_orig.structures) - structs_no_struct.pop(structure_index) - sim_no_structure = sim_orig.updated_copy(structures=structs_no_struct) - eps_no_structure = sim_no_structure.epsilon( - box=plane_eps, coord_key="centers", freq=freq_adj - ) - - # get permittivity with structures on top of an infinite version of this structure - structs_inf_struct = list(sim_orig.structures)[structure_index + 1 :] - sim_inf_structure = sim_orig.updated_copy( - structures=structs_inf_struct, - medium=structure.medium, - monitors=[], - ) - eps_inf_structure = sim_inf_structure.epsilon( - box=plane_eps, coord_key="centers", freq=freq_adj - ) + # auto permittivity detection + sim_orig = sim_data_orig.simulation + plane_eps = eps_fwd.monitor.geometry + + # get permittivity without this structure + structs_no_struct = list(sim_orig.structures) + structs_no_struct.pop(structure_index) + sim_no_structure = sim_orig.updated_copy(structures=structs_no_struct) + eps_no_structure = sim_no_structure.epsilon( + box=plane_eps, coord_key="centers", freq=freq_adj + ) - else: - eps_no_structure = eps_inf_structure = None + # get permittivity with structures on top of an infinite version of this structure + structs_inf_struct = list(sim_orig.structures)[structure_index + 1 :] + sim_inf_structure = sim_orig.updated_copy( + structures=structs_inf_struct, + medium=structure.medium, + monitors=[], + ) + eps_inf_structure = sim_inf_structure.epsilon( + box=plane_eps, coord_key="centers", freq=freq_adj + ) # get minimum intersection of bounds with structure and sim struct_bounds = rmin_struct, rmax_struct = structure.geometry.bounds