Obtaining the phase from field monitor for inverse design problem

[Nikolaos-Matthaiakakis]

Hello, I am a little bit confused about how i can extract some data from a field monitor in a way that is differentiable for autograd for an inverse design issue.

In more detail I face the following problem. I use

# Loss function
def loss_fn(params: anp.ndarray, beta: float) -> tuple[float, dict]:
    """Loss function for the design, the difference in intensity + the feature size penalty."""

    freqs_inv  = np.linspace(3, 6, 3) * THz  # frequency range of interest 
    # construct and run the simulation
    sim_inv = make_sim(params, beta=beta,freqs=freqs_inv)
    sim_data_inv = web.run(sim_inv, task_name="loss_calc", verbose=False)

    # calculate the error
    error=error_fn(sim_data_inv )

    # grab the respective and total losses
    return error, sim_data_inv.to_static()

# construct a function of `params` and `beta` that returns the loss value, gradient, and the aux_data
loss_fn_val_grad = value_and_grad(loss_fn, has_aux=True)

# call this on our initial parameters
(val, grad), sim_data = loss_fn_val_grad(params0, beta=beta0)

which calls an error function that tries to obtain the phase of the fields from a monitor and compare them to a desired broadband phase profile. The issue is I am not sure how I can obtain this data in a way suitable for autograd as

def error_fn(sim_data):
    # Extract raw arrays from xarray
    Ey_Tr = sim_data["Tf"].Ey.data
    Ey_Tr_ref = ref_results["Tf"].Ey.data
    
    # Average over x and y dimensions
    Ey_Tr_avg = anp.mean(Ey_Tr[:, :, 0, :], axis=(0, 1))
    Ey_Tr_avg_ref = anp.mean(Ey_Tr_ref[:, :, 0, :], axis=(0, 1))

    # Phase in radians
    phase_Ey_Tr = anp.angle(Ey_Tr_avg)
    phase_Ey_ref = anp.angle(Ey_Tr_avg_ref)

does not seem to properly provide the data in the correct form for this problem.

[momchil-flex]

Could you clarify the problem: does not seem to properly provide the data in the correct form for this problem. do you mean that the data is not what you expect, that the gradient does not seem right, or that the gradient computation errors?

[Nikolaos-Matthaiakakis]

Hello, I apologize for the very unclear first comment. At first I thought I was retrieving the data in a wrong way that caused a variety of errors with autograd. In reality, I was making an error on how I handled the data later in my code and the result ended up being non differentiable. I believe I have fixed this issue now but I have some other questions. For example for a broadband optimization (like for achieving a suitable phase profile for an achromatic metalens) does the adjoint solver run an adjoint simulation for each of the wavelengths captured by the monitor? Are the results then weighted together somehow? I am a little bit confused on how the adjoint method is implemented here for such simulations.

Thank you very much for your time.

[momchil-flex]

Good that you got it to work! Your question is pretty good and the details are actually complicated. In the worst case, we do need one adjoint simulation per monitor frequency. In some cases even with broadband, a single adjoint simulation can be used with some tricks. However I think (at least currently) if you are differentiating a FieldMonitor, it would be one adjoint sim per frequency.. so use as few as possible. @groberts-flex can confirm if this is true.

[momchil-flex]

No, in any case this is handled under the hood. Your objective function should be scalar-valued. But if it depends on multiple frequency outputs, this is detected automatically, and the adjoint simulations will be created as needed. It's just that in some cases (e.g. a single ModeMonitor in the objective function) we can even use a single adjoint simulation, but again, this is detected automatically.

[Nikolaos-Matthaiakakis]

I understand, thank you very much.

[groberts-flex]

Just wanted to confirm that the above is true! When using a FieldMonitor, there will be one adjoint simulation per frequency that is used in your objective function. You can have more frequencies in your monitor, but the ones that get used in your objective will each be given their own adjoint simulation. Depending on how you weight them together to achieve a scalar value in the end, this will determine the relative weights of all those adjoint simulations. As was pointed out, this all happens automatically for you. The main thing to take into account is to limit the number of frequencies that enter the objective when possible to reduce the number of simulations you need when using FieldMonitors.

[Nikolaos-Matthaiakakis]

Thank you for the verification, I saw that there is a limit of 10 frequencies by default. Currently I do not need to use that many, but if I needed would there be a way to overcome this limit?

[groberts-flex]

Yes, you can use the argument max_num_adjoint_per_fwd to your run or run_async call to overcome the limit.

[Nikolaos-Matthaiakakis]

Hello, sorry to revive this thread but I am having a related issue, probably due to my lack of experience with autograd. I have tried several types of loss functions to achieve a target phase (mse, cosine loss, etc) but I seem to be facing issues such as, 0 gradients, value getting stuck at the worst value, value getting stuck in plateaus, etc.

Phase is wrapped so this can cause issues and I probably haven't found an ideal way to overcome this, unwrapping seems to cause issues also. I was wondering if you have could offer some advice on how to handle loss functions that have phase instead of intensity as a target if you have any relative experience. Thank you very much in advance.

[Nikolaos-Matthaiakakis]

i checked now and print(td.version) gives 2.10.0 .

[groberts-flex]

Thanks for sending that! We had a fix in 2.10.1 to improve accuracy of CustomMedium gradients so if possible it would be good to try things at that version too.

I checked out the notebook and made a few modifications and tested on our current state. Feel free to give these suggestions a try in your version or 2.10.1:

1. the initial design seems to be quite close to the desired phase at step 1. I changed the phase target to np.pi / 4 so there's a little more room to optimize the loss
2. I would try some smaller learning rates. I believe the first step jumping really high is likely from being near the minimum and then taking a large step. I took the learning rate down to 0.05.
3. For the beta parameter, I would recommend starting it even smaller and then ramping it up more slowly. Each time you change that value, it will abruptly change your permittivity distribution. Usually it takes a few iterations to recover from that at the same beta value. I tried just fixing the beta value to beta0=0.1 to make sure the optimization was working, but once you feel confident I would try and increase it every 10 iterations or so.

With those changes, I get a loss curve that looks like the one below. Let me know if you have any questions or are having trouble getting any of that working!

[Nikolaos-Matthaiakakis]

Thank you, this looks amazing. I upgraded to 2.10.2 just now. I will give this a try tomorrow.

If it is not too much effort, regarding performing a similar optimization for a material like gold or Si in the visible, what would be the best way to set the material?

[groberts-flex]

Happy to hear that looks good and curious to hear about what you find out in our test.

I think this may be a good example notebook to check out for those other types of material: https://www.flexcompute.com/tidy3d/examples/notebooks/Autograd15Antenna/

You can set up custom medium type media that are dispersive or lossy. You can set up a custom pole residue or a custom Drude model for example and compute gradients with respect to their parameters. Let me know if that sounds right for your application and if you run into any issues with it!

[Nikolaos-Matthaiakakis]

Thank you, I believe that for now I am covered with the answers