Skip to content

Conversation

@groberts-flex
Copy link
Contributor

This PR updates and fixes the autograd level set notebook so that it matches the optimization happening in the adjoint version (with Jax). The major changes were to (1) remove references to the Jax backend and replace with autograd, (2) fix parameters that were set incorrectly, (3) match function calls more correctly between Jax and autograd (i.e. - the np.gradient with a certain mesh size), (4) and fix an underlying bug in the tidy3d codebase to get gradient magnitudes correct in the CustomMedium class (note that an update to the server python code needs to be in effect for this code to run as it does locally).

Copy link
Contributor

@yaugenst-flex yaugenst-flex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @groberts-flex thanks this looks great!

Just very minor comments:

  • There are some debugging comments I think left in the notebook (# does optax convert it to a jax array?)
  • from tidy3d.plugins.autograd.differential_operators import value_and_grad should work as a top level import so from tidy3d.plugins.autograd import value_and_grad would be preferred
  • While we're at it, could you run ruff format on the notebook? We should start doing this for every notebook.

@groberts-flex groberts-flex force-pushed the groberts-flex/fix_level_set_nb branch from face266 to 515293e Compare January 17, 2025 16:38
@groberts-flex
Copy link
Contributor Author

great, thank you for these comments! think I got all the debugging comments out of there now, fixed the import, and ran ruff format on the notebook file!

@yaugenst-flex yaugenst-flex self-requested a review January 17, 2025 19:13
Copy link
Contributor

@yaugenst-flex yaugenst-flex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

@groberts-flex groberts-flex merged commit 30db01c into develop Jan 20, 2025
@groberts-flex groberts-flex deleted the groberts-flex/fix_level_set_nb branch January 20, 2025 16:53
@yaugenst-flex yaugenst-flex linked an issue Apr 16, 2025 that may be closed by this pull request
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Fix level set notebook for autograd

3 participants