NanoComp / meep

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

Enhancement Request - Index Monitor in Adjoint Solver #2234

Open Andeloth opened 2 years ago

Andeloth commented 2 years ago

In the adjoint sub-package, it would be nice to have a monitor in the objective functions, or some other way to pass in the refractive index/design parameter into an objective function.

Example: I want to multiply a quantity by the design parameter array to make the FOM function care about the places where the topology has a high index material and I want it to ignore places where it has a low index material. Currently I have to reference the OptimizationProblem object and retrieve the stored design region's weights. However, I'm not sure whether this method is properly auto -differentiable for the gradient calculation.

stevengj commented 2 years ago

The stored design region weights are the same as the input parameters, and hence are differentiable. Unfortunately, they are not the same as the material parameters if you use filtering and projection.

smartalecH commented 2 years ago

As @stevengj suggests, this is rather straightforward out of the box (although one could always make it easier for the user...).

As described in Ch. 2 of my dissertation, the complete (linear) adjoint gradient is defined by

$$ \frac{df}{d\textbf{u}} = \frac{\partial f}{\partial \textbf{u}} + \lambda^T\left(\frac{\partial A}{\partial \textbf{u}}x + \frac{\partial b}{\partial \textbf{u}}\right) $$

where $f$ is your FOM as a function of the $E$ fields, $H$ fields, or even the design variables themselves (i.e. a function of the permittivity) such that

$$ f(\textbf{u},E,H) = \dots $$

( $\textbf{u}$ are the design variables, $A$ is the maxwell operator, $b$ are the defined current sources, $x$ are the fields from the forward solve and $\lambda$ are the fields from the adjoint solve).

Normally, meep assumes that $\frac{\partial b}{\partial \textbf{u}}=0$ and $\frac{\partial f}{\partial \textbf{u}} =0$. In other words, the gradient that meep returns is just the $\lambda^T\frac{\partial A}{\partial \textbf{u}}x$ term.

In your case, to compute the total gradient, you need simply compute $\frac{\partial f}{\partial \textbf{u}}$ yourself (or by using Autograd, Jax, PyTorch, etc.) and add it to the gradient Meep computes.

But be warned, as @stevengj said, you need to fully backpropagate through your mapping function, just like you do with Meep's gradient. You must backpropagate each term before doing the summation (unless your mapping function is linear... then the order shouldn't matter).

So again, it's really not too hard to do what you manually. But if you (or someone else) wanted to create a unified API that did this automatically, that would be great! (if you use the new subpixel smoothing features with the built in tanh functions, then you really want to automate this and create a routine in the C++ code that does the backprop for you... but for now I would just not use the subpixel smoothing routine as it's rather nascent still)

Andeloth commented 2 years ago

That explanation with the math written out helps a lot, thanks Alec :D. One day I'd like to make enhancements to the code base myself, but I feel like at the moment my understanding of the C++ backend is too weak to be able to contribute quality code 😅. But I hope to get there eventually!