NanoComp / meep

free finite-difference time-domain (FDTD) software for electromagnetic simulations
GNU General Public License v2.0
1.16k stars 596 forks source link

First-order subpixel smoothing for density-based TO #2741

Closed smartalecH closed 5 months ago

smartalecH commented 6 months ago

Here we implement first-order subpixel smoothing for density-based TO. This approach allows us to treat the density formulation as a level set, such that the user can now continuously increase β→∞.

This approach is 100% in python and leverages autograd to do all of the backpropagation. It's very straightforward, but currently only works for 2D degrees of freedom. Adding capability for 1D and 3D is trivial -- we just need to ensure the filters work in those dimension, and we need to add in the right fill factor kernel (analytically derived by assuming the smoothing kernel is a sphere). This is a "simple version" of what's implemented in #1951

Here's an example:

image

While the approach works well, it is still sensitive to numerical roundoff errors. Autograd doesn't really have any tooling in place to track where things are breaking down (particularly in the backward pass). So we'll have to get creative. Here's a plot that shows the norm of the gradient which breaks down due to numerical error with increasing β:

image

Also, still needs a test.

codecov-commenter commented 6 months ago

Codecov Report

Attention: 21 lines in your changes are missing coverage. Please review.

Comparison is base (295fa63) 74.06% compared to head (7fb6c6b) 73.50%. Report is 8 commits behind head on master.

:exclamation: Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## master #2741 +/- ## ========================================== - Coverage 74.06% 73.50% -0.56% ========================================== Files 18 18 Lines 5395 5549 +154 ========================================== + Hits 3996 4079 +83 - Misses 1399 1470 +71 ``` | [Files](https://app.codecov.io/gh/NanoComp/meep/pull/2741?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | | |---|---|---| | [python/adjoint/filters.py](https://app.codecov.io/gh/NanoComp/meep/pull/2741?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-cHl0aG9uL2Fkam9pbnQvZmlsdGVycy5weQ==) | `69.76% <4.54%> (-7.44%)` | :arrow_down: | ... and [1 file with indirect coverage changes](https://app.codecov.io/gh/NanoComp/meep/pull/2741/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)
oskooi commented 6 months ago

Also, still needs a test.

I suppose a test for this feature would be based on @stevengj's proposal in #1854 (comment) but this time demonstrating first-order convergence.

smartalecH commented 6 months ago

but this time demonstrating first-order convergence.

I don't think we really need to test the first-order convergence, since that's not really the feature here. Rather we should probably check the accuracy of the gradient for a few finite betas, and when beta = np.inf. Maybe a few tests that check for correctness, type sanity too, etc.

The real novelty with this PR is the ability to differentiate with any beta. So that's probably the test we should implement.

stevengj commented 6 months ago

So we'll have to get creative. Here's a plot that shows the norm of the gradient which breaks down due to numerical error with increasing β:

What's the source of this error?

Why does the gradient seem to be blowing up as β increases? Isn't the point of smoothing to get a finite nonzero gradient?

smartalecH commented 6 months ago

Isn't the point of smoothing to get a finite nonzero gradient?

Right, exactly. I would have expected it to be monotonic too, but maybe there's no reason for it to be monotonic (but it should smoothly converge at least to something finite).

What's the source of this error?

So during the backward pass, autograd's numpy throws some overflow warnings for a few different functions. But again, the tooling and logging capabilities for autograd are really nonexistent, so debugging will be pretty manual.

But if I had to guess, here are some areas where things could blow up:

(In the docstring, I have an example which can be used to recreate the gradient norm plot above -- just sweep beta and use the autograd grad function).

smartalecH commented 5 months ago

Alrighty, I think I resolved all the above issues. Just had to implement a few "double where" tricks to sanitize the backprop. (also had a bug in the effective material part of the algorithm). The smoothed materials themselves look much better:

image

And, if we do a convergence check as β→∞, we see that the norm of the gradient converges smoothly to a finite, non-zero value:

image

I've also checked the pathological case where we have a uniform design field (such that the spatial gradient is zero). Normally, this would create a divide by zero error when computing the distance to the interface, and further complicate the backward pass. But I've sanitized things enough that everything looks good.

This should be a drop-in replacement for the tanh_projection() function, with the requirement of an additional resolution argument (to compute the spatial gradient).

Any ideas for a test?