add notebook for parallel adjoint

[marcorudolphflex]

This PR adds a new notebook, Autograd30ParallelAdjoint.ipynb, to introduce the parallel adjoint feature in Tidy3D with a simple waveguide-bend example.

The notebook is written as a user-facing tutorial and focuses on a single gradient evaluation rather than an optimization loop. It starts with a short explanation of the standard adjoint workflow, then shows when parallel adjoint can reduce runtime: if the adjoint setup does not depend on forward results, one or more canonical basis simulations can be launched in parallel with the forward simulation and postprocessed afterward. In favorable cases, this can approach a 2x speedup, although the observed gain also depends on fixed overheads such as upload, download, and local pre/postprocessing.

Still missing are a thumbnail image, notebook metadata and adding this to the index of autograd notebooks, this is just meant to be coarsely reviewed for now.

Which ideas do you have for a thumbnail image? Maybe a small diagram to explain the parallel adjoint idea with some timeline?

[github-actions]

Spell Check Report

Autograd30ParallelAdjoint.ipynb:

Cell 1, Line 7: 'gradient-relevant'
  > To use parallel adjoint efficiently, the simulation has to be set up such that it is already clear beforehand which adjoint simulations will be needed. In practice, this means using only supported and gradient-relevant monitor outputs, and restricting those outputs to the data that actually enters the objective.

---

Checked 1 notebook(s). Found spelling errors in 1 file(s).
Generated by GitHub Action run: https://github.com/flexcompute/tidy3d-notebooks/actions/runs/24405668643

[groberts-flex]

This is great @marcorudolphflex, thanks for putting this together! I left a few comments below:

1. "We do not want to focus on the bend parametrization itself in this notebook. The bend can be seen in the simulation view in the next section, and Autograd8WaveguideBend is the better reference if you want to study a more deliberate waveguide-bend parametrization."

there are a few references to not paying attention too closely to the bend parameterization. I think we could state that once up top and this one is a good one to use since it points to another waveguide bend example.

2. "Select the Mode of Interest"

This section might be able to be simplified. Is there a reason the target mode isn't showing up as mode_index=0 from the start?

3. this comment in objective function section in both cells: "# note that parallel adjoint is only supported with local gradients yet"

could be "# note that parallel adjoint is currently only supported with local gradients"

One additional question on this - do we warn currently if someone is using parallel_adjoint but not local_gradient?

4. In this block:

with config as cfg: # optional to ensure temporary changes
cfg.adjoint.local_gradient = (
True # note that parallel adjoint is only supported with local gradients yet
)

Does this mean we don't need to be setting local_gradient in the objective function?

5. The "runtime to compute gradients" plot could probably be removed since the speedup is shown nicely in the table above.

6. "If any unsupported monitor is present in the differentiable simulation, Tidy3D falls back to the sequential adjoint pipeline for that run. Not supported:"

The "Not supported:" part could be changed to something like "Unsupported monitors include:"

7. "This is useful when the intended direction is ambiguous, but it increases the mode-monitor part of the parallel adjoint work."

For this it might be useful to clarify what would make the intended direction ambiguous in the case of the assume_outgoing.

[marcorudolphflex]

This is great @marcorudolphflex, thanks for putting this together! I left a few comments below:

1. "We do not want to focus on the bend parametrization itself in this notebook. The bend can be seen in the simulation view in the next section, and Autograd8WaveguideBend is the better reference if you want to study a more deliberate waveguide-bend parametrization."

there are a few references to not paying attention too closely to the bend parameterization. I think we could state that once up top and this one is a good one to use since it points to another waveguide bend example.

2. "Select the Mode of Interest"

This section might be able to be simplified. Is there a reason the target mode isn't showing up as mode_index=0 from the start?

3. this comment in objective function section in both cells: "# note that parallel adjoint is only supported with local gradients yet"

could be "# note that parallel adjoint is currently only supported with local gradients"

One additional question on this - do we warn currently if someone is using parallel_adjoint but not local_gradient?

4. In this block:

with config as cfg: # optional to ensure temporary changes cfg.adjoint.local_gradient = ( True # note that parallel adjoint is only supported with local gradients yet )

Does this mean we don't need to be setting local_gradient in the objective function?

5. The "runtime to compute gradients" plot could probably be removed since the speedup is shown nicely in the table above.
6. "If any unsupported monitor is present in the differentiable simulation, Tidy3D falls back to the sequential adjoint pipeline for that run. Not supported:"

The "Not supported:" part could be changed to something like "Unsupported monitors include:"

7. "This is useful when the intended direction is ambiguous, but it increases the mode-monitor part of the parallel adjoint work."

For this it might be useful to clarify what would make the intended direction ambiguous in the case of the assume_outgoing.

Thanks for this review!

1. redundant toy explanation: True, I shortened that.
2. mode index 2: I just wanted to show a non-obvious case which may be hard/confusing for users to construct themselves as they have to specify mode_index=0 (not 2) in the monitor. Or is that case of mode_index=2 super unusual/weird?
3. local_gradient comment: changed that. Yes, we warn on this.
4. local_gradient redundancy: Yep, that was redundant. Removed the redundant setting in web.run.
5. bar chart: Thought it's a good eyecatcher for people who just scroll over the notebook. But should be already clear from the table.
6. "Unsupported"-change: changed that
7. ambigious direction: Tbh not sure if they are such cases - I guess your question implies that this is quite unusual. Just remove this comment?

[groberts-flex]

"mode index 2: I just wanted to show a non-obvious case which may be hard/confusing for users to construct themselves as they have to specify mode_index=0 (not 2) in the monitor. Or is that case of mode_index=2 super unusual/weird?"

I see what you're saying here! In this example, you're showing how to set things up so that only one parallel adjoint source is created because if the ModeSpec has multiple modes and then the objective function selects only one of them, then there will be extra parallel adjoint runs. I think it's a good thing to show and demonstrate, but could be accompanied with more explanation of why it's necessary. My guess is most users won't immediately know the idea of a canonical adjoint source and when you do or don't need to know the objective function to create adjoint sources. I would suggest expanding on the flowchart and initial explanation you have explaining the difference between parallel and sequential adjoint to try and give readers a better model of how to think about setting up their simulation for parallel adjoint. And then in this section, you could explain what would happen if they had a ModeSpec with 3 modes and selected one of them afterwards in sequential versus parallel.

"bar chart: Thought it's a good eyecatcher for people who just scroll over the notebook. But should be already clear from the table."

I'm good either way on it! It is eye catching, so totally happy to leave it in there.

"ambigious direction: Tbh not sure if they are such cases - I guess your question implies that this is quite unusual. Just remove this comment?"

I think it is worth having the comment, but maybe just a little more detail on what would cause ambiguity. My understanding is the direction is computed based on the position in the simulation with respect to the center, so you could just explain that part of the assume_outgoing policy explicitly.

[groberts-flex]

I like the changes and think things are really clear for users on how to best use parallel adjoint effectively!

looks great