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...

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.

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

Open

Unifying pertubr_curves_systematic and perturb_curves_statistic into …

24b1009

…one function perturb_curves.



		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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

eduardolneto Jul 1, 2026 •

edited

Loading