En/review eg/coil perturbations fix by eduardolneto · Pull Request #37 · uwplasma/ESSOS

eduardolneto · 2026-06-26T22:09:42Z

Changing coil_perturbation and coils for coil perturbation and stochastic optimization functionalities to comply with custom_losses

…stic optimization functionalities to comply with custom_losses

EstevaoMGomes

This PR does not seem master-ready. More exactly, the scaling implementation is an important feat, but it needs to be implemented properly.

EstevaoMGomes · 2026-06-27T14:38:36Z

+        return
+
+    raise TypeError(f"Unsupported type {type(curves)}. Expected Curves, Coils, or CoilsFromGamma.")
    #return curves  


EstevaoMGomes · 2026-06-27T14:44:57Z

+        perturbed_base_gamma = base_gamma + perturbation[:, 0, :, :]
+        dofs_new, _ = fit_dofs_from_coils(perturbed_base_gamma, curves.order, curves.n_segments,assume_uniform=True)
+        curves.dofs = dofs_new
+        return


I thought the source code was implemented in a simpler way, like:

dofs, skeleton = ravel_pytree(curves) new_dofs = dofs + perturbation new_curves = skeleton(new_dofs)

where perturbation is a vector with the same shape as dofs.

Btw, this creates a new object instead of altering the last one in-place. I think this is better.

@EstevaoMGomes I am not sure what you want with this comment. the purturbation is usually in cosntruction of the coils, so it is usually given in terms of the gamma quantity and not in fourier coeficients. if the instance is curves_from_Gammas it does what you are suggesting.

EstevaoMGomes · 2026-06-27T14:47:38Z



-def perturb_curves_statistic(curves: Curves,sampler:GaussianSampler, key=None):
+def perturb_curves_statistic(curves, sampler:GaussianSampler, key=None):


This works, but since both functions have similar implementations, wouldn't it make more sense to add another input to indicate whether the perturbation is statistical or systematic?

EstevaoMGomes · 2026-06-27T14:48:40Z

-                 stellsym:    bool = True):
+                 stellsym:    bool = True,
+                 scaling_type: int = 2,
+                 scaling_factor: float = 0,


Set 0.0 as it is a float

EstevaoMGomes · 2026-06-27T14:49:23Z

        self._nfp = nfp
        self._stellsym = stellsym

+        self._scaling_type = scaling_type  # 1 for L-1 norm, 2 for L-2 norm, jnp.inf for L-infinity norm


Maybe should be -1 for L-infinity

We can do that but I woudl think this is more confusing for the user? Surfaces.py was implemented like this as well, which I think is because the options used in jnp.linalg.norm functions are also these ones and not -1.

If it is well documented, it is not confusing. Using jnp.inf may lead to problems as not all infinities are the same (e.g., np.inf vs jnp.inf)

EstevaoMGomes · 2026-06-27T14:52:40Z

    @property
    def dofs(self):
-        return jnp.array(self._dofs)
+        # Apply scaling to each coordinate (X, Y, Z) independently


This makes no sense. When using the new loss class, the dofs that are used are those that form the PyTree, i.e., self._dofs. When you scale only the dofs property, this is only useful for the old optimizations. I think the scaling is a good addition, but as it is implemented, it should not be merged.
Furthermore, it seems overcomplicated for the user.

I am not sure I understand the problem here. Both the property and setter are scaled, as it is also done on surfaces.py, which was already implemented.

The property is different from the attribute. _tree_flatten() gets the property, not the attribute. The scaling is just wrong as is.

I understand that. But this is exactly what previous implemented scalling of surfaces.py do as well, which is what this is mimicking. @EstevaoMGomes

Since that behaviour for surfaces comes from the original branch. Do you want to change both here or chnag eit after this has been pulled to the original branch?

EstevaoMGomes · 2026-06-27T14:56:17Z

-        # Allow downstream code (e.g. coil_perturbation) to override the
-        # computed gamma; otherwise compute from Fourier coefficients.
-        if self._gamma is not None:
-            return self._gamma


What is the reason for removing the delayed initialization? I am becoming more and more confused with these changes...

Addresed, only discretized coils are eager because they are dofs

EstevaoMGomes · 2026-06-27T14:58:37Z

-    return Curves(curves, n_segments=n_segments, nfp=nfp, stellsym=stellsym)
+    return Curves(curves, n_segments=n_segments, nfp=nfp, stellsym=stellsym, scaling_type=scaling_type, scaling_factor=scaling_factor, scale_fixed=scale_fixed)
+
+def extract_axis_from_surface(surface, n_samples: int = 200):


Shouldn't this be in surface.py as a method of the class?

EstevaoMGomes · 2026-06-27T15:00:38Z

+
+    return axis_gamma
+
+def CreateCoilsAroundAxis(n_coils: int,


I am not saying this is wrong but the code looks way too complex for what it does...

EstevaoMGomes · 2026-06-27T15:08:58Z

-    return dofs, gamma_uni
+    return dofs, gamma_uni
+
+class CoilsFromGamma:


Just so I know, is this supposed to be an optimizable or not?

I see it is registered as a pytree, so I suppose so. Then what are the dofs? Is it all the gammas?

rogeriojorge

A couple things I suggest:

Remove unused variables such as old_scaling in the order setter and R_axis in the surface Frenet-frame helper.
Curves use self._dofs, but Coils.to_simsopt() iterates over self.dofs_curves, which goes through the scaled Curves.dofs property. That can make exported SIMSOPT coils differ from the actual geometry when scaling is active.
Curves.dofs now returns scaled DOFs, but _tree_flatten() still exposes self._dofs, so custom_losses/JAX tree operations will see the unscaled variables, not the scaled public dofs. Also, scale_fixed is not included in the PyTree aux data, and scaling_type is stored but not actually used when building self.scaling.

rogeriojorge · 2026-06-29T20:59:45Z

-    This means taht an independent perturbation is applied to the each unique coil
-    Then, the required symmetries are applied to the perturbed unique set of coils


I actually liked these descriptions better, "systematic" is not well defined and this defined it pretty well

rogeriojorge · 2026-06-29T21:01:55Z

@@ -177,14 +227,8 @@ def create_data(order: int) -> jnp.ndarray:
    # gamma_dashdash property
    @property
    def gamma_dashdash(self):
-        if self._gamma_dashdash is not None:
-            return self._gamma_dashdash
        return self._compute_gamma_dashdash()

-    @gamma_dashdash.setter
-    def gamma_dashdash(self, value):
-        self._gamma_dashdash = value
-


This looks important

rogeriojorge · 2026-06-29T21:02:16Z


    @classmethod
-    def from_simsopt(cls, simsopt_curves, nfp=1, stellsym=True):
+    def from_simsopt(cls, simsopt_curves, nfp=1, stellsym=True, scaling_type=2, scaling_factor=0.0):


Not sure what scaling type is? Perhaps should be given in the docstring

rogeriojorge · 2026-06-29T21:06:00Z

+
+    return axis_gamma
+
+def CreateCoilsAroundAxis(n_coils: int,


This seems like something that would help with pyQSC, when optimizing for coils with the near-axis approach. But I would probably move this for a different PR, unless it is really needed for the coils perturbation PR.

…one function perturb_curves.

…cretizedCoils. Removed fucntion to intialize coils frtom axis which will added innstead in future pull request together with other implementations.

…fty or -1. Adding docstrings to clealry state the scalling. Fixing the unflatten and flatten definitions for the surface pytree. So that custom_losses now see the scaled dofs, and unflatten unscales the dofs. This should make custom_losses take gradients against scaled variables while loss values will be still obtained by applying the loss function to the unscaled dofs-

…asses. Note that s fixed scaled option is added and explained on docstrings. In order to also have the intended behaviour in custom_losses/tree flatten/ unflatten interaction (that is loss(unflatten dofs) but gradient of scaled dofs) the current_scale behaviour had to be slighlty changed. Now it can be given as an input, and if it is None, the current_scale will be calculated at the beggining and kept frozen for the rest of the optimization. This is needed because otherwise the tracing gets confused and gives the wrong gradient.

…and currents dofs consistent with Coils/Curves

eduardolneto · 2026-07-04T03:20:02Z

@rogeriojorge @EstevaoMGomes I fixed the scaling dofs behaviour. It is explained in the commits summary. But basically, now the dofs are unscaled consistently during unflatten, and the scaled ones are flatten. This was changed for both surfaces and coils. This means that the loss is calculated with unscaled and the gradients against the scaled ones. For the coils this meant that the current_scale had to be given as an input or, if not, it is calculated at the beginning and frozen. Before it was updated every time, but this would mean that the gradient tracer would see d loss(scaling(sofs) d dofs) which woudl give wrong gradient. DiscretizedCoils now also do the same buit the geometry dofs are the unique gammas which are scaled always to the maximum value of X,Y,Z like currents.

eduardolneto · 2026-07-04T03:20:41Z

I think it should be good to merge. And then test more with the examples once merged into #29 .

rogeriojorge

Looks good to me

Changing coil_perturnation and coils for coil perturbation and stocha…

20442c0

…stic optimization functionalities to comply with custom_losses

eduardolneto requested review from EstevaoMGomes and rogeriojorge June 26, 2026 22:09

eduardolneto changed the base branch from main to eg/analysis June 27, 2026 14:29

EstevaoMGomes reviewed Jun 27, 2026

View reviewed changes

rogeriojorge requested changes Jun 29, 2026

View reviewed changes

Adding missing static metadate to coils pytree

4feff61

eduardolneto mentioned this pull request Jul 2, 2026

En/review eg/fields surfaces objective functions mod #40

Merged

eduardolneto added 5 commits July 3, 2026 05:45

Unifying pertubr_curves_systematic and perturb_curves_statistic into …

24b1009

…one function perturb_curves.

Adding lazy intizalization to CoildFromGamma which was renamed to Dis…

f073dcd

…cretizedCoils. Removed fucntion to intialize coils frtom axis which will added innstead in future pull request together with other implementations.

DiscretizedCoils now also has an option to scale the geometry gammas …

13ca4a4

…and currents dofs consistent with Coils/Curves

rogeriojorge approved these changes Jul 6, 2026

View reviewed changes

eduardolneto merged commit ac3ffa4 into eg/analysis Jul 7, 2026



		def perturb_curves_statistic(curves: Curves,sampler:GaussianSampler, key=None):
		def perturb_curves_statistic(curves, sampler:GaussianSampler, key=None):

		This means taht an independent perturbation is applied to the each unique coil
		Then, the required symmetries are applied to the perturbed unique set of coils

Uh oh!

Conversation

eduardolneto commented Jun 26, 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!

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

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!

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!

eduardolneto Jul 1, 2026 •

edited

Loading