En/review eg/augmented lagrangian by eduardolneto · Pull Request #36 · uwplasma/ESSOS

eduardolneto · 2026-06-26T21:49:49Z

Mofiying augmented_lagrangian.py so that constraints comply with custom_losses structure.

EstevaoMGomes

Overall looks good

…angian compatible with custom_losses structure

rogeriojorge

A few things:

With a few changes the caching can be made more future-proof. The existing losses.py pattern clears cached DOFs when dependencies change. The new CompositeConstraint and SelectiveConstraint cache _starting_dofs / _dofs_to_pytree, but their dependencies setters do not clear those caches. If dependencies are reassigned after first access, the flattened DOFs can silently refer to old objects.
My take would be to reuse the losses.py cache/dependency pattern instead of reimplementing it. losses.py already has base_loss, custom_loss, dependency setters, starting_dofs, and dofs_to_pytree behavior. The new constraint classes duplicate much of that behavior.
I didn't test it locally, but seems that default arguments for ALM_model_jaxopt_LevenbergMarquardt might fail, as it has model_mu='Mu_Tolerance', but the shown implementation only defines update_fn inside the if model_mu=='Mu_Adaptative_2' branch, then returns partial(update_fn, ...) always. That means the default path can hit an undefined local variable.
The tests should follow the new api change of LagrangeMultiplier from 3 fields to 5 required fields: value, penalty, omega, eta, sq_grad
eq() and ineq() repeatedly call fun(*args, **kwargs) just to build zeros_like arrays. For eq, the init path calls fun several times in one expression. For ineq, it computes out once, then still calls fun repeatedly for the multiplier fields. This can become:

def _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad):
z = jnp.zeros_like(out)
return LagrangeMultiplier(
value=multiplier + z,
penalty=penalty + z,
omega=omega + z,
eta=eta + z,
sq_grad=sq_grad + z,
)

Then both eq and ineq init functions become much shorter and cheaper.

Not now, but maybe in the future, we could refactor some of the code to make it easier to manage and easier to catch bugs. The Optax, JAXopt LBFGS-B, JAXopt LBFGS, JAXopt LM, and Optimistix LM paths all repeat the same pattern: build lagrangian, compute omega_min, run inner solver, recompute grad/info, update lagrange params, recompute grad/info again. You could extract:

def _make_lagrangian(constraints, loss, model_lagrangian):
...

def _omega_min(lagrange_params):
...

def _lagrange_updates(lagrange_params, grad, info, old_info, model_mu, model_lagrangian, ...):
...

def _finish_outer_step(main_params, lagrange_params, lag_updates, lagrangian, **kwargs):
...

That would reduce the duplicated branches.

rogeriojorge · 2026-06-29T20:52:28Z

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


Perhaps parallelization is not needed here?

rogeriojorge · 2026-06-29T20:52:34Z

@@ -0,0 +1,249 @@
+import os
+number_of_processors_to_use = 1 # Parallelization, this should divide ntheta*nphi
+os.environ["XLA_FLAGS"] = f'--xla_force_host_platform_device_count={number_of_processors_to_use}'


rogeriojorge · 2026-06-29T20:53:20Z

+LENGTH_WEIGHT = 1.; LENGTH_TARGET = 40.
+CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 0.5
+NORMAL_FIELD_WEIGHT = 1.


As in the other examples, the main parameters should have one line each, and be at the top before initializing other variables/classes

rogeriojorge · 2026-06-29T20:53:46Z

+def BdotN_constraint(field,surface,target_tol=1.e-6):
+    bdotn_over_b = BdotN_over_B(surface, field)
+    bdotn_over_b_loss = jnp.sqrt(jnp.sum(jnp.maximum(jnp.square(bdotn_over_b)-target_tol,0.0)))
+    #bdotn_over_b_loss = jnp.sqrt(jnp.sum(jnp.maximum((jnp.square(bdotn_over_b)-target_tol)/target_tol,0.0)))


This comment, and the ones below it, perhaps could be deleted to make the file simpler?

rogeriojorge · 2026-06-29T20:53:53Z

+
+def loss_curvature_contraint(field):
+    return jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET)
+    #return jnp.mean(jnp.maximum(0, (field.coils.curvature - CURVATURE_TARGET)/CURVATURE_TARGET))


rogeriojorge · 2026-06-29T20:54:18Z

+
+
+
+
+


Too many line breaks, can be deleted

rogeriojorge · 2026-06-29T20:55:20Z

+#curvature_constraint=alm.ScaledConstraint(alm.eq(loss_curvature_contraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad))
+#length_constraint=alm.ScaledConstraint(alm.eq(loss_length_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad))
+#field_constraint=alm.ScaledConstraint(alm.eq(BdotN_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad))


If these comments are meant to show the user how to use different constraints/objectives, can you add a comment here saying that?

rogeriojorge · 2026-06-29T20:55:30Z

+
+
+
+
+
+


Again, line breaks

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

+
+
+
+
+
+


rogeriojorge · 2026-06-29T20:56:18Z

+plt.tight_layout()
+plt.show()
+
+EXPORT = False


Again, input parameters like these should be at the top

I have addressed most comments. The only thing is that I kept the different method as these will be addressed in future modification. So I think it is better to keep them and modifying after, than erase those extra methods and then add corrected versions later.

…osses, adding tests for it

…educing number of zero_like intialization when initializing equality and inequality constraints.

…_comparison.py

rogeriojorge

While there are some still unresolved comments (too many comments, line breaks, parameters at the top, and optional metods for future) these are small. After the update_fn default path issue is resolved we can swiftly move forward for merge.

rogeriojorge · 2026-07-03T08:49:48Z

+def ALM_model_jaxopt_LevenbergMarquardt(constraints: BaseConstraint,#List of constraints
    loss= lambda x: 0.,                    #function which represents the loss   (Callable, default 0.)
+    model_lagrangian='Squared',            #Model to use for updating lagrange multipliers
+    model_mu='Mu_Tolerance',               #Model to use for updating lagrange multipliers    


The main thing I still see is in ALM_model_jaxopt_LevenbergMarquardt(): the default is now model_mu="Mu_Tolerance", but in the current diff update_fn is only defined inside the if model_mu == "Mu_Adaptative_2". Then the function returns partial(update_fn, ...) always. With the default path, that can still fail with update_fn undefined.

This is already addressed

…_lagrangian.py

eduardolneto · 2026-07-03T15:48:01Z

Addressed last comment

rogeriojorge

Looks good to me

Changing augmented_lagrangian to comply with custom_losses

04999b3

eduardolneto requested review from EstevaoMGomes and rogeriojorge June 26, 2026 21:49

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

EstevaoMGomes reviewed Jun 27, 2026

View reviewed changes

Comment thread essos/augmented_lagrangian.py Outdated

eduardolneto added 2 commits June 29, 2026 17:02

Removing commented section

d64823c

Adding example of coils-vmec surface optimization with augmented lagr…

4c24857

…angian compatible with custom_losses structure

rogeriojorge requested changes Jun 29, 2026

View reviewed changes

eduardolneto added 4 commits July 1, 2026 02:49

Making cahce behaviour of constraints custom class closer to custom_l…

378a818

…osses, adding tests for it

Correcting number of fields for LagrangeMultiplier object in tests, r…

6df0a09

…educing number of zero_like intialization when initializing equality and inequality constraints.

Correcting example optimizing_coils_vmec_surface_augmented_lagrangian…

2c7e480

…_comparison.py

Changes

3760c57

eduardolneto mentioned this pull request Jul 2, 2026

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

Merged

rogeriojorge requested changes Jul 3, 2026

View reviewed changes

Addressed bug with model_mu in LavembergMarquardt method in augmented…

972e16e

…_lagrangian.py

eduardolneto added 2 commits July 5, 2026 17:08

Removing breaklines on examples

5e5bd94

Removing comments on examples

b02e9f3

rogeriojorge approved these changes Jul 6, 2026

View reviewed changes

eduardolneto merged commit 43dde59 into eg/analysis Jul 7, 2026

eduardolneto mentioned this pull request Jul 8, 2026

Fixing all Examples #42

Merged

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

Uh oh!

Conversation

eduardolneto commented Jun 26, 2026

Uh oh!

EstevaoMGomes left a comment

Choose a reason for hiding this comment

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!

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!

eduardolneto commented Jul 3, 2026

Uh oh!

rogeriojorge left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants