Major refactoring for JAX-style classes. by EstevaoMGomes · Pull Request #29 · uwplasma/ESSOS

EstevaoMGomes · 2025-12-06T15:30:44Z

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

- 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

… examples

…in integrators

… class for improved step size handling

…s in examples

…for adaptive step size

…ters for performance

…s scripts

… evaluation

…ance plots, and improve layout for better visualization

…d ESSOS

…arate gamma computation

… logic

EstevaoMGomes · 2025-12-06T15:34:47Z

Tests need fixing;
All examples should be fixed;
Future (minor) changes may affect some examples, such as turning gamma from a property to a function;
Coils should have the normalized dofs as leaves and should not update the normalization at runtime.

Copilot

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, and BiotSavart classes as JAX PyTrees with proper tree flattening/unflattening
Added essos/losses.py with custom_loss and composite_loss classes for differentiable loss functions
Updated API: Coils_from_json() → Coils.from_json(), tracing.energy property → 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.

Copilot · 2025-12-06T15:36:21Z

+    @property
+    def dependencies_buffer(self):
+        if self._dependencies_buffer is None:
+            self._dependencies_buffer = tree_util.tree_map(lambda x: jnp.zeros_like(x), self.dependencies)


This 'lambda' is just a simple wrapper around a callable object. Use that object directly.

Suggested change

self._dependencies_buffer = tree_util.tree_map(lambda x: jnp.zeros_like(x), self.dependencies)

self._dependencies_buffer = tree_util.tree_map(jnp.zeros_like, self.dependencies)

Copilot · 2025-12-06T15:36:21Z

+        json_file_stel = curves_stel
+        field_simsopt = load(json_file_stel)
+        coils_simsopt = field_simsopt.coils
+        curves_simsopt = [coil.curve for coil in coils_simsopt]


This assignment to 'curves_simsopt' is unnecessary as it is redefined before this value is used.

Copilot · 2025-12-06T15:36:28Z

+compile_tracing = Tracing('GuidingCenter', field_essos, tmax_gc, timesteps=100, method='Dopri5',
+                  stepsize='adaptive', tol_step_size=trace_tolerance_array[0], particles=particles)
+block_until_ready(compile_tracing.trajectories)
+
+for index, trace_tolerance_ESSOS in enumerate(trace_tolerance_array):
+    num_steps_essos = avg_steps_SIMSOPT_array[index]
+    print(f'Tracing ESSOS guiding center with tolerance={trace_tolerance_ESSOS}')
+    start_time = time()
+    tracing = Tracing('GuidingCenter', field_essos, tmax_gc, timesteps=num_steps_essos, method='Dopri5',
+                    stepsize='adaptive', tol_step_size=trace_tolerance_ESSOS, particles=particles)


Keyword argument 'timesteps' is not a supported parameter name of Tracing.init.
Keyword argument 'tol_step_size' is not a supported parameter name of Tracing.init.
Keyword argument 'method' is not a supported parameter name of Tracing.init.
Keyword argument 'stepsize' is not a supported parameter name of Tracing.init.

Suggested change

compile_tracing = Tracing('GuidingCenter', field_essos, tmax_gc, timesteps=100, method='Dopri5',

stepsize='adaptive', tol_step_size=trace_tolerance_array[0], particles=particles)

block_until_ready(compile_tracing.trajectories)

for index, trace_tolerance_ESSOS in enumerate(trace_tolerance_array):

num_steps_essos = avg_steps_SIMSOPT_array[index]

print(f'Tracing ESSOS guiding center with tolerance={trace_tolerance_ESSOS}')

start_time = time()

tracing = Tracing('GuidingCenter', field_essos, tmax_gc, timesteps=num_steps_essos, method='Dopri5',

stepsize='adaptive', tol_step_size=trace_tolerance_ESSOS, particles=particles)

compile_tracing = Tracing('GuidingCenter', field_essos, tmax_gc, particles=particles)

block_until_ready(compile_tracing.trajectories)

for index, trace_tolerance_ESSOS in enumerate(trace_tolerance_array):

num_steps_essos = avg_steps_SIMSOPT_array[index]

print(f'Tracing ESSOS guiding center with tolerance={trace_tolerance_ESSOS}')

start_time = time()

tracing = Tracing('GuidingCenter', field_essos, tmax_gc, particles=particles)

Copilot · 2025-12-06T15:36:28Z

+tracing_fo = Tracing(field=field, model='FullOrbit', particles=particles, maxtime=tmax_fo,
+                     timesteps=timesteps_fo, tol_step_size=trace_tolerance)


Keyword argument 'timesteps' is not a supported parameter name of Tracing.init.
Keyword argument 'tol_step_size' is not a supported parameter name of Tracing.init.

Copilot · 2025-12-06T15:36:28Z

+tracing_gc = Tracing(field=field, model='GuidingCenter', particles=particles, maxtime=tmax_gc,
+                     timesteps=timesteps_gc, tol_step_size=trace_tolerance)


Keyword argument 'timesteps' is not a supported parameter name of Tracing.init.
Keyword argument 'tol_step_size' is not a supported parameter name of Tracing.init.

Copilot · 2025-12-06T15:36:28Z

+                                   nfp=number_of_field_periods, stellsym=True)
+coils_essos = Coils(curves=curves_essos, currents=[current_on_each_coil]*number_coils_per_half_field_period)
+field_essos = BiotSavart(coils_essos)
+surface_essos = SurfaceRZFourier_ESSOS(vmec, ntheta=ntheta, nphi=nphi, close=False)


Call to SurfaceRZFourier.init with too few arguments; should be no fewer than 5.

Suggested change

surface_essos = SurfaceRZFourier_ESSOS(vmec, ntheta=ntheta, nphi=nphi, close=False)

surface_essos = SurfaceRZFourier_ESSOS(vmec, order_Fourier_series_coils, ntheta=ntheta, nphi=nphi, close=False)

Copilot · 2025-12-06T15:36:29Z

This statement is unreachable.

Suggested change

EXPORT = False

EXPORT = True

EstevaoMGomes · 2026-02-28T18:35:05Z

The name of the file is not correct. It does not take the gradients of any particle confinement loss.

…t for some reason it gets stuck running

EstevaoMGomes added 30 commits April 20, 2025 20:12

Improve loss functions speed

f86f5ee

Refine integrator analysis

408c2c5

Optimize rotation matrix computation in RotatedCurve function

6ff79db

Edit fo and gc integration analysis plots

fdb1508

Improve loss functions computational efficiency

415bbae

Feature: gradient analysis

f38fdea

Add block_until_ready to integrators analysis

079cdb7

Minor tweaks to gradient analysis plot

fe5d81d

Create Poincare Plots analysis

9595ab4

Refactor integrators analysis: add cyclotron frequency calculation an…

108123e

…d adjust num_steps based on dt

Change loss functions to be quadratic & implement separation loss

c80ff2c

Minor improvements in coils class

ea0af03

Add join method to Particles class and optimize tracing parameters in…

5617e7f

… examples

Simplify curve appending logic in apply_symmetries_to_curves

0bb4462

Update cyclotron frequency calculation and adjust tracing parameters …

91f8cd0

…in integrators

Enhance Particles class with phase angle parameter and update Tracing…

549a43f

… class for improved step size handling

Refactor coil separation loss function and optimize tracing parameter…

23bf271

…s in examples

Tentative to incorporate pytree methods for BiotSavart class

dcef625

Refactor Tracing class initialization and improve parameter handling …

3c96714

…for adaptive step size

Add gc_vs_fo.py for particle tracing and visualization; adjust parame…

b1a227e

…ters for performance

Add output directory creation and update file saving paths in analysi…

8332e1e

…s scripts

Add comparison_coils.py for BiotSavart field analysis and performance…

c64a9c1

… evaluation

Fix errors derived from dynamics refactoring

f867857

Enhance plotting in gc_integrators.py: adjust figure sizes, add toler…

7d7be44

…ance plots, and improve layout for better visualization

Change dynamics to accept Integrator name

53682b6

Add comparison_gc.py for comparing gc trajectories between SIMSOPT an…

093f11d

…d ESSOS

Fixed energy calculation for fo trajectories

1c6584c

Finalize gc analysis ESSOS vs SIMSOPT

32b84ad

Add comparison script for fo tracing & improve gc script plots

f902e22

EstevaoMGomes added 12 commits October 27, 2025 19:34

Fix: minor fixes

c7378bb

Merge branch 'en/surface_modes_hotfix' into eg/analysis

d3b073c

Fix(coils): assertion removal in Coils class; Perf(coils): vmap & sep…

d40fa9c

…arate gamma computation

Fix(analysis): precompile coil gammas for comparison with simsopt

37252ef

Fix(surfaces): PyTree able Surfaces

df5a525

Fix(fields,surfaces): fixed mpol in vmec import

a4fad74

Fix(surfaces): jit and pjit fix

ce7caf7

Feat (losses): Create losses files

43fd404

Fix (losses): Allow for grad usage during optimization

ac5c7bc

Clean (losses): Move example to correct file

c9efcf7

Refactor (example): Simplify coils_from_vmec_surface example with new…

41a0c1a

… logic

Refactor (coils): Decache gamma calculation

2324cfe

EstevaoMGomes requested review from Copilot, eduardolneto and rogeriojorge December 6, 2025 15:30

EstevaoMGomes self-assigned this Dec 6, 2025

Copilot started reviewing on behalf of EstevaoMGomes December 6, 2025 15:31 View session

EstevaoMGomes marked this pull request as draft December 6, 2025 15:31

Copilot AI reviewed Dec 6, 2025

View reviewed changes

rogeriojorge and others added 3 commits December 9, 2025 09:59

refactored examples folder

a2af497

Feat (example): import optimize_surface_qs from rj/quasisymmetry_surface

455356e

[fix] get_derivatives_coils_particle_confinement_guidingcenter.py

1163af7

EstevaoMGomes commented Feb 28, 2026

View reviewed changes

EstevaoMGomes added 4 commits February 28, 2026 20:24

[fix] trace_particles_vmec

08a581a

[fix] trace_particles_vmec_Electric_field

143aa01

[fix] trace_particles_coils_guidingcenter

e88e480

[wip] trace_particles_coils_guidingcenter_with_classifier now runs bu…

85f1f13

…t for some reason it gets stuck running

rogeriojorge mentioned this pull request Jun 5, 2026

Pm updates #32

Open

Tejas7007 mentioned this pull request Jun 10, 2026

Fix all examples for PR #29 #35

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Major refactoring for JAX-style classes.#29

Major refactoring for JAX-style classes.#29
EstevaoMGomes wants to merge 71 commits into
mainfrom
eg/analysis

EstevaoMGomes commented Dec 6, 2025

Uh oh!

EstevaoMGomes commented Dec 6, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

Copilot AI Dec 6, 2025

Uh oh!

EstevaoMGomes Feb 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	self._dependencies_buffer = tree_util.tree_map(lambda x: jnp.zeros_like(x), self.dependencies)
	self._dependencies_buffer = tree_util.tree_map(jnp.zeros_like, self.dependencies)

		tracing_fo = Tracing(field=field, model='FullOrbit', particles=particles, maxtime=tmax_fo,
		timesteps=timesteps_fo, tol_step_size=trace_tolerance)

		tracing_gc = Tracing(field=field, model='GuidingCenter', particles=particles, maxtime=tmax_gc,
		timesteps=timesteps_gc, tol_step_size=trace_tolerance)

	surface_essos = SurfaceRZFourier_ESSOS(vmec, ntheta=ntheta, nphi=nphi, close=False)
	surface_essos = SurfaceRZFourier_ESSOS(vmec, order_Fourier_series_coils, ntheta=ntheta, nphi=nphi, close=False)

Conversation

EstevaoMGomes commented Dec 6, 2025

Uh oh!

EstevaoMGomes commented Dec 6, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 6, 2025

Choose a reason for hiding this comment

Uh oh!

EstevaoMGomes Feb 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants