Inverse Design Plugin

[tylerflex]

Notebook PR:
flexcompute/tidy3d-notebooks#94

To Do list as of Mar 26

• Finalize API
• Use API in tutorial notebook problem.
• Add all docstrings
• Add all pd.Field descriptions
• 100% test coverage.
• Nice to haves
  • import common postprocess function constructors.
  • try single postprocess fn with batch data for multi-objective
  • standardize to_simulation / data for design and results.
• Polish up notebook with better text and nice images.
• Get feedback. Rinse and repeat above if needed

[lucas-flexcompute]

I don't see anything that looks out of place here. The user interface is pretty great! Definitely easier to use than going trough adjoint directly.

[momchil-flex]

Overall looks great, should be very useful!

I see you've added the invdes.png image here and also in the notebooks PR, and it may not be used here?

One thing I noticed when I worked with the similar history handling that we have in our notebooks is that in some cases (esp. if I e.g. push things to large design regions like the 12x14um grating coupler I optimized) the memory needed to store the history can be quite large. This can make both the file large so it is inconvenient to handle (e.g. if you want to keep a bunch of optimization files they can take up a lot of space), and there's even a risk of running our of cpu memory completely since you store the whole history specifically for grad, params, and also the opt_state which itself is of a similar or larger size to params. Along those lines I feel like some more control over what exactly is stored, and some good default settings, may provide for a better experience?

For example, I probably always want to store the history of the objective_fn_val, penalty, and post_process_val. In most cases I probably just want the last params and opt_state so that I can examine my last design and continue an optimization. I probably almost never actually need to look at the grad. Now, I realize that using these as defaults (do not store history for params, grad, opt_state) limits functionality a bit (e.g. if I want to make a gif of the params evolution, or if I want to start from a state that's not the latest one). But if we allow the option to enable those, do you think it may be better?

[momchil-flex]

I always found this tricky previously when we wrote similar optimization history scripts. Here it seems like the params that you write do not match all the other values that you have written in history, i.e. the grad, objective function, etc. were computed for the params before the update, while here you write them to history after the update. You can certainly move history["params"].append(params) to be above with the other history updates but I think there was some reason to put it later, maybe related to what happens when you continue an optimization? I remember I thought about this once and it was not as straightforward as one would expect.

[tylerflex]

So what I did is when the optimizer is run from an empty state Optimizer.run(params), it will store the initial params and opt_state before running this loop.

tidy3d/tidy3d/plugins/invdes/optimizer.py

Line 76 in 1844968

return InverseDesignResult(design=self.design, opt_state=[opt_state], params=[params0])

So actually for num_steps steps, we end up having num_steps+1 values for params and opt_state in the history but only num_steps values for the other stuff.

It's not always clear whether people want the final state / parameters to be after the final update or what. But I decided to include it here just for completeness. And also so that if we continue a run it will not waste that gradient calculation.

[momchil-flex]

Yeah I think maybe not wasting the gradient computation is what was tripping me up before too. But if I understand correctly this can definitely be confusing now? First of all having history contain lists of different lengths is already not great. But probably the biggest problem is that if someone does result.get_last("objective_fn_val"), this value will not correspond to the simulation I get from result.get_sim(), if I then do a manual run on this and compute the objective. This seems like something that would happen often.

[tylerflex]

I see you've added the invdes.png image here and also in the notebooks PR, and it may not be used here?

One is used for the notebook itself, and the other is used in the README.md for invdes, but I suppose I could try making them point to the same place.

One thing I noticed when I worked with the similar history handling that we have in our notebooks is that in some cases (esp. if I e.g. push things to large design regions like the 12x14um grating coupler I optimized) the memory needed to store the history can be quite large. This can make both the file large so it is inconvenient to handle (e.g. if you want to keep a bunch of optimization files they can take up a lot of space), and there's even a risk of running our of cpu memory completely since you store the whole history specifically for grad, params, and also the opt_state which itself is of a similar or larger size to params.

A few things:

1. The history saving is optional (in case that wasn't apparent). so I suppose users could turn this off if they don't want it.
2. Without the opt_state or params, the history will not be able to be continued, so I feel it should probably be included, but I was considering adding a few options to further customize the optimization loop, which I'll explain below.
3. We still restrict the number of parameters in the custom medium to 500,000 I think, so my feeling is it shouldn't be too big of an issue? Correct me if I'm wrong though. Maybe the issue iwas more storing he whole SimulationData object? which we did in some notebooks but I decided not to store here.

Along those lines I feel like some more control over what exactly is stored, and some good default settings, may provide for a better experience?

I think changing the fields in the InverseDesignResults will lead to a bunch of annoying edge cases internally, but I'm considering the following two things:

1. add a callback function option to the Optimizer, so the user could in principle do something like

history = dict(params=[], ...)
def callback(current_results, step_index: int) -> None:
    history['params'].append(params)

optimizer = Optimizer(..., callback=callback)
optimizer.run(params0)

history['params'] 

2. Allowing the user to specify a post_process_fn that returns aux_data, so they could even store the eg. SimulationData as

def postprocess(sim_data, ...):
    ...
    return objective, dict(data=sim_data)

history = dict(data=[])
def callback(..., aux_data):
    history['data'].append(aux_data['data'])

For example, I probably always want to store the history of the objective_fn_val, penalty, and post_process_val. In most cases I probably just want the last params and opt_state so that I can examine my last design and continue an optimization. I probably almost never actually need to look at the grad. Now, I realize that using these as defaults (do not store history for params, grad, opt_state) limits functionality a bit (e.g. if I want to make a gif of the params evolution, or if I want to start from a state that's not the latest one). But if we allow the option to enable those, do you think it may be better?

I think the callback / aux_data could be good enough to support this, but the learning curve is a bit high. What do you think?

[momchil-flex]

One thing I noticed when I worked with the similar history handling that we have in our notebooks is that in some cases (esp. if I e.g. push things to large design regions like the 12x14um grating coupler I optimized) the memory needed to store the history can be quite large. This can make both the file large so it is inconvenient to handle (e.g. if you want to keep a bunch of optimization files they can take up a lot of space), and there's even a risk of running our of cpu memory completely since you store the whole history specifically for grad, params, and also the opt_state which itself is of a similar or larger size to params.

A few things:

1. The history saving is optional (in case that wasn't apparent). so I suppose users could turn this off if they don't want it.

You mean the history saving to file? That's true, but it's nice to have. And apart from that it's always stored in memory currently.

2. Without the opt_state or params, the history will not be able to be continued, so I feel it should probably be included, but I was considering adding a few options to further customize the optimization loop, which I'll explain below.

Yeah but you only need the last value, not the whole history. My suggestion was to store only the last value by default.

3. We still restrict the number of parameters in the custom medium to 500,000 I think, so my feeling is it shouldn't be too big of an issue? Correct me if I'm wrong though. Maybe the issue iwas more storing he whole SimulationData object? which we did in some notebooks but I decided not to store here.

Storing SimulationData could also be an issue, but that's not what I was thinking. But yeah let's try to estimate this. Let's say you have 500k params and those, and you have those in params, grad, and I think my observation was that opt_state contains at least another 2x. So say you have 2M double-precision floats at every iteration, and you do 100 iterations. That's 1.6GB. Not terrible, but also not that great.

Along those lines I feel like some more control over what exactly is stored, and some good default settings, may provide for a better experience?

I think changing the fields in the InverseDesignResults will lead to a bunch of annoying edge cases internally, but I'm considering the following two things:

1. add a callback function option to the Optimizer, so the user could in principle do something like

history = dict(params=[], ...)
def callback(current_results, step_index: int) -> None:
    history['params'].append(params)

optimizer = Optimizer(..., callback=callback)
optimizer.run(params0)

history['params'] 

2. Allowing the user to specify a post_process_fn that returns aux_data, so they could even store the eg. SimulationData as

def postprocess(sim_data, ...):
    ...
    return objective, dict(data=sim_data)

history = dict(data=[])
def callback(..., aux_data):
    history['data'].append(aux_data['data'])

For example, I probably always want to store the history of the objective_fn_val, penalty, and post_process_val. In most cases I probably just want the last params and opt_state so that I can examine my last design and continue an optimization. I probably almost never actually need to look at the grad. Now, I realize that using these as defaults (do not store history for params, grad, opt_state) limits functionality a bit (e.g. if I want to make a gif of the params evolution, or if I want to start from a state that's not the latest one). But if we allow the option to enable those, do you think it may be better?

I think the callback / aux_data could be good enough to support this, but the learning curve is a bit high. What do you think?

This does get more complicated. Maybe for starters we leave as is. But I was more driving towards something like an e.g. store_complete_history flat set to False by default, in which case store only the history of the objective function and penalty, while only the last values for grad, params, and opt_state. Or it could even be True by default. Not sure how much internal complexity this adds though.

[tylerflex]

@momchil-flex sounds good for all points, I'll play around with a store_full_history : bool = True field and see how it works internally.

I think more complicated callbacks can be added later, including for more customized history storage, display, and parameter scheduling.

[tylerflex]

@momchil-flex I implemented something in this commit. Could you take a look when convenient? thanks.