Fix all examples for PR #29 by Tejas7007 · Pull Request #35 · uwplasma/ESSOS

Tejas7007 · 2026-06-10T07:11:28Z

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.

…loss wrappers, solver kwarg in Tracing, override setters for Curves gamma

rogeriojorge

Done with review of this PR.

rogeriojorge · 2026-06-10T15:39:31Z

@@ -1,6 +1,6 @@
 import os
 import gc
-number_of_processors_to_use = 1 # Parallelization, this should divide nparticles


Any reason to remove this comment?

rogeriojorge · 2026-06-10T15:39:51Z

 plt.grid(axis='y', which='major', linestyle='--', linewidth=0.6)
 plt.tight_layout()
 plt.savefig(os.path.join(output_dir, 'fo_integration.pdf'))
-plt.show()


We shouldn't be removing code that works

rogeriojorge · 2026-06-10T15:43:24Z

                    dt0=self.timestep,#self.maxtime / self.timesteps,
                    y0=initial_condition,
-                    solver=diffrax.Dopri8(),
+                    solver=(self.solver if self.solver is not None else diffrax.Dopri8()),


Is this if statement differentiable? Could you test if we could do the optimize coils particle confinement script with this solver and get gradients?

Tested it.. gradients flow fine. I traced with both the default solver and an explicit solver=Tsit5() and took jax.grad through each; both give finite, identical gradients. The branch is on self.solver (a Python attribute, not a traced value), so it resolves at trace time and doesn't break differentiability. Added a note in the code.

I think this should be further tested because self.solver may be a tracer when the Tracing class is used as a PyTree.

rogeriojorge · 2026-06-10T15:44:15Z

            energy = vmap(compute_energy)(self.trajectories)

-        elif self.model == 'FullOrbit':
+        elif self.model == 'FullOrbit' or self.model == 'FullOrbit_Boris':


I'm not sure if FullOrbit and FullOrbit_Boris refer to the same model, but if they do, then there's no need for a second redundant name.

They're different integrators for the same physics..Basically, FullOrbit uses the diffrax adaptive solver, FullOrbit_Boris uses the Boris pusher. The energy computation is identical for both, so they share that branch, but the names aren't redundant.

rogeriojorge · 2026-06-10T15:49:28Z

-
+            nfp, mpol, ntor: static
+
+            Backward-compat: rc may also be a filename (string) or a Vmec-like object,


@eduardolneto can you review this init? Do we need such backwards compatibility?

rogeriojorge · 2026-06-10T15:54:51Z

 field = BiotSavart(coils)

-# Particle parameters
-nparticles = number_of_processors_to_use


Wasn't it used before? I guess perhaps not, but it might be good to trace more than one particle.

rogeriojorge · 2026-06-10T15:55:02Z

+cyclotron_frequency = ELEMENTARY_CHARGE * 0.3 / mass
+print("cyclotron period:", 1 / cyclotron_frequency)

-# Particles initialization


Comments like these one are important!

rogeriojorge · 2026-06-10T15:55:31Z

-plt.show()
-
-## Save results in vtk format to analyze in Paraview
-# tracing.to_vtk('trajectories')
-# coils.to_vtk('coils')


All this is important.

rogeriojorge · 2026-06-10T15:57:07Z

@@ -1,13 +1,13 @@
 import os
-number_of_processors_to_use = 12 # Parallelization, this should divide ntheta*nphi


If this change is done to make the examples uniform, that's ok, if it's here because parallelization is not working, then it should be addressed.

…or counts, verify solver differentiability

rogeriojorge · 2026-06-16T14:57:06Z

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

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.

EstevaoMGomes · 2026-06-16T17:23:57Z

+            return self._gamma
        return self._compute_gamma()

+    @gamma.setter


gamma having a setter has the risk of only updating gamma and not gamma_dash and gamma_dash_dash. This should only be used within the coils class and thus should be updated directly via Coil._gamma = ...

EstevaoMGomes · 2026-06-16T17:24:24Z

+            return self._gamma_dash
        return self._compute_gamma_dash()

+    @gamma_dash.setter


The same as for tha gamma setter

EstevaoMGomes · 2026-06-16T17:24:37Z

+            return self._gamma_dashdash
        return self._compute_gamma_dashdash()

+    @gamma_dashdash.setter


The same as for tha gamma setter

EstevaoMGomes · 2026-06-16T17:27:31Z

                    dt0=self.timestep,#self.maxtime / self.timesteps,
                    y0=initial_condition,
-                    solver=diffrax.Dopri8(),
+                    solver=(self.solver if self.solver is not None else diffrax.Dopri8()),


I think this should be further tested because self.solver may be a tracer when the Tracing class is used as a PyTree.

EstevaoMGomes · 2026-06-16T17:29:39Z

 @partial(jit, static_argnums=(1, 4, 5, 6, 7))
-def loss_BdotN(x, vmec, dofs_curves, currents_scale, nfp, max_coil_length=42,
-               n_segments=60, stellsym=True, max_coil_curvature=0.1):
+def loss_BdotN(x, vmec=None, dofs_curves=None, currents_scale=None, nfp=None, max_coil_length=42,


Why set the args as kwargs?

EstevaoMGomes · 2026-06-16T17:33:54Z

+    """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


Does this work with taking the PyTree correspondent to a Coils class? Otherwise it is not compatible with losses.py

EstevaoMGomes · 2026-06-16T17:45:12Z

+            if ntor is None: ntor = nml['ntor']
+            rc = jnp.ravel(nested_lists_to_array(nml['rbc']))[2:]
+            zs = jnp.ravel(nested_lists_to_array(nml['zbs']))[2:]
+        elif rc is not None and not hasattr(rc, '__len__') and hasattr(rc, 'nfp') and hasattr(rc, 'rmnc'):


These will give problems. When rc is dynamic, it is passed downstream as a tracer, which does not have the attributes len, nfp, or rmnc because it is a PyTree. Even if the current examples don't notice the problem, you will see it soon. Delayed initialization is needed in these cases. Also, another useful trick to know is initialization from class methods (like the Coils.from_simsopt)

EstevaoMGomes · 2026-06-16T17:48:50Z


 class SurfaceRZFourier:
-    def __init__(self, rc, zs, nfp, mpol, ntor, ntheta=30, nphi=30, close=True, range_torus='full torus',
+    def __init__(self, rc=None, zs=None, nfp=None, mpol=None, ntor=None, ntheta=30, nphi=30, close=True, range_torus='full torus',


Also here, is there a reason to make the args kwargs? I do not think it makes sense to make empty surface classes

EstevaoMGomes · 2026-06-16T17:52:01Z

+        """ rc, zs: dynamic arrays
+            nfp, mpol, ntor: static
+
+            As a convenience, the first argument (rc) may instead be a filename


I really think this should be done as a class method, like

@class_method def from_file(nml: str, ...)

Also, why give all the nml from rc? rc is just the dynamical vector, so there should be another input as I showed in this class method

eduardolneto · 2026-06-16T19:20:04Z

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.

Fix examples for PR uwplasma#29: API alignment, path fixes, restored …

326f6bf

…loss wrappers, solver kwarg in Tracing, override setters for Curves gamma

rogeriojorge reviewed Jun 10, 2026

View reviewed changes

Address review: restore comments, clarify docstrings, justify process…

f9105f8

…or counts, verify solver differentiability

rogeriojorge requested review from EstevaoMGomes and eduardolneto June 16, 2026 14:56

EstevaoMGomes reviewed Jun 16, 2026

View reviewed changes

EstevaoMGomes mentioned this pull request Jun 18, 2026

Coil_pertubation fix by adding Coils_from_gamma #30

Open


		nfp, mpol, ntor: static

		Backward-compat: rc may also be a filename (string) or a Vmec-like object,

		@@ -1,13 +1,13 @@
		import os
		number_of_processors_to_use = 12 # Parallelization, this should divide ntheta*nphi

Conversation

Tejas7007 commented Jun 10, 2026

Uh oh!

rogeriojorge left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rogeriojorge commented Jun 16, 2026

Uh oh!

EstevaoMGomes left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eduardolneto commented Jun 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants