static jax structures & dispersive adjoint

[tylerflex]

Support for dispersive materials in adjoint (with shape / boundary optimization).
Adds two new adjoint structure classes:

• JaxStructureStaticMedium combines a jax geometry with a tidy3d medium.
• JaxStructureStaticGeometry combines a tidy3d geometry with a jax medium.

The former can be used to pair a differentiable geometry (JaxPolySlab) with any medium, eg a dispersive medium or something from the material library (like a metal).

This PR mainly generalizes the JaxGeometry internals to an abstract class and then splits into the three JaxStructure, JaxStructureStaticMedium, and JaxStructureStaticGeometryto handle the field definitions and to do different adjoint handling depending on their _differentiable_fields class attribute.

Among another static geometry structure, the tests have a JaxStructureStaticMedium containing a JaxBox and silver from our material library. So I think we can begin using this for eg metallic antenna design.

[momchil-flex]

Awesome. Somewhat related, and also related to our discussion yesterday, we should (or do we already?) validate that JaxSimulations do not have nonlinear media in either their static structures or their JaxStructureStaticMedium ones.

[momchil-flex]

Did you forget to uncomment the silver?

[tylerflex]

Yea good point. I can either do that in this PR or a separate one but would be a good idea. We could also consider just warning because the gradients might be approximately correct in some contexts?

[momchil-flex]

I guess warning could be fine too yeah.

[e-g-melo]

Hi @tylerflex!

I tried the JaxStructureStaticMedium and JaxStructureStaticGeometry within the adjoint examples 5 and 3, respectively, and both finished with an error (16:54:02 -03 ERROR: Error running task!).

In the JaxStructureStaticMedium class, the error happened when inserting both a non-dispersive or a dispersive medium (['Si3N4']['Philipp1973Sellmeier']): task id

In the JaxStructureStaticGeometry test, I've replaced the JaxBox by a Box object, and the same error has happened: task id.

[tylerflex]

This is probably because we'd need to build a solver to use this front end commit. Can you try with a run_local instead of run?

[tylerflex]

Note this PR adds a new field to the Jax info datastructure to remember what kind of input structures were used. So we need to make sure this is handled in the backend too. It might not even work with run local. @momchil-flex if it doesn't maybe could you build a solver using this commit so i can test?

[e-g-melo]

Ok.. got it!

[momchil-flex]

I've built a release and associated it to frontend version 2.5.0, modify version.py to use.

[tylerflex]

Thanks!

[e-g-melo]

Hi, @tylerflex and @momchil-flex!

It is working now! The gradients were calculated with no errors in both examples.

Thanks for this feature Tyler! Taking into account dispersive materials in the simulation adds another level of interest to the adjoint.

[tylerflex]

Great! Maybe we can update those examples to use the new structure and materials from your test. Could you open a PR in the docs repo? (Assuming everything else worked the same too, if there were some complications like poor performance, we can leave them out for now). Also at some point it might be good to do a metallic inverse design example. Maybe an antenna. I will get to field projection support hopefully in the coming weeks

[e-g-melo]

Sure, I will run the complete optimizations and update them in docs.

[tylerflex]

Thanks Emerson. @momchil-flex ill do nonlinear warnings in a separate PR