🗻 autograd tidy3d (tidygrad?)

[tylerflex]

basic framework for autogradding tidy3d.

1. Adds autograd.ArrayBox type annotation in traced fields. For now just Medium.permittivity.
2. Adds json_encoder for this type, to write it to json as an untraced float.
3. Adds a autograd.extend.primitive version of web.run() that tells autograd how to connect the traced-fields to the xarray values in SimulationData.data through an adjoint simulation. Only works for AmpsDataArray. Only uses emulated run.

Todo Done

• Runs with emulated test
• Runs with web API
• Runs with extraneous monitors, structures
• Grad roughly matches numerical.
• Do all existing tidy3d tests pass with the changes so far?
• Test data array selection, interpolation in objective function.
• Where will all of the code live?
• Put floating helper methods as class methods.
• Auto handle traced vs untraced simulation fields.
• Auto handle traced vs untraced data fields.
• Handle integration with box bounds.
• Implement PIC structure fields
  • Medium
  • CustomMedium
  • PolySlab
  • Box
• Implement PIC adjoint sources
  • DiffractionMonitor
  • ModeMonitor
• run_batch support
• Replace existing web API
• Autograd dependency

Todo Before 2.7.0

• Invdes
  • penalty and transformation classes (ya) <-
  • make sure we like the final API
• Autograd plugin at tidy3d/plugins/autograd
  • Where do we import autograd from?
  • Functional autograd utilities (ya)
• Polishing
  • Warn or disallow autograd usage for some contains (eg. num structures).
  • yield instead of return dict in run_async and invdes
  • why is the bend notebook slow?
• Testing
  • Test speed of many structures. [think 400 max is a safe limit (2 min)]
  • Test validators and edge cases.
  • Test numerical accuracy.
• Reviews
  • Only use autograd.numpy where we really need it.
  • Thorough code reviews.
• Documentation
  • Update all compatible notebooks.
  • Write readme, including limitations and supported features.
  • Deprecate adjoint / jax stuff.
  • Document limitations of autograd.

Todo for 2.8.0

• Multi-frequency -> new notebook
• Web API
  • Server-side processing.
  • Selective ignoring of non-traced batch elements (if necessary)
• Feature support
  • GeometryGroup -> metalens
  • ComplexPolyslab
  • AnisotropicMedium
  • FieldMonitor -> metalens
  • .flux
  • get_intensity()
  • get_poynting()
• Invdes
  • Optimizer extensions
    • scipy / LBFGS
    • nlopt
    • pytorch
    • optax
  • Postprocess function definitions
• Jax, PyTorch, tensorflow converters.

Todo eventually

• DataArray improvements
  • differentiable interpolation with autograd ArrayBox dtypes.
  • Be able to return just the single element DataArray, right now need to grab values for it to register.
• Invdes
  • GUI integration
  • web API support
• Feature support
  • utility to look at gradient at some location? gradient monitor?

[tylerflex]

note: default behavior of copy changed so now you CAN set deep=False (although True is default).

[tylerflex]

I will need to be able to pass the parent tasks to run() for the server-side implementation.