Fix all examples for PR #29#35
Conversation
…loss wrappers, solver kwarg in Tracing, override setters for Curves gamma
rogeriojorge
left a comment
rogeriojorge
left a comment
There was a problem hiding this comment.
Done with review of this PR.
|
|
||
| nfp, mpol, ntor: static | ||
|
|
||
| Backward-compat: rc may also be a filename (string) or a Vmec-like object, |
There was a problem hiding this comment.
@eduardolneto can you review this init? Do we need such backwards compatibility?
…or counts, verify solver differentiability
|
Hi @EstevaoMGomes, @eduardolneto, please review this today or tomorrow. If it looks good (check the new examples locally), then please approve the review or add your comments. |
EstevaoMGomes
left a comment
EstevaoMGomes
left a comment
There was a problem hiding this comment.
There are some changes needed. I'm not sure if in this or another PR, but I can help with understanding PyTrees and tracers. Nevertheless, the changes are appreciated! Thanks for fixing the examples :)
I'll try to join the next meeting to give you some input on losses.py and how we actually should differentiate over PyTrees. When I started coding ESSOS, I also did not know exactly how this worked, so the repo may still be confusing when mixing the 2 approaches.
| """Curvature penalty as a function of the optimization vector x. | ||
|
|
||
| Unlike loss_coil_curvature, which takes a Coils object, this version takes | ||
| the flat degrees-of-freedom vector x used by the optimizers. It rebuilds the |
There was a problem hiding this comment.
Does this work with taking the PyTree correspondent to a Coils class? Otherwise it is not compatible with losses.py
There was a problem hiding this comment.
The object-based losses (loss_coil_curvature(coils)) already work with custom_loss directly. So migrating the examples onto losses.py and dropping these flat-vector wrappers seems like the right follow-up.
There was a problem hiding this comment.
This could be tweeked with the changes of #30 , after this one is merged. Most of the exmaples and changes in objective funcions in there have the format @EstevaoMGomes is askign for.
|
There was an open pull request here #30 that has conflicts with some parts of this pull request. I think the changes in that pull request should be merged before as they were already fixing a lot of these issues including coil_perturbation, coils, surfaces and augmented_lagragian/losses iteraction. The exampels were also working . @EstevaoMGomes can you recheck that one first please. |
…hic init, revert loss_BdotN kwargs
eduardolneto
left a comment
eduardolneto
left a comment
There was a problem hiding this comment.
I left my comments, I am finsihed with the review of the PR.
4 participants
Got all 44 examples in PR #29 running end-to-end (was 3/44).
Verified locally: Fig 3 (gc_vs_fo), Fig 5 (poincare_plots), and the autodiff figure (gradients) all reproduce correctly.
Source-level changes (all additive / backward-compatible):
essos/coils.py: override setters for Curves gamma / gamma_dash / gamma_dashdash (unblocks coil perturbation with the immutable Curves model)
essos/surfaces.py: SurfaceRZFourier.init accepts filename string or Vmec object as first arg (existing explicit-arg calls unchanged)
essos/dynamics.py: optional solver= kwarg in Tracing (defaults preserve Dopri8); energy() now handles FullOrbit_Boris
essos/objective_functions.py: restored loss_coil_curvature_new, loss_coil_length_new, loss_particle_r_cross_max_constraint as flat-vector wrappers; loss_BdotN accepts surface= alternative to vmec=; fixed shape mismatch in loss_optimize_coils_for_particle_confinement
essos/augmented_lagrangian.py: 16 jax.jax.tree_util -> jax.tree_util typo fixes
essos/coil_perturbation.py: jnp.clip(eigvals, a_min=0) -> min=0 (NumPy 2.x API)
Example fixes:
Path bugs (10+ files): name -> file typos, "../examples/input_files" -> "../input_files", missing "../" when running from subdirs
API renames: Coils_from_json -> Coils.from_json (7 files), Coils_from_simsopt -> Coils.from_simsopt (2 files)
Tracing API alignment in poincare_plots.py (timesteps=count -> timestep=dt, tol_step_size -> atol/rtol)
gradients.py: hardcoded 8 processors -> 1 (was failing on shape mismatch)
Rewrote paper/fo_integrators.py and paper/gc_integrators.py for current keyword API + new solver= kwarg (Fig 4 unblocked)
.energy attribute -> .energy() method form
_new aliasing in adam / augmented-lagrangian examples
Optional deps the examples need: simsopt, h5py, plotly, booz_xform (latter needs libnetcdf-dev + gfortran on the system). Happy to add these as extras to pyproject.toml in a follow-up if useful.
Caveats: every example now starts and runs without crashing. End-to-end completion verified for the 3 paper figures above. The longer-running optimizers (LBFGS, augmented Lagrangian) should be confirmed on full runs.