Major refactoring for JAX-style classes.#29
Open
EstevaoMGomes wants to merge 77 commits into
Open
Conversation
- Implemented `fo_integrators.py` for full orbit tracing with various methods and parameters. - Implemented `gc_integrators.py` for guiding center dynamics with adaptative and constant step sizes. - Enhanced `Tracing` class in `dynamics.py` to support multiple methods and step sizes.
…d adjust num_steps based on dt
… class for improved step size handling
…for adaptive step size
…ters for performance
…ance plots, and improve layout for better visualization
Collaborator
Author
|
Tests need fixing; |
There was a problem hiding this comment.
Pull request overview
This PR implements a major refactoring to make coils and surfaces proper JAX PyTrees, enabling automatic differentiation. It introduces a new loss wrapper system for gradient-based optimization and adds comprehensive analysis and validation code comparing ESSOS with SIMSOPT.
Key changes:
- Refactored
Coils,Curves,SurfaceRZFourier, andBiotSavartclasses as JAX PyTrees with proper tree flattening/unflattening - Added
essos/losses.pywithcustom_lossandcomposite_lossclasses for differentiable loss functions - Updated API:
Coils_from_json()→Coils.from_json(),tracing.energyproperty →tracing.energy()method - Added extensive analysis scripts for validation against SIMSOPT
Reviewed changes
Copilot reviewed 32 out of 32 changed files in this pull request and generated 39 comments.
Show a summary per file
| File | Description |
|---|---|
essos/losses.py |
New module implementing base_loss, custom_loss, and composite_loss classes for automatic differentiation |
essos/surfaces.py |
Refactored SurfaceRZFourier as PyTree with cached properties and improved initialization |
essos/coils.py |
Refactored Curves and Coils as PyTrees with cached properties, changed to classmethod constructors |
essos/fields.py |
Added MagneticField base class and registered BiotSavart as PyTree |
essos/dynamics.py |
Changed energy from cached property to method, added Particles.join() method |
essos/objective_functions.py |
Removed deprecated loss functions, added new coil separation and curvature losses |
essos/optimization.py |
Updated surface instantiation to include mpol/ntor parameters |
examples/optimize_coils_vmec_surface.py |
Major rewrite using new loss wrapper API instead of old optimization functions |
examples/trace_particles_coils_guidingcenter.py |
Updated imports and API calls (from_json, energy method) |
examples/trace_fieldlines_coils.py |
Updated to use Coils.from_json() |
examples/optimize_coils_particle_confinement_fullorbit.py |
Minor formatting and parameter updates |
examples/optimize_coils_and_surface.py |
Added mpol/ntor parameters, simplified loss calculations |
examples/input_files/*. |
Updated VMEC input file coefficients |
examples/comparisons_SIMSOPT/*.py |
Deleted old comparison scripts |
examples/compare_guidingcenter_fullorbit.py |
Updated particle initialization and energy calculation |
analysis/*.py |
New analysis scripts for validation and benchmarking |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
EstevaoMGomes
commented
Feb 28, 2026
Open
…ppers, solver kwarg in Tracing, override setters for Curves gamma
…or counts, verify solver differentiability
…hic init, revert loss_BdotN kwargs
Copilot review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refactor of coils & surfaces to be proper PyTrees;
Added a loss wrapper to differentiate with respect to the dogs (PyTree leaves);
Added analysis & validation of the code