fix[autograd]: sample fields along slab height and polygon edges #2418

yaugenst-flex · 2025-04-28T15:26:06Z

Previously face gradients came from a fixed 30 × 30 grid so the accuracy depended on the size of the polygon. Edge gradients were evaluated only at the center of each face, so long edges and tall slabs produced incorrect adjoint results depending on the simulation.

Changed:

Adaptive, spacing-based sampling for both faces and vertices/edges.
Extra edge-cases covered: slab fully outside sim, faces at +/- inf, and fix for correct normal sign for axis==1.

I also added some analytic tests for the face and edge integration.

The sample spacing is currently set to 50nm, which I think should be reasonable in most cases but is another one of those that we might either want to make adaptive later or expose in a config. For vertices spaced <50nm apart, the results are identical to before (tested on the shape optimization notebook).

@bzhangflex fyi

yaugenst-flex · 2025-04-28T17:43:58Z

looks like some tests are failing due to the previously unhandled edge cases. have to double check to be sure though:

FAILED tests/test_components/test_autograd.py::test_multi_frequency_equivalence[polyslab] - AssertionError
FAILED tests/test_components/test_autograd.py::test_autograd_async_some_zero_grad[polyslab-<ALL>] - AssertionError: some gradients are 0
FAILED tests/test_components/test_autograd.py::test_autograd_server[polyslab-<ALL>] - AssertionError: some gradients are 0
FAILED tests/test_components/test_autograd.py::test_autograd_objective[polyslab-<ALL>] - AssertionError: some gradients are 0
FAILED tests/test_components/test_autograd.py::test_autograd_async[polyslab-<ALL>] - AssertionError: some gradients are 0
FAILED tests/test_components/test_autograd.py::test_autograd_async_server[polyslab-<ALL>] - AssertionError: some gradients are 0

tylerflex

Thanks @yaugenst-flex. This is pretty complex, I think it will require a lot of testing for both performance and gradient accuracy. Probably in all of the notebooks. I'd say maybe best left to 2.9 to be safe. how much testing in notebooks have you done already? either way will be great to have more accurate gradients for PolySlab!

CHANGELOG.md

tylerflex · 2025-04-28T19:08:57Z

tests/test_components/test_autograd.py

+    We test ``PolySlab.compute_derivatives`` by feeding it a dummy ``DerivativeInfo``
+    whose ``grad_surfaces()`` method returns the analytically chosen surface integrand
+
+        k(x, y) = ax + by + c


nice, this is smart

how do we test z / slab_bounds?

the integrand depends only on the in-plane coordinates so the derivative w.r.t. slab_bounds is just the surface integral over k. we test the slab_bounds for the three edge cases of fully inside, partially inside, and fully outside faces, where the face that lies outside of the simulation (ie the slab_bounds for that index) get a derivative of zero for that side.

or do you mean testing the integration along z by giving a variable z? that I should probably add yeah

or do you mean testing the integration along z by giving a variable z? that I should probably add yeah

yea this

done. the analytical tests got kind of out of hand but I think everything works...

tylerflex · 2025-04-28T19:11:12Z

tidy3d/components/geometry/polyslab.py

@@ -47,8 +47,8 @@

 _MIN_POLYGON_AREA = fp_eps

-# number of points per dimension when discretizing polyslab faces for slab bounds gradients
-_NUM_PTS_DIM_SLAB_BOUNDS = 30
+# spacing between points when discretizing polyslab faces and edges for gradients


probably just a points per wavelength would make most sense generally. For example this could easily become an issue with RF wavelengths?

so my thinking here was that this is much better than before, and checking for wavelength in material adds some additional complexity because we can't only depend on DerivativeInfo.frequency, or at least it's not clear to me how to make this very robust. in principle the best solution would be adaptive integration, which was actually the first implementation that I had but it was pretty slow because it ended up needing a lot of sampling points to actually converge. so basically what I'm trying to say is I'd leave this as is for now and do the adaptive sampling in a separate pr if it ends up being a problem...

what we did in adjoint was pass the min wavelength in the Simulation.sources and divide by the minimum refractive index in the Simulation.

tidy3d/tidy3d/plugins/adjoint/components/structure.py

Lines 101 to 106 in 1000bda

freq_max = float(max(grad_data_eps.eps_xx.f))

eps_in = self.medium.eps_model(frequency=freq_max)

ref_ind = np.sqrt(np.max(np.real(eps_in)))

ref_ind = max([1.0, abs(ref_ind)])

wvl_free_space = C_0 / freq_max

wvl_mat = wvl_free_space / ref_ind

https://github.com/flexcompute/tidy3d/blob/develop/tidy3d/plugins/adjoint/components/geometry.py#L203

does this work?

interesting approach 😄 this is a bit at odds with the performance concern though isnt it? because it can end up severely oversampling the polygon. what happens if there are metals in the simulation?

for metals, the ref_ind would evaluate to 1. I think we can effectively cap the number of points , but this at least scales the sampling to the amount of detail in the fields so I think that's the proper way to do it.

I'll second that I think this spacing should scale with wavelength as well so that gradient accuracy is not dependent on wavelength (and agreed that in the case of RF, we will end up with really fine sampling otherwise)

my suspicion for metals is that the vjp sampling might need to be based on the skin depth, ie I think we might actually need really fine sampling. I doubt we'd get away with 1/um for metals... anyhow. I added wavelength-adaptive sampling to the DerivativeInfo and also implemented the skin-depth based approach. however, I'm setting the minimum allowable sampling distance to 10nm currently so the skin depth will actually never be resolved. once we figure this out you'll want to take another look at this @groberts-flex

that sounds great to me! will definitely do some testing with the pec/metal version to make sure the sampling is right

tidy3d/components/geometry/polyslab.py

tylerflex · 2025-04-28T19:27:53Z

tidy3d/components/geometry/polyslab.py

+        ax_val = ax_min if min_max_index == 0 else ax_max
+        accum = 0.0
+
+        # intentionally not vectorized to limit memory use


have you tested this for speed?

I have, yes. I originally had a fully vectorized implementation that ran basically as fast as the single sample point but it'd eat up way too much memory for large polygons, so I decided to go back to the non-vectorized version. in the limit where the vertices are already sampled finely (like in the shape optimized taper notebook), the speed is pretty much identical to before because it does the same thing (takes just an edge center sample). it basically only kicks in when the edges are longer than the sampling distance, in which case it is slower than previously, but correct 😄 for thin slabs it's basically the same thing as the "hack" of subdividing the vertices.

tidy3d/components/geometry/polyslab.py

tylerflex · 2025-04-28T19:41:08Z

tidy3d/components/geometry/polyslab.py

+                if not mask.any():
+                    continue
+
+                mesh = DerivativeSurfaceMesh(


I worry this approach will be quite slow. it's the same as the adjoint plugin implementation (also vectorized over each edge xy only, but looped over z points and edges). We had complaints about speed back then. have you tested this on some problems? such as the taper?

it depends on the problem. but yeah it's ~n times slower than without the z sampling, where n is the number of z samples. for a typical SOI device it shouldn't be making a huge difference I think. for a tall slab it will be significantly slower, but especially in that case it's necessary.

just curious to understand speed of things like this versus simulation time length? are they approaching the length of the simulations or are they just relatively slow compared to other parts of the gradient calculation? This will be dependent on the exact simulation setting, but curious for some average 3D case.

tidy3d/components/geometry/polyslab.py

groberts-flex

This is really exciting and will be a huge improvement. Is it possible to run some of the numerical testing and possibly modify the polyslab numerical test file as well to make sure we can test going forward the surface integration (and keep testing cases where it is adding the most improvement).

tidy3d/components/autograd/derivative_utils.py

groberts-flex

updates look good to me!

groberts-flex · 2025-04-30T21:02:56Z

tidy3d/components/geometry/polyslab.py

@@ -47,9 +47,6 @@

 _MIN_POLYGON_AREA = fp_eps

-# number of points per dimension when discretizing polyslab faces for slab bounds gradients
-_NUM_PTS_DIM_SLAB_BOUNDS = 30
-

 class PolySlab(base.Planar):


I think we've talked already before about having configuration options exposed to the user. Would we want to expose config options for how the integration is performed on the level of the geometric objects themselves? (i.e. - each polyslab could potentially have it's own config or if not specified would use a default config). It isn't something that would go into this PR anyway, just curious if ultimately it will be useful for users to be able to have the ability to config these integrations for each object in the simulation or if eventually we would just let them specify some more global options for controlling integrations?

I would have only considered a global option, but I guess you can come up with a setup where it would be beneficial to control this explicitly. Seems like a pretty advanced use case though. But yeah it might make sense to have this machinery.

groberts-flex · 2025-05-07T18:45:05Z

Was in the structure.py file and noticed that in make_adjoint_monitors (

tidy3d/tidy3d/components/structure.py

Line 258 in 6b0c432

def make_adjoint_monitors(

), the monitor size for polyslabs along the extrusion axis is set to 0. I am assuming to get the full integration to work, we would need that to be the height of the polyslab.

of only at center for VJP calculation

yaugenst-flex requested a review from groberts-flex April 28, 2025 15:26

yaugenst-flex self-assigned this Apr 28, 2025

yaugenst-flex force-pushed the yaugenst-flex/polyslab-grads branch from f37734b to 00b99de Compare April 28, 2025 17:24

yaugenst-flex requested a review from tylerflex April 28, 2025 17:24

yaugenst-flex marked this pull request as ready for review April 28, 2025 17:25

yaugenst-flex force-pushed the yaugenst-flex/polyslab-grads branch 2 times, most recently from d73a485 to d188b31 Compare April 28, 2025 17:42

tylerflex reviewed Apr 28, 2025

View reviewed changes

groberts-flex reviewed Apr 29, 2025

View reviewed changes

tidy3d/components/geometry/polyslab.py Show resolved Hide resolved

groberts-flex reviewed Apr 29, 2025

View reviewed changes

tidy3d/components/geometry/polyslab.py Outdated Show resolved Hide resolved

groberts-flex reviewed Apr 29, 2025

View reviewed changes

yaugenst-flex mentioned this pull request Apr 30, 2025

feat: autograd support for rotated Box #2362

Open

groberts-flex reviewed Apr 30, 2025

View reviewed changes

tidy3d/components/autograd/derivative_utils.py Show resolved Hide resolved

groberts-flex reviewed Apr 30, 2025

View reviewed changes

yaugenst-flex force-pushed the yaugenst-flex/polyslab-grads branch 2 times, most recently from 332d8f4 to c0ab734 Compare May 2, 2025 07:25

fix[autograd]: sample fields along slab height and polygon edges instead

95c1619

of only at center for VJP calculation

yaugenst-flex force-pushed the yaugenst-flex/polyslab-grads branch from c0ab734 to 95c1619 Compare May 9, 2025 19:04

fix polyslab monitor z and add interval_space for adjoint monitors

84753f5

yaugenst-flex added the 2.9 will go into version 2.9.* label May 9, 2025

yaugenst-flex added 2 commits May 10, 2025 20:13

fix tests

eef5d03

separate interval_space for medium derivatives

633b247

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix[autograd]: sample fields along slab height and polygon edges #2418

fix[autograd]: sample fields along slab height and polygon edges #2418

yaugenst-flex commented Apr 28, 2025 •

edited

Loading

yaugenst-flex commented Apr 28, 2025

tylerflex left a comment

tylerflex Apr 28, 2025

tylerflex Apr 28, 2025

yaugenst-flex Apr 29, 2025

yaugenst-flex Apr 29, 2025

tylerflex Apr 29, 2025

yaugenst-flex Apr 29, 2025

tylerflex Apr 28, 2025

yaugenst-flex Apr 29, 2025

tylerflex Apr 29, 2025

yaugenst-flex Apr 29, 2025

tylerflex Apr 29, 2025

groberts-flex Apr 29, 2025

yaugenst-flex Apr 30, 2025 •

edited

Loading

groberts-flex Apr 30, 2025

tylerflex Apr 28, 2025

yaugenst-flex Apr 29, 2025

tylerflex Apr 28, 2025

yaugenst-flex Apr 29, 2025

groberts-flex Apr 29, 2025

groberts-flex left a comment

groberts-flex left a comment

groberts-flex Apr 30, 2025

yaugenst-flex May 2, 2025

groberts-flex commented May 7, 2025

	freq_max = float(max(grad_data_eps.eps_xx.f))
	eps_in = self.medium.eps_model(frequency=freq_max)
	ref_ind = np.sqrt(np.max(np.real(eps_in)))
	ref_ind = max([1.0, abs(ref_ind)])
	wvl_free_space = C_0 / freq_max
	wvl_mat = wvl_free_space / ref_ind

fix[autograd]: sample fields along slab height and polygon edges #2418

Are you sure you want to change the base?

fix[autograd]: sample fields along slab height and polygon edges #2418

Conversation

yaugenst-flex commented Apr 28, 2025 • edited Loading

yaugenst-flex commented Apr 28, 2025

tylerflex left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yaugenst-flex Apr 30, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

groberts-flex left a comment

Choose a reason for hiding this comment

groberts-flex left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

groberts-flex commented May 7, 2025

yaugenst-flex commented Apr 28, 2025 •

edited

Loading

yaugenst-flex Apr 30, 2025 •

edited

Loading