Right now, doing eg jax.jacfwd(lambda eq: eq.solve()) doesn't work, and this makes interfacing with lots of other codes a lot more annoying. If we want to do optimization in desc we have workarounds, but if other codes just want the derivative of an equilibrium then there's a ton of extra glue code that needs to be written (and it's gotten a lot more complicated since the last time I had to do something like this).
It would be super useful if we got this to work correctly. The idea is that the input would be an equilibrium with whatever boundary/profiles you want, and the output of jvp would be another equilibrium where eq.R_lmn etc is really the tangent at F=0. This would make composing it with other jax functions super easy and not rely on knowing internals of desc in order to use our derivative information.
Right now, doing eg
jax.jacfwd(lambda eq: eq.solve())doesn't work, and this makes interfacing with lots of other codes a lot more annoying. If we want to do optimization in desc we have workarounds, but if other codes just want the derivative of an equilibrium then there's a ton of extra glue code that needs to be written (and it's gotten a lot more complicated since the last time I had to do something like this).It would be super useful if we got this to work correctly. The idea is that the input would be an equilibrium with whatever boundary/profiles you want, and the output of
jvpwould be another equilibrium whereeq.R_lmnetc is really the tangent atF=0. This would make composing it with other jax functions super easy and not rely on knowing internals of desc in order to use our derivative information.