Fixing all Examples by eduardolneto · Pull Request #42 · uwplasma/ESSOS

eduardolneto · 2026-07-08T21:13:22Z

This is a pull request to clean examples and make them work with PR #36 to PR #40. Due to the new changed requesting during pull request some of them stopped working. So I had to add some changes to coils and surface classes to make them work again. There is however, still a problem with examples using stochastic optimization and the classifier object.

rogeriojorge

A couple of changes asked for

rogeriojorge · 2026-07-09T17:04:16Z

    @patch('essos.objective_functions.BiotSavart', return_value=DummyField())
-    @patch('essos.objective_functions.perturb_curves', side_effect=lambda curves, sampler, key=None, perturbation_type=None: curves)
-    def test_perturbed_field_and_coils_from_dofs(self, pc, bs, coils, curves):
+    @patch('essos.objective_functions.perturb_curves_systematic')


Does this perturb_curves_systematic exist? I'm also checking in the objective_functions.py and it appears in line 206, but it is never imported. Doesn't this error out for you?

The same for perturb_curves_statistic, pertubred_field_from_dofs, field_from_dofs, and coils_from_dofs, as the current objective_functions.py only imports perturb_curves and has moved toward object-based losses

rogeriojorge · 2026-07-09T17:06:17Z

+class PytreeCoils:
+    def __init__(self, gamma, gamma_dash, gamma_dashdash, currents, quadpoints):
+        self.gamma = gamma
+        self.gamma_dash = gamma_dash
+        self.gamma_dashdash = gamma_dashdash
+        self.currents = currents
+        self.quadpoints = quadpoints
+
+    def __len__(self):
+        return self.gamma.shape[0]
+
+    def tree_flatten(self):
+        children = (self.gamma, self.gamma_dash, self.gamma_dashdash, self.currents, self.quadpoints)
+        return children, None
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        return cls(*children)
+


Doesn't this error out for you? Seems like PytreeCoils is missing attributes used by the tested losses. loss_coil_surface_distance() needs coils.stellsym and coils.nfp; loss_linkingnumber() and loss_lorentz_force_coils() need coils.curves.quadpoints. What about using a real Coils object in these tests or add a small .curves.quadpoints, .stellsym, and .nfp to the dummy object?

rogeriojorge · 2026-07-09T17:07:05Z


 # Load coils and field
-json_file = os.path.join(os.path.dirname(__file__), '../input_files', 'ESSOS_biot_savart_LandremanPaulQA.json')
+json_file = os.path.join(os.path.dirname(__name__), 'input_files', 'ESSOS_biot_savart_LandremanPaulQA.json')


Didn't we see that this kind of change breaks for other users? Do you really need it?

rogeriojorge · 2026-07-09T17:07:15Z


 # Load coils and field
-wout_file = os.path.join(os.path.dirname(__file__), "../input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc")
+wout_file = os.path.join(os.path.dirname(__name__), "input_files", "wout_LandremanPaul2021_QA_reactorScale_lowres.nc")


rogeriojorge · 2026-07-09T17:07:18Z


 #Load electric field
-Er_file=os.path.join(os.path.dirname(__file__), '../input_files','Er.h5')
+Er_file=os.path.join(os.path.dirname(__name__), 'input_files','Er.h5')


rogeriojorge · 2026-07-09T17:08:54Z

+    R_axis=field.r_axis
+    Z_axis=field.z_axis
+    #Ideally here one would differentiate in time through diffrax !TODO
+    r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,0])+jnp.square(xyz[:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,2]-Z_axis+1.e-12))


I thought xyz is (nparticles, ntimes, 3). Shouldn't this be xyz[:,:,0]?

@eduardolneto

rogeriojorge · 2026-07-09T17:09:18Z

+
+
+
+
+
+


Many line breaks?

rogeriojorge · 2026-07-09T17:10:19Z

+MAXTIME_TRACING = 1e-4
+NUMBER_COILS_PER_HALF_FIELD_PERIOD = 3
+NUMBER_OF_FIELD_PERIODS = 2
+MODEL = 'GuidingCenterAdaptative' 


Is this used later? Seems like it calls particles.to_full_orbit(field) inside the loss. What about importing/using the existing loss_particle_radial_drift() from objective_functions.py instead of redefining it locally?

I kept the exmaples with independent functions on purpose, for people to see what the example is doing more clearly wihtout having to do to the code @rogeriojorge .

rogeriojorge · 2026-07-09T17:11:30Z

+    coils = copy_coils_from_field(field)
+    base_key = jax.random.key(key)
+    split_keys = jax.random.split(base_key, 2)
+    perturb_curves_systematic(coils, sampler, key=split_keys[0])


This is not assigned anywhere? so systematic perturbations are ignored? Should it be coils = perturb_curves_systematic(...) before applying the statistical perturbations?

… trace workflow

…Tree traversal, fixing 'None is not a valid value for jnp.array' crash in tracing (from_simsopt and from_json paths)

Guard Coils.dofs_currents against None/bool sentinels (fixes tracing crash on #42)

eduardolneto · 2026-07-15T18:35:22Z

The only example not working is with from_simsopt method. but I dopnt have a solution for this yet. I would merge this and solve that after. The change fo PR #43 works also on older diffrax version. @rogeriojorge @EstevaoMGomes

rogeriojorge

A couple of small things, and then it's ready to merge

rogeriojorge · 2026-07-15T21:25:13Z

+    @patch("essos.objective_functions.perturb_curves_systematic", create=True)
+    @patch("essos.objective_functions.perturb_curves_statistic", create=True)


I think this is masking the fact that in objective_functions perturb_curves_systematic and perturb_curves_statistic are not imported but still used, which leads to error. Why not use the API already imported there, in objective_functions?

coils = perturb_curves(
coils,
sampler,
key=split_keys[0],
perturbation_type="systematic",
)
coils = perturb_curves(
coils,
sampler,
key=split_keys[1],
perturbation_type="statistical",
)

rogeriojorge · 2026-07-15T21:25:25Z

+    coils = perturb_curves_systematic(coils, sampler, key=split_keys[0])
    coils = perturb_curves_statistic(coils, sampler, key=split_keys[1])


See previous comment:

coils = perturb_curves(
coils,
sampler,
key=split_keys[0],
perturbation_type="systematic",
)
coils = perturb_curves(
coils,
sampler,
key=split_keys[1],
perturbation_type="statistical",
)

rogeriojorge · 2026-07-15T21:26:07Z

+    R_axis=field.r_axis
+    Z_axis=field.z_axis
+    #Ideally here one would differentiate in time through diffrax !TODO
+    r_cross=jnp.sqrt(jnp.square(jnp.sqrt(jnp.square(xyz[:,0])+jnp.square(xyz[:,1]))-R_axis+1.e-12)+jnp.square(xyz[:,2]-Z_axis+1.e-12))


@eduardolneto

rogeriojorge · 2026-07-15T21:27:28Z

-    ax3.plot(tracing.times, trajectory[:, 3]*SPEED_OF_LIGHT/jnp.sqrt(tracing.energy[i]/mass*2.), label=f'Particle {i+1}')    
+for i in subset_indices:
+    trajectory = trajectories[i]
+    ax1.plot(trajectory[:,0], trajectory[:,1], trajectory[:,2], label=f'Particle {i+1}')


I think for this to work you need

ax1 = fig.add_subplot(221, projection="3d")

The same in the other collision examples

Slightly correcting bugs in modifications coils.py

079abb6

rogeriojorge requested changes Jul 9, 2026

View reviewed changes

EstevaoMGomes changed the title ~~Slightly correcting bugs in modifications coils.py~~ Fixing all examples Jul 10, 2026

EstevaoMGomes changed the title ~~Fixing all examples~~ Fixing all Examples Jul 10, 2026

eduardolneto and others added 4 commits July 11, 2026 18:38

Adding sfaeguard to coils and surfaces flatten/ubnflatten methods for…

22c0c39

… trace workflow

Removing repeated cosntructor in from_sismopt methods

df2b80e

Adding some chanmges

10cf2e4

Guard Coils.dofs_currents against None/bool sentinel leaves during Py…

78c6997

…Tree traversal, fixing 'None is not a valid value for jnp.array' crash in tracing (from_simsopt and from_json paths)

Tejas7007 mentioned this pull request Jul 15, 2026

Guard Coils.dofs_currents against None/bool sentinels (fixes tracing crash on #42) #43

Merged

eduardolneto and others added 3 commits July 15, 2026 17:48

Merge pull request #43 from Tejas7007/guard-dofs-currents-for-42

f642c13

Guard Coils.dofs_currents against None/bool sentinels (fixes tracing crash on #42)

Fixing test for objective functions

f94c079

Removing line braks and changing path from file to name

5a10318

eduardolneto requested a review from EstevaoMGomes July 15, 2026 18:35

rogeriojorge requested changes Jul 15, 2026

View reviewed changes

eduardolneto added 2 commits July 16, 2026 00:30

Removing line braks and changing path from file to name

b2e6f48

Fixing collisions plotting bug

9536e33

eduardolneto merged commit 074cf7f into eg/analysis Jul 16, 2026

		@patch("essos.objective_functions.perturb_curves_systematic", create=True)
		@patch("essos.objective_functions.perturb_curves_statistic", create=True)

		coils = perturb_curves_systematic(coils, sampler, key=split_keys[0])
		coils = perturb_curves_statistic(coils, sampler, key=split_keys[1])

Uh oh!

Conversation

eduardolneto commented Jul 8, 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!

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!

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 Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

eduardolneto commented Jul 15, 2026 •

edited

Loading