Adjoint Solver: Optimization problem with an objective function dependent on separate simulations.

[rafael-fuente]

I'm currently working on an optimization project involving an objective function that is dependent on the DFT fields of completely independent simulations.

That is, denoting the DFT fields of a specific region of the simulation number $s$ by $\mathbf{E_s}$, I need the gradients with respect to the design region permittivity $\epsilon$ of an objective function $\mathcal{F}$ of the form:
$$\mathcal{F}(\mathbf{E_1} , \mathbf{E_2}, \mathbf{E_3} ... \mathbf{E_N} )$$

However, meep adjoint OptimizationProblem object only allows an objective function whose arguments are exclusive from the output of a single simulation object.

Is there a way to handle my problem without substantially changing the meep adjoint module?

One possibility is computing the gradients via the following expression:
$$\frac{\partial \mathcal{F}}{\partial \epsilon_i} = \sum_{s} \frac{\partial \mathcal{F} }{\partial \mathbf{E_s}} \frac{\partial \mathbf{E_s}}{\partial \epsilon_i} $$

Where the chain rule was used and $\frac{\partial \mathbf{E_s}}{\partial \epsilon_i}$ is the jacobian of the DFT fields of the simulation number $s$.
How could I compute that Jacobian by using meep adjoint module?

[ianwilliamson]

I'm currently working on an optimization project involving an objective function that is dependent on the DFT fields of completely independent simulations.

This supported by the Meep JAX wrapper. Essentially, you can create multiple MeepJaxWrapper objects wrapping separate simulations. You can then compose calls to these wrapper objects into any JAX loss function / computational graph that you want and then apply jax.value_and_grad(). The main limitation is that all of the involved simulations, although completely independent of each other, will be performed sequentially.

[rafael-fuente]

Thank you for your answer!

I have been testing Meep JAX wrapper.

Is there a way to make it accept a monitor mpa.FourierFields that returns an array of the DFT field evaluated at multiple points?

For example, when I use as mpa.MeepJaxWrapper monitors argument:
monitors = [mpa.FourierFields(sim,volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(x=0)), component = mp.Ez)] which return the DFT Ez component of the field at an single point, everything works fine.

But when I try to use mpa.MeepJaxWrapper monitors argument with a non-zero size Volume:
monitors = [mpa.FourierFields(sim,volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(x=2)), component = mp.Ez)] which should return an array of the DFT Ez component of the field evaluated along the volume specified; I'm getting an AssertionEror:

If I want to evaluate the DFT fields in multiple locations, I could, in principle, use multiple monitors. But that would be computationally expensive if the array is large. Is there a way to fix this problem?

[smartalecH]

This supported by the Meep JAX wrapper.

This is probably the easiest approach.

The main limitation is that all of the involved simulations, although completely independent of each other, will be performed sequentially.

Luckily you can indeed do this in parallel using the approach described below...

I'm currently working on an optimization project involving an objective function that is dependent on the DFT fields of completely independent simulations.

If everything is linear, you can cascade things together rather easily. The trick is to compute the gradient of each independent FOM, and then recombine them together like you suggest.

We do this quite often with the current infrastructure (no changes to the OptimizationProblem needed). If you have enough hardware, you can even run these in parallel (which is what we do). This is especially useful for robust optimization, where you are optimizing multiple design fields simultaneously.

We simply run my_group = mp.divide_parallel_processes(N) where N is the number of "groups" you want to create and my_group is the MPI group index of the current group. Once this starts, you can then create multiple simulations for different groups using something like..

if my_group == 0:
   geometry = [<first geometry>]
   sources = [<first source>]

elif my_group == 1:
   geometry = [<second geometry>]
   sources = [<second source>]

Then, you compute your gradients as expected and can broadcast all the gradients from each group to each other, such that you can now optimize over all design fields etc. using something like

f, df = opt([x])
all_grads = mp.merge_subgroup_data(df)

This is a bit intimidating at first, but it's a truly powerful form of parallelism (which we call "in-the-loop" parallelism in our paper) that dramatically simplifies multiobjective optimization.