Skip to content

Commit

Permalink
fix bug in autograd differentiation cylinder.center
Browse files Browse the repository at this point in the history
  • Loading branch information
tylerflex committed Nov 1, 2024
1 parent 76ab553 commit b0a92fd
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 28 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Fixed
- Regression in local field projection leading to incorrect results for `far_field_approx=True`.
- Bug when differentiating with respect to `Cylinder.center`.


## [2.7.6] - 2024-10-30

Expand Down
66 changes: 45 additions & 21 deletions tests/test_components/test_autograd.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,10 +395,10 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
custom_pole_res = td.Structure(geometry=box, medium=custom_med_pole_res)

radius = 0.4 * (1 + anp.abs(vector @ params))
cyl_center_y = 0 # vector @ params
cyl_center_z = 0 # -vector @ params
cyl_center_y = vector @ params
cyl_center_z = -vector @ params
cylinder_geo = td.Cylinder(
radius=radii,
radius=anp.mean(radii) * 0.5,
center=(0, cyl_center_y, cyl_center_z),
axis=0,
length=LX / 2 if IS_3D else td.inf,
Expand Down Expand Up @@ -770,43 +770,67 @@ def objective(*args):
def test_autograd_polyslab_cylinder(use_emulated_run, monitor_key):
"""Test an objective function through tidy3d autograd."""

fn_dict_polyslab = get_functions("polyslab", monitor_key)
make_sim_polyslab = fn_dict_polyslab["sim"]
t = 1.0
axis = 0

fn_dict_cylinder = get_functions("cylinder", monitor_key)
make_sim_cylinder = fn_dict_cylinder["sim"]
num_pts = 89

postprocess = fn_dict_cylinder["postprocess"]
monitor, postprocess = make_monitors()[monitor_key]

def make_cylinder(radius, x0, y0):
return td.Cylinder(
center=td.Cylinder.unpop_axis(0.0, (x0, y0), axis=axis),
radius=radius,
length=t,
axis=axis,
) # .to_polyslab(num_pts)

def make_polyslab(radius, x0, y0):
phis = anp.linspace(0, 2 * np.pi, num_pts + 1)[:-1]

xs = radius * anp.cos(phis) + x0
ys = radius * anp.sin(phis) + y0

vertices = anp.stack((xs, ys), axis=-1)

return td.PolySlab(
vertices=vertices,
axis=axis,
slab_bounds=(-t / 2, t / 2),
)

def objective_polyslab(*args):
def make_sim(params, geo_maker):
geo = geo_maker(*params)
structure = td.Structure(geometry=geo, medium=td.Medium(permittivity=2))

return SIM_BASE.updated_copy(structures=[structure], monitors=[monitor])

p0 = [1.0, 0.0, 0.0]

def objective_polyslab(params):
"""Objective function."""
sim = make_sim_polyslab(*args)
sim = make_sim(params, geo_maker=make_polyslab)
if PLOT_SIM:
plot_sim(sim, plot_eps=True)
data = run(sim, task_name="autograd_test", verbose=False)
value = postprocess(data)
return value
return anp.sum(anp.abs(data[monitor.name].amps)).item()

val_polyslab, grad_polyslab = ag.value_and_grad(objective_polyslab)(params0)
val_polyslab, grad_polyslab = ag.value_and_grad(objective_polyslab)(p0)
print(val_polyslab, grad_polyslab)
assert anp.all(grad_polyslab != 0.0), "some gradients are 0"

def objective_cylinder(*args):
def objective_cylinder(params):
"""Objective function."""
sim = make_sim_cylinder(*args)
sim = make_sim(params, geo_maker=make_cylinder)
if PLOT_SIM:
plot_sim(sim, plot_eps=True)
data = run(sim, task_name="autograd_test", verbose=False)
value = postprocess(data)
return value
return anp.sum(anp.abs(data[monitor.name].amps)).item()

val_cylinder, grad_cylinder = ag.value_and_grad(objective_cylinder)(params0)
val_cylinder, grad_cylinder = ag.value_and_grad(objective_cylinder)(p0)
print(val_cylinder, grad_cylinder)
assert anp.all(grad_cylinder != 0.0), "some gradients are 0"

# just make sure they're somewhat close (use different discretizations)
assert np.allclose(grad_cylinder, grad_polyslab, rtol=0.3)


@pytest.mark.parametrize("structure_key, monitor_key", args)
def test_autograd_server(use_emulated_run, structure_key, monitor_key):
Expand Down
13 changes: 6 additions & 7 deletions tidy3d/components/geometry/primitives.py
Original file line number Diff line number Diff line change
Expand Up @@ -301,14 +301,13 @@ def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldM
vjps_vertices_xs, vjps_vertices_ys = vjps_polyslab[("vertices",)].T

# transform polyslab vertices derivatives into Cylinder parameter derivatives
xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference)
vjp_xs = np.sum(xs_ * vjps_vertices_xs)
vjp_ys = np.sum(ys_ * vjps_vertices_ys)

vjps = {}
for path in derivative_info.paths:
if path == ("radius",):
xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference)

vjp_xs = np.sum(xs_ * vjps_vertices_xs)
vjp_ys = np.sum(ys_ * vjps_vertices_ys)

vjps[path] = vjp_xs + vjp_ys

elif "center" in path:
Expand All @@ -322,9 +321,9 @@ def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldM

_, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis)
if center_index == index_x:
vjps[path] = np.sum(vjp_xs)
vjps[path] = np.sum(vjps_vertices_xs)
elif center_index == index_y:
vjps[path] = np.sum(vjp_ys)
vjps[path] = np.sum(vjps_vertices_ys)
else:
raise ValueError(
"Something unexpected happened. Was asked to differentiate "
Expand Down

0 comments on commit b0a92fd

Please sign in to comment.