From 04999b31995444ae9ce30a1e99ab23c7b199c1f2 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Fri, 26 Jun 2026 22:46:49 +0100 Subject: [PATCH 01/10] Changing augmented_lagrangian to comply with custom_losses --- essos/augmented_lagrangian.py | 1512 +++++++++++++++++++++++++-------- 1 file changed, 1161 insertions(+), 351 deletions(-) diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index 4bbad19..7ab42e8 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -15,150 +15,493 @@ class LagrangeMultiplier(NamedTuple): """A class containing constrain parameters for Augmented Lagrangian Method""" value: Any penalty: Any + omega: Any + eta: Any sq_grad: Any #For updating squared gradient in case of adaptative penalty and multiplier evolution -#This is used for the usual augmented lagrangian form -def update_method(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): - """Different methods for updating multipliers and penalties +# #This is used for the usual augmented lagrangian form +# def update_method(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): +# """Different methods for updating multipliers and penalties +# """ + + +# pred = lambda x: isinstance(x, LagrangeMultiplier) +# if model_mu=='Constant': +# #jax.debug.print('{m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Monotonic': +# #jax.debug.print('{m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Conditional_True': +# #jax.debug.print('True {m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Conditional_False': +# #jax.debug.print('False {m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Tolerance_True': +# #jax.debug.print('Standard True {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=eta/mu_average**(0.1) +# #omega=omega/mu_average +# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) +# omega=jnp.maximum(omega/mu_average,omega_tol) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_False': +# #jax.debug.print('Standard False {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=1./mu_average**(0.1) +# #omega=1./mu_average +# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) +# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) +# #jax.debug.print('HMMMMMM eta {m}', m=eta) +# omega=jnp.maximum(1./mu_average,omega_tol) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_True': +# #jax.debug.print('Standard True {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=eta/mu_average**(0.1) +# #omega=omega/mu_average +# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) +# omega=jnp.maximum(omega/mu_average,omega_tol) +# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_False': +# #jax.debug.print('Standard False {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=1./mu_average**(0.1) +# #omega=1./mu_average +# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) +# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) +# #jax.debug.print('HMMMMMM eta {m}', m=eta) +# omega=jnp.maximum(1./mu_average,omega_tol) +# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.maximum(x.penalty,gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_Linear_True': +# #jax.debug.print('Standard True {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=eta/mu_average**(0.1) +# #omega=omega/mu_average +# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) +# omega=jnp.maximum(omega/mu_average,omega_tol) +# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*jnp.abs(y.value)),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_Linear_False': +# #jax.debug.print('Standard False {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=1./mu_average**(0.1) +# #omega=1./mu_average +# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) +# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) +# #jax.debug.print('HMMMMMM eta {m}', m=eta) +# omega=jnp.maximum(1./mu_average,omega_tol) +# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.maximum(x.penalty,x.penalty*(1.+gamma*(alpha*x.sq_grad+(1.-alpha)*jnp.abs(y.value)))),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*jnp.abs(y.value)),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_SOTA_True': +# #jax.debug.print('Standard True {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=eta/mu_average**(0.1) +# #omega=omega/mu_average +# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) +# omega=jnp.maximum(omega/mu_average,omega_tol) +# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_SOTA_False': +# #jax.debug.print('Standard False {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=1./mu_average**(0.1) +# #omega=1./mu_average +# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) +# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) +# #jax.debug.print('HMMMMMM eta {m}', m=eta) +# omega=jnp.maximum(1./mu_average,omega_tol) +# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #util_tree=jax.jax.tree_util.tree_map(lambda x: jax.nn.softmax(jnp.abs(x)),updates,is_leaf=pred) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(x.penalty*(1.+jnp.abs(y.value)),mu_max),0.0*x.sq_grad),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Adaptative': +# #jax.debug.print('True {m}', m=model_mu) +# #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*y.value,-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred) +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty**2.)+epsilon)*y.value,-x.penalty+gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty**2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty**2.),params,updates,is_leaf=pred) +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0),params,updates,is_leaf=pred) + + + +# #This is used for the squared form of the augmented Lagrangioan +# def update_method_squared(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): +# """Different methods for updating multipliers and penalties) +# """ + + +# pred = lambda x: isinstance(x, LagrangeMultiplier) +# if model_mu=='Constant': +# #jax.debug.print('{m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier((y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Monotonic': +# #jax.debug.print('{m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Conditional_True': +# #jax.debug.print('True {m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Conditional_False': +# #jax.debug.print('False {m}', m=model_mu) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) +# elif model_mu=='Mu_Tolerance_True': +# #jax.debug.print('Squared True {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=eta/mu_average**(0.1) +# #omega=omega/mu_average +# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) +# omega=jnp.maximum(omega/mu_average,omega_tol) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_False': +# #jax.debug.print('Squared False {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=1./mu_average**(0.1) +# #omega=1./mu_average +# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) +# omega=jnp.maximum(1./mu_average,omega_tol) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_True': +# #jax.debug.print('Squared True {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=eta/mu_average**(0.1) +# #omega=omega/mu_average +# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) +# omega=jnp.maximum(omega/mu_average,omega_tol) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Tolerance_Adaptative_False': +# #jax.debug.print('Squared False {m}', m=model_mu) +# mu_average=penalty_average(params) +# #eta=1./mu_average**(0.1) +# #omega=1./mu_average +# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) +# omega=jnp.maximum(1./mu_average,omega_tol) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega +# elif model_mu=='Mu_Adaptative': +# #jax.debug.print('True {m}', m=model_mu) +# #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) +# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*(y.value-x.value/x.penalty),-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*(y.penalty*2.+(x.value/x.penalty)**2)),params,updates,is_leaf=pred) +# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*(y.value-x.value/x.penalty),-x.penalty+gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*(y.penalty*2.+(x.value/x.penalty)**2)),params,updates,is_leaf=pred) + + + + +# def lagrange_update(model_lagrangian='Standard'): +# """A gradient transformation for Optax that prepares an MDMM gradient +# descent ascent update from a normal gradient descent update. + +# It should be used like this with a base optimizer: +# optimizer = optax.chain( +# optax.sgd(1e-3), +# mdmm_jax.optax_prepare_update(), +# ) + +# Returns: +# An Optax gradient transformation that converts a gradient descent update +# into a gradient descent ascent update. +# """ +# def init_fn(params): +# del params +# return optax.EmptyState() + +# # def update_fn(lagrange_params,updates, state,eta,omega, params=None,model_mu='Constant',beta=2.,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): +# # del params +# # if model_lagrangian=='Standard' : +# # return update_method(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state +# # elif model_lagrangian=='Squared' : +# # return update_method_squared(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state +# # else: +# # print('Lagrangian model not available please select Standard or Squared ') +# # os._exit(0) + +# return optax.GradientTransformation(init_fn, update_fn) + + + + + +class BaseConstraint: + """A minimal mutable container holding `init` and `loss` callables for a constraint. + + This mirrors the simple tuple-like behavior used elsewhere but allows + attribute access and matches the `base_loss` style in `losses.py`. """ + def __init__(self, init: Callable, loss: Callable): + self.init = init + self.loss = loss - pred = lambda x: isinstance(x, LagrangeMultiplier) - if model_mu=='Constant': - #jax.debug.print('{m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Monotonic': - #jax.debug.print('{m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_True': - #jax.debug.print('True {m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_False': - #jax.debug.print('False {m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Tolerance_True': - #jax.debug.print('Standard True {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=eta/mu_average**(0.1) - #omega=omega/mu_average - eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) - omega=jnp.maximum(omega/mu_average,omega_tol) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Tolerance_False': - #jax.debug.print('Standard False {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=1./mu_average**(0.1) - #omega=1./mu_average - eta=jnp.maximum(1./mu_average**(0.1),eta_tol) - #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) - #jax.debug.print('HMMMMMM eta {m}', m=eta) - omega=jnp.maximum(1./mu_average,omega_tol) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - #return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Adaptative': - #jax.debug.print('True {m}', m=model_mu) - #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*y.value,-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred) - - - -#This is used for the squared form of the augmented Lagrangioan -def update_method_squared(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): - """Different methods for updating multipliers and penalties) - """ +class CompositeConstraint: + """Mutable composite constraint container. - - pred = lambda x: isinstance(x, LagrangeMultiplier) - if model_mu=='Constant': - #jax.debug.print('{m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier((y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Monotonic': - #jax.debug.print('{m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_True': - #jax.debug.print('True {m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Conditional_False': - #jax.debug.print('False {m}', m=model_mu) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) - elif model_mu=='Mu_Tolerance_True': - #jax.debug.print('Squared True {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=eta/mu_average**(0.1) - #omega=omega/mu_average - eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) - omega=jnp.maximum(omega/mu_average,omega_tol) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Tolerance_False': - #jax.debug.print('Squared False {m}', m=model_mu) - mu_average=penalty_average(params) - #eta=1./mu_average**(0.1) - #omega=1./mu_average - eta=jnp.maximum(1./mu_average**(0.1),eta_tol) - omega=jnp.maximum(1./mu_average,omega_tol) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - #return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega - elif model_mu=='Mu_Adaptative': - #jax.debug.print('True {m}', m=model_mu) - #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) - return jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*(y.value-x.value/x.penalty),-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*(y.penalty*2.+(x.value/x.penalty)**2)),params,updates,is_leaf=pred) - - - - -def lagrange_update(model_lagrangian='Standard'): - """A gradient transformation for Optax that prepares an MDMM gradient - descent ascent update from a normal gradient descent update. - - It should be used like this with a base optimizer: - optimizer = optax.chain( - optax.sgd(1e-3), - mdmm_jax.optax_prepare_update(), - ) - - Returns: - An Optax gradient transformation that converts a gradient descent update - into a gradient descent ascent update. + Exposes `init` and `loss` callables (same as `Constraint`) while + allowing attaching metadata like `arg_names`, `_dependencies`, and + `selective_map`. `set_dependencies` will propagate dependencies to any + contained `SelectiveConstraint` instances. """ - def init_fn(params): - del params - return optax.EmptyState() - - def update_fn(lagrange_params,updates, state,eta,omega, params=None,model_mu='Constant',beta=2.,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): - del params - if model_lagrangian=='Standard' : - return update_method(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state - elif model_lagrangian=='Squared' : - return update_method_squared(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state + def __init__(self, init_fn: Callable, loss_fn: Callable, selective_map=None, arg_names=None): + self.init = init_fn + self.loss = loss_fn + self.selective_map = selective_map or {} + # keep ordered list of dependency names + self.arg_names = list(arg_names) if arg_names is not None else [] + self._dependencies = {} + self._starting_dofs = None + self._dofs_to_pytree = None + + @property + def dependencies(self): + return self._dependencies + + @dependencies.setter + def dependencies(self, value): + if not isinstance(value, dict): + raise TypeError("dependencies must be a dictionary mapping names to arrays") + self._dependencies = value + for selective in self.selective_map.values(): + selective.dependencies = value + + def set_dependencies(self, deps): + self.dependencies = deps + + @property + def starting_dofs(self): + if self._starting_dofs is None: + if not self._dependencies: + raise RuntimeError("dependencies must be set on composite before accessing starting_dofs") + vals = tuple(self._dependencies[name] for name in self.arg_names) + self._starting_dofs, self._dofs_to_pytree = jax.flatten_util.ravel_pytree(vals) + return self._starting_dofs + + @property + def dofs_to_pytree(self): + if self._dofs_to_pytree is None: + _ = self.starting_dofs + return self._dofs_to_pytree + + +class SelectiveConstraint: + """Wraps a constraint with selective named dependencies, similar to custom_loss. + + This allows constraints to only depend on a subset of the available degrees of freedom + by name, enabling combination of constraints with different argument requirements. + No need to specify indices - just the dependency names! + + Why filtering is necessary: + - Different constraints may require different subsets of arguments (e.g., one depends on + 'field' only, another on 'coil' only, a third on both). + - Named filtering ensures each constraint receives only its required arguments, avoiding: + * Unnecessary computations with unused data + * Constraints expecting different signatures from breaking + * Wasteful memory transfer of irrelevant arrays + - Enables flexible constraint composition where DOF dependencies vary. + + Attributes: + constraint: The underlying Constraint (init_fn, loss_fn) tuple + arg_names: Tuple of argument names that this constraint depends on + dependencies: Dictionary mapping dependency names to their arrays/objects + + Example: + # Create constraints on different DOF subsets - no indices needed! + field_constraint = alm.eq(lambda field: jnp.sum(field**2)) + surface_constraint = alm.eq(lambda surface: jnp.sum(surface**2)) + + selective1 = SelectiveConstraint(field_constraint, 'field') + selective2 = SelectiveConstraint(surface_constraint, 'surface') + + combined = alm.combine(selective1, selective2) + # Set dependencies by name + combined.dependencies = {'field': field_array, 'surface': surface_array} + """ + def __init__(self, constraint: BaseConstraint, *arg_names, **kwargs): + """Initialize a SelectiveConstraint with named dependencies. + + Args: + constraint: A Constraint (init_fn, loss_fn) tuple + *arg_names: Names of the arguments this constraint depends on (order matters!) + """ + if not (hasattr(constraint, 'init') and hasattr(constraint, 'loss')): + raise TypeError(f"constraint must provide `init` and `loss` callables, got {type(constraint)}") + if not arg_names: + raise ValueError("At least one argument name must be provided") + + self.constraint = constraint + self.arg_names = arg_names + self.kwargs = kwargs + self._dependencies = {} + self._starting_dofs = None + self._dofs_to_pytree = None + + @property + def dependencies(self): + """Get the dependencies dictionary.""" + return self._dependencies + + @dependencies.setter + def dependencies(self, value): + """Set dependencies (mapping of arg_names to their values).""" + if not isinstance(value, dict): + raise TypeError("dependencies must be a dictionary mapping names to arrays") + self._dependencies = value + + def _get_filtered_args(self): + """Extract only the required arguments from dependencies by name.""" + filtered_args = [] + for name in self.arg_names: + if name not in self._dependencies: + raise KeyError( + f"SelectiveConstraint '{self.arg_names}' depends on '{name}', " + f"but it's not in dependencies dict. Available: {list(self._dependencies.keys())}" + ) + filtered_args.append(self._dependencies[name]) + return tuple(filtered_args) + + @property + def starting_dofs(self): + if self._starting_dofs is None: + if not self._dependencies: + raise RuntimeError("dependencies must be set before accessing starting_dofs") + vals = tuple(self._dependencies[name] for name in self.arg_names) + self._starting_dofs, self._dofs_to_pytree = jax.flatten_util.ravel_pytree(vals) + return self._starting_dofs + + @property + def dofs_to_pytree(self): + if self._dofs_to_pytree is None: + _ = self.starting_dofs + return self._dofs_to_pytree + + def init(self, *args, **kwargs): + """Initialize constraint parameters using current dependencies.""" + # Prefer explicit call-time args/kwargs; otherwise use stored dependencies and constructor kwargs + if args or kwargs: + return self.constraint.init(*args, **kwargs, **self.kwargs) + args = self._get_filtered_args() + return self.constraint.init(*args, **self.kwargs) + + def loss(self, params, *args, **kwargs): + """Compute loss using current dependencies. + + Args: + params: Constraint parameters from init() + + Returns: + (loss_value, constraint_info) tuple + """ + # Prefer explicit call-time args/kwargs; otherwise use stored dependencies and constructor kwargs + if args or kwargs: + return self.constraint.loss(params, *args, **kwargs, **self.kwargs) + args = self._get_filtered_args() + return self.constraint.loss(params, *args, **self.kwargs) + + +class ScaledConstraint: + """Wraps a constraint with automatic or user-specified scaling. + + Useful when combining constraints with vastly different magnitudes. + Automatically normalizes constraint output by its initial norm or element-wise absolute values. + + Attributes: + constraint: The underlying Constraint (init_fn, loss_fn) tuple + scale_factor: Scalar to multiply constraint output (auto-computed or user-specified) + auto_scale: Whether to auto-compute scale_factor from initial constraint norm + elementwise_scale: If True, scale by 1/abs(constraint) element-wise; if False, scale by 1/norm + + Example: + # Auto-scale based on initial constraint norm (default) + constraint = alm.eq(lambda x: jnp.sum(x**2)) + scaled = ScaledConstraint(constraint, auto_scale=True, elementwise_scale=False) + + # Auto-scale element-wise + scaled = ScaledConstraint(constraint, auto_scale=True, elementwise_scale=True) + + # Or use fixed scaling + scaled = ScaledConstraint(constraint, scale_factor=1.0, auto_scale=False) + """ + def __init__(self, constraint: BaseConstraint, scale_factor=1.0, auto_scale=True, elementwise_scale=False): + """Initialize a ScaledConstraint. + + Args: + constraint: A Constraint (init_fn, loss_fn) tuple + scale_factor: Fixed scaling factor (ignored if auto_scale=True) + auto_scale: If True, automatically compute scale_factor from initial constraint + elementwise_scale: If True, scale element-wise by 1/abs(constraint); if False, scale by 1/norm + """ + if not (hasattr(constraint, 'init') and hasattr(constraint, 'loss')): + raise TypeError(f"constraint must provide `init` and `loss` callables, got {type(constraint)}") + + if not auto_scale and scale_factor <= 0: + raise ValueError(f"scale_factor must be positive, got {scale_factor}") + + self.constraint = constraint + self.scale_factor = scale_factor + self.auto_scale = auto_scale + self.elementwise_scale = elementwise_scale + self._initial_scale = None + self._initial_scale_scalar = None # For loss_value scaling + + def init(self, *args, **kwargs): + """Initialize constraint parameters and compute scale factor if needed.""" + params = self.constraint.init(*args, **kwargs) + + # If auto_scale, compute initial constraint scaling + if self.auto_scale: + _, constraint_info = self.constraint.loss(params, *args, **kwargs) + constraint_flat, unflatten_fn = jax.flatten_util.ravel_pytree(constraint_info) + + if self.elementwise_scale: + # Element-wise scaling: scale = 1 / abs(constraint_flat) + scale_flat = 1.0 / jnp.maximum(jnp.abs(constraint_flat), 1.) + # Unflatten to match constraint_info structure + self._initial_scale = unflatten_fn(scale_flat) + # For loss_value, use mean of element-wise scales + self._initial_scale_scalar = jnp.linalg.norm(scale_flat) + else: + # Norm-based scaling: scale = 1 / norm(constraint_flat) + scale_scalar = 1.0 / jnp.maximum(jnp.linalg.norm(constraint_flat), 1.) + self._initial_scale = scale_scalar + self._initial_scale_scalar = scale_scalar + + return params + + def loss(self, params, *args, **kwargs): + """Compute scaled constraint loss. + + Args: + params: Constraint parameters from init() + *args: Arguments to constraint + **kwargs: Keyword arguments + + Returns: + (scaled_loss_value, constraint_info) tuple + """ + loss_value, constraint_info = self.constraint.loss(params, *args, **kwargs) + + # Apply scaling + if self.auto_scale and self._initial_scale is not None: + # Scale loss_value with scalar + loss_scale = self._initial_scale_scalar + # Scale constraint_info with potentially element-wise scale + info_scale = self._initial_scale else: - print('Lagrangian model not available please select Standard or Squared ') - os._exit(0) - - return optax.GradientTransformation(init_fn, update_fn) - + loss_scale = self.scale_factor + info_scale = self.scale_factor + + return loss_value * loss_scale, jax.tree_util.tree_map(lambda x: x * info_scale, constraint_info) +def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,omega=1.0,eta=1.0,sq_grad=0., weight=1., reduction=jnp.sum): -class Constraint(NamedTuple): - """A pair of pure functions implementing a constraint. - - Attributes: - init: A pure function which, when called with an example instance of - the arguments to the constraint functions, returns a pytree - containing the constraint's learnable parameters. - loss: A pure function which, when called with the the learnable - parameters returned by init() followed by the arguments to the - constraint functions, returns the loss value for the constraint. - """ - init: Callable - loss: Callable - - -def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,sq_grad=0., weight=1., reduction=jnp.sum): """Represents an equality constraint, g(x) = 0. Args: @@ -177,7 +520,11 @@ def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,sq_grad=0., we """ def init_fn(*args, **kwargs): - return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.zeros_like(fun(*args, **kwargs)))} + return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)), + eta+jnp.zeros_like(fun(*args, **kwargs)), + omega+jnp.zeros_like(fun(*args, **kwargs)), + sq_grad+jnp.zeros_like(fun(*args, **kwargs)))} + #return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.maximum(jnp.square(fun(*args, **kwargs)),1.e-4))} if model_lagrangian=='Standard': def loss_fn(params, *args, **kwargs): @@ -188,10 +535,10 @@ def loss_fn(params, *args, **kwargs): inf = fun(*args, **kwargs) return weight * reduction(-params['lambda'].value * inf + params['lambda'].penalty* inf ** 2 / 2+ params['lambda'].value**2 /(2.*params['lambda'].penalty)), inf - return Constraint(init_fn, loss_fn) + return BaseConstraint(init_fn, loss_fn) -def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1., sq_grad=0.,weight=1., reduction=jnp.sum): +def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1.,omega=1.0,eta=1.0, sq_grad=0.,weight=1., reduction=jnp.sum): """Represents an inequality constraint, h(x) >= 0, which uses a slack variable internally to convert it to an equality constraint. @@ -212,8 +559,11 @@ def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1., sq_grad=0., def init_fn(*args, **kwargs): out = fun(*args, **kwargs) - return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.zeros_like(fun(*args, **kwargs))), - 'slack': jax.nn.relu(out) ** 0.5} + return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)), + eta+jnp.zeros_like(fun(*args, **kwargs)), + omega+jnp.zeros_like(fun(*args, **kwargs)), + sq_grad+jnp.zeros_like(fun(*args, **kwargs))), + 'slack': jax.nn.relu(out) ** 0.5} if model_lagrangian=='Standard': def loss_fn(params, *args, **kwargs): @@ -224,28 +574,179 @@ def loss_fn(params, *args, **kwargs): inf = fun(*args, **kwargs) - params['slack'] ** 2 return weight * reduction(-params['lambda'].value * inf + params['lambda'].penalty * inf ** 2 / 2+ params['lambda'].value**2 /(2.*params['lambda'].penalty)), inf - return Constraint(init_fn, loss_fn) + return BaseConstraint(init_fn, loss_fn) def combine(*args): - """Combines multiple constraint tuples into a single constraint tuple. + """Combines constraints with selective named dependencies, mirroring losses.py. + + Each SelectiveConstraint specifies which named arguments it depends on. + Regular Constraints still work (they receive all arguments positionally). + + The returned combined constraint supports both: + - Old style: `combined.init(arg1, arg2); combined.loss(params, arg1, arg2)` + - New style: Set `combined.dependencies = {...}` and call `combined.init/loss()` + + Implementation optimizes for JIT by pre-wrapping constraint functions at combination time, + avoiding dynamic control flow inside the loss/init functions. + + Validation: + - All constraints must be Constraint or SelectiveConstraint objects + - Constraint outputs must be compatible (all scalars, all same shape, etc.) Args: - *args: A series of constraint (init_fn, loss_fn) tuples. + *args: A series of constraint (init_fn, loss_fn) tuples or SelectiveConstraint objects. Returns: - A single (init_fn, loss_fn) tuple that wraps the input constraints. + A combined Constraint with optional `.dependencies` and `.arg_names` attributes. + + Raises: + ValueError: If no constraints provided or constraint types are invalid + TypeError: If constraint objects are not Constraint or SelectiveConstraint """ - init_fns, loss_fns = zip(*args) + if not args: + raise ValueError("At least one constraint must be provided to combine()") + + # Separate init_fns and loss_fns, tracking SelectiveConstraints + constraints_list = [] + selective_map = {} # Maps position to SelectiveConstraint + all_arg_names = [] # ordered list of unique arg names for named dependency mode + + for i, arg in enumerate(args): + # Extract the underlying callable pair and record selective wrappers + if isinstance(arg, SelectiveConstraint): + c = arg.constraint + selective_map[i] = arg + # preserve order and uniqueness of arg names + for name in arg.arg_names: + if name not in all_arg_names: + all_arg_names.append(name) + else: + c = arg + + # Accept tuple (init_fn, loss_fn) or objects with .init and .loss + if isinstance(c, tuple) and len(c) == 2: + init_fn, loss_fn = c + elif hasattr(c, 'init') and hasattr(c, 'loss'): + init_fn, loss_fn = c.init, c.loss + else: + raise TypeError( + f"Constraint {i} must be SelectiveConstraint, BaseConstraint-like, or (init_fn, loss_fn) tuple, got {type(arg)}" + ) + constraints_list.append((init_fn, loss_fn)) + + init_fns, loss_fns = zip(*constraints_list) + + # Pre-wrap constraint functions to bake in filtering at combination time + # This avoids dynamic control flow (if/else) inside JIT-compiled functions + wrapped_init_fns = [] + wrapped_loss_fns = [] + + for i, (init_fn, loss_fn) in enumerate(zip(init_fns, loss_fns)): + if i in selective_map: + selective = selective_map[i] + + # Wrap init_fn to use named dependencies + def make_init_wrapper(s, f): + def wrapped_init(*a, **kw): + # Merge constructor kwargs with call-time kwargs, prefer call-time + merged_kw = {**s.kwargs, **kw} + if a or kw: + # If first arg looks like the flat dofs (array or tracer), map it + first = a[0] if a else None + is_object_like = hasattr(first, 'B') or hasattr(first, 'coils') or hasattr(first, 'dofs') or isinstance(first, dict) + looks_flat = (hasattr(first, 'ndim') or hasattr(first, 'shape') or hasattr(first, 'aval')) + if a and (looks_flat and not is_object_like): + # Try per-selective unravel first; if it fails, fall back to composite unravel + try: + pytrees = s.dofs_to_pytree(a[0]) + return f(*pytrees, **merged_kw) + except Exception: + # fall back to composite-level mapping if available + if hasattr(s, '_composite_index_map') and hasattr(combined, 'dofs_to_pytree'): + all_pytrees = combined.dofs_to_pytree(a[0]) + pytrees = tuple(all_pytrees[idx] for idx in s._composite_index_map) + return f(*pytrees, **merged_kw) + raise + if is_object_like: + return f(*a, **merged_kw) + # fall through to use stored dependencies + filtered_args = s._get_filtered_args() + return f(*filtered_args, **s.kwargs) + return wrapped_init + wrapped_init_fns.append(make_init_wrapper(selective, init_fn)) + + # Wrap loss_fn to use named dependencies + def make_loss_wrapper(s, f): + def wrapped_loss(p, *a, **kw): + merged_kw = {**s.kwargs, **kw} + if a or kw: + first = a[0] if a else None + is_object_like = hasattr(first, 'B') or hasattr(first, 'coils') or hasattr(first, 'dofs') or isinstance(first, dict) + looks_flat = (hasattr(first, 'ndim') or hasattr(first, 'shape') or hasattr(first, 'aval')) + if a and (looks_flat and not is_object_like): + try: + pytrees = s.dofs_to_pytree(a[0]) + return f(p, *pytrees, **merged_kw) + except Exception: + if hasattr(s, '_composite_index_map') and hasattr(combined, 'dofs_to_pytree'): + all_pytrees = combined.dofs_to_pytree(a[0]) + pytrees = tuple(all_pytrees[idx] for idx in s._composite_index_map) + return f(p, *pytrees, **merged_kw) + raise + if is_object_like: + return f(p, *a, **merged_kw) + # fall through to stored dependencies + filtered_args = s._get_filtered_args() + return f(p, *filtered_args, **s.kwargs) + return wrapped_loss + wrapped_loss_fns.append(make_loss_wrapper(selective, loss_fn)) + else: + # Regular constraints: pass-through wrappers that accept all args + def make_init_wrapper_all(f): + def wrapped_init(*args, **kwargs): + return f(*args, **kwargs) + return wrapped_init + wrapped_init_fns.append(make_init_wrapper_all(init_fn)) + + def make_loss_wrapper_all(f): + def wrapped_loss(p, *args, **kwargs): + return f(p, *args, **kwargs) + return wrapped_loss + wrapped_loss_fns.append(make_loss_wrapper_all(loss_fn)) + + # Now init_fn and loss_fn are clean - no control flow, just list comprehensions def init_fn(*args, **kwargs): - return tuple(fn(*args, **kwargs) for fn in init_fns) + """Initialize all constraints using pre-wrapped functions.""" + results = [fn(*args, **kwargs) for fn in wrapped_init_fns] + return tuple(results) def loss_fn(params, *args, **kwargs): - outs = [fn(p, *args, **kwargs) for p, fn in zip(params, loss_fns)] + """Compute total loss from all constraints using pre-wrapped functions.""" + outs = [fn(p, *args, **kwargs) for p, fn in zip(params, wrapped_loss_fns)] return sum(x[0] for x in outs), tuple(x[1] for x in outs) - return Constraint(init_fn, loss_fn) + combined = CompositeConstraint(init_fn, loss_fn, selective_map=selective_map, arg_names=all_arg_names) + + + # keep legacy-style attributes for compatibility + combined._dependencies = {} + combined.selective_map = selective_map + + # Precompute mapping from composite arg_names -> indices for each selective + for sel in selective_map.values(): + sel._composite_index_map = tuple(combined.arg_names.index(n) for n in sel.arg_names) + + # Add property for dependency management (keeps previous API) + def set_dependencies(deps): + combined._dependencies = deps + for selective in selective_map.values(): + selective.dependencies = deps + + combined.set_dependencies = set_dependencies + + return combined @@ -271,6 +772,178 @@ def penalty_average(tree): return jnp.average(penalty[0]) +def apply_mu_tolerance_per_constraint(constraint_dict, grad_dict, constraint_info, constraint_info_prev=None, model_lagrangian='Standard', model_mu='Mu_Adaptative_1', beta=2.0, mu_max=1.e4, alpha=0.99, gamma=1.e-2, epsilon=1.e-8, eta_tol=1.e-4, omega_tol=1.e-6, decrease_tol=0.75): + """Apply Mu update rule for vectorized constraints. + + Supports three strategies: + - Mu_Tolerance: All elements updated together based on norm, uses average penalty + - Mu_Adaptative_1: Element-wise updates using eta parameter + - Mu_Adaptative_2: Element-wise updates using decrease tolerance criterion + + Args: + constraint_dict: dict with 'lambda' and optionally 'slack' keys containing LagrangeMultiplier objects + grad_dict: dict with gradients matching the constraint_dict structure + constraint_info: current constraint violation information (array or scalar) + constraint_info_prev: previous constraint violation info (needed for Mu_Adaptative_2) + model_lagrangian: 'Standard' or 'Squared' lagrangian formulation + model_mu: 'Mu_Tolerance' (global norm, avg penalty), 'Mu_Adaptative_1' (element-wise eta), or 'Mu_Adaptative_2' (element-wise decrease) + decrease_tol: tolerance for decrease criterion in Mu_Adaptative_2 (default 0.75 = 25% decrease) + """ + pred = lambda x: isinstance(x, LagrangeMultiplier) + + # Extract eta and omega from the first LagrangeMultiplier + first_key = list(constraint_dict.keys())[0] + eta_val = constraint_dict[first_key].eta + omega_val = constraint_dict[first_key].omega + + # Extract penalty values + constraint_penalty = jax.tree_util.tree_map(lambda x: x.penalty, constraint_dict, is_leaf=pred) + penalty_val = constraint_penalty[first_key] + + # Get constraint absolute values (element-wise) + constraint_abs = jnp.abs(constraint_info) + + + # Element-wise strategies (Mu_Adaptative_1 or Mu_Adaptative_2) + # Compute updated eta and omega (element-wise) + eta_updated = jnp.maximum(eta_val / jnp.power(penalty_val, 0.1), eta_tol) + omega_updated = jnp.maximum(omega_val / penalty_val, omega_tol) + + eta_updated_false = jnp.maximum(1. / jnp.power(penalty_val, 0.1), eta_tol) + omega_updated_false = jnp.maximum(1.0 / penalty_val, omega_tol) + + # Determine if constraints are satisfied based on model_mu strategy + if model_mu == 'Mu_Adaptative_1': + # Strategy 1: Use eta parameter - constraint satisfied if abs(constraint) < eta + is_satisfied = constraint_abs < eta_val + + elif model_mu == 'Mu_Adaptative_2': + # Strategy 2: Use decrease tolerance - constraint satisfied if abs(constraint) <= decrease_tol * abs(constraint_prev) + if constraint_info_prev is None: + # If no previous constraint, fall back to eta-based check + is_satisfied = constraint_abs < eta_val + else: + constraint_abs_prev = jnp.abs(constraint_info_prev) + # Check if constraint decreased by specified tolerance factor + is_satisfied = constraint_abs <= decrease_tol * constraint_abs_prev + else: + # Default to Mu_Adaptative_1 + is_satisfied = constraint_abs < eta_val + + if model_lagrangian == 'Standard': + # For satisfied constraints: update lambda value, set penalty change to 0 + # For unsatisfied constraints: set lambda to 0, increase penalty + updated_dict = jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * y.value, 0.0 * x.value), # lambda update + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), # penalty change + jnp.where(is_satisfied, -x.omega+omega_updated, -x.omega+omega_updated_false), # omega + jnp.where(is_satisfied, -x.eta+eta_updated, -x.eta+eta_updated_false), # eta + 0.0 * x.sq_grad + ), + constraint_dict, grad_dict, is_leaf=pred + ) + + elif model_lagrangian == 'Squared': + # For Squared lagrangian: lambda update differs + updated_dict = jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * (y.value - x.value / x.penalty), 0.0 * x.value), # lambda update + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), # penalty change + jnp.where(is_satisfied, -x.omega+omega_updated, -x.omega+omega_updated_false), # omega + jnp.where(is_satisfied, -x.eta+eta_updated, -x.eta+eta_updated_false), # eta + 0.0 * x.sq_grad + ), + constraint_dict, grad_dict, is_leaf=pred + ) + + return updated_dict + + +def apply_mu_tolerance_all_constraints(constraint_dicts_map, grad_dicts_map=None, constraint_infos_map=None, model_lagrangian='Standard', beta=2.0, mu_max=1.e4, alpha=0.99, gamma=1.e-2, epsilon=1.e-8, eta_tol=1.e-4, omega_tol=1.e-6): + """Apply Mu_Tolerance across all constraints (JAX-compatible, no Python loops or dict indexing). + + This mirrors the per-constraint updater but computes a single global + decision using all constraint infos, then applies the same update to + every constraint's Lagrange multipliers. + Args: + constraint_dicts_map: pytree mapping keys -> per-constraint lagrange dicts + grad_dicts_map: optional pytree of matching shapes with gradient/info dicts + constraint_infos_map: optional pytree mapping keys -> constraint info pytrees + Returns: + pytree with updated LagrangeMultiplier dicts matching `constraint_dicts_map`. + """ + pred = lambda x: isinstance(x, LagrangeMultiplier) + + # Build global flat vector of all constraint infos + if constraint_infos_map is None: + global_norm = jnp.array(0.0) + else: + flat_map = jax.tree_util.tree_map(lambda info: jax.flatten_util.ravel_pytree(info)[0], constraint_infos_map) + leaves = jax.tree_util.tree_leaves(flat_map) + global_norm = jnp.linalg.norm(jnp.concatenate(leaves) if leaves else jnp.array(0.0)) + + # Extract eta/omega/penalty from all LagrangeMultipliers directly - no nested tree_map + # This matches the pattern in penalty_average function which works inside JIT + eta_pytree = jax.tree_util.tree_map(lambda x: x.eta, constraint_dicts_map, is_leaf=pred) + omega_pytree = jax.tree_util.tree_map(lambda x: x.omega, constraint_dicts_map, is_leaf=pred) + penalty_pytree = jax.tree_util.tree_map(lambda x: x.penalty, constraint_dicts_map, is_leaf=pred) + + # Flatten to get scalar values + eta_flat = jax.flatten_util.ravel_pytree(eta_pytree)[0] + omega_flat = jax.flatten_util.ravel_pytree(omega_pytree)[0] + penalty_flat = jax.flatten_util.ravel_pytree(penalty_pytree)[0] + + global_eta = jnp.mean(eta_flat) if eta_flat.size > 0 else eta_tol + global_omega = jnp.mean(omega_flat) if omega_flat.size > 0 else omega_tol + mu_average = jnp.mean(penalty_flat) if penalty_flat.size > 0 else 1.0 + + is_satisfied = global_norm < global_eta + + # Compute updated eta and omega based on satisfaction + eta_updated = jnp.maximum(global_eta / jnp.power(mu_average, 0.1), eta_tol) + omega_updated = jnp.maximum(global_omega / mu_average, omega_tol) + + eta_updated_false = jnp.maximum(1. / jnp.power(mu_average, 0.1), eta_tol) + omega_updated_false = jnp.maximum(1.0 / mu_average, omega_tol) + + # Prepare grad_map fallback (zeroed) if not provided + if grad_dicts_map is None: + def _zero_leaf(leaf): + return LagrangeMultiplier(jnp.zeros_like(leaf.value), jnp.zeros_like(leaf.penalty), jnp.zeros_like(leaf.omega), jnp.zeros_like(leaf.eta), jnp.zeros_like(leaf.sq_grad)) + grad_map_full = jax.tree_util.tree_map(lambda c: jax.tree_util.tree_map(_zero_leaf, c, is_leaf=pred), constraint_dicts_map) + else: + grad_map_full = grad_dicts_map + + def _update_single(cdict, gdict): + if model_lagrangian == 'Standard': + return jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * y.value, 0.0 * x.value), + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), + jnp.where(is_satisfied, -x.omega + omega_updated, -x.omega + omega_updated_false), + jnp.where(is_satisfied, -x.eta + eta_updated, -x.eta + eta_updated_false), + 0.0 * x.sq_grad + ), + cdict, gdict, is_leaf=pred + ) + else: + return jax.tree_util.tree_map( + lambda x, y: LagrangeMultiplier( + jnp.where(is_satisfied, x.penalty * (y.value - x.value / x.penalty), 0.0 * x.value), + jnp.where(is_satisfied, 0.0 * x.penalty, jnp.minimum(beta * x.penalty, mu_max) - x.penalty), + jnp.where(is_satisfied, -x.omega + omega_updated, -x.omega + omega_updated_false), + jnp.where(is_satisfied, -x.eta + eta_updated, -x.eta + eta_updated_false), + 0.0 * x.sq_grad + ), + cdict, gdict, is_leaf=pred + ) + + # Map over constraint dicts, treating dicts as leaves (not LagrangeMultipliers) + updated_map = jax.tree_util.tree_map(_update_single, constraint_dicts_map, grad_map_full, is_leaf=lambda x: isinstance(x, dict)) + return updated_map + + @@ -285,10 +958,10 @@ class ALM(NamedTuple): #This can use optax gradient descent optimizers with different mu updating methods def ALM_model_optax(optimizer: optax.GradientTransformation, #an optimizer from OPTAX - constraints: Constraint, #List of constraints + constraints: BaseConstraint, #List of constraints loss= lambda x: 0., #function which represents the loss (Callable, default 0.) model_lagrangian='Standard' , #Model to use for updating lagrange multipliers - model_mu='Constant' , #Model to use for updating lagrange multipliers + model_mu='Mu_Tolerance' , #Model to use for updating lagrange multipliers beta=2.0, mu_max=1.e4, alpha=0.99, @@ -299,87 +972,48 @@ def ALM_model_optax(optimizer: optax.GradientTransformation, #an optimizer from **kargs, #Extra key arguments for loss ): - - if model_mu=='Mu_Tolerance_LBFGS': - @jax.jit - def init_fn(params,**kargs): - main_params,lagrange_params=params - main_state = optimizer.init(main_params) - lag_state=lagrange_update(model_lagrangian=model_lagrangian).init(lagrange_params) - opt_state=main_state,lag_state - value,grad=jax.value_and_grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - return opt_state,grad,value[0],value[1] - else: - @jax.jit - def init_fn(params,**kargs): - main_params,lagrange_params=params - main_state = optimizer.init(main_params) - lag_state=lagrange_update(model_lagrangian=model_lagrangian).init(lagrange_params) - opt_state=main_state,lag_state - grad,info=jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - return opt_state,grad,info + @jax.jit + def init_fn(params,**kargs): + main_params,lagrange_params=params + main_state = optimizer.init(main_params) + lag_state=None + grad,info=jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) + return (main_state,lag_state),grad,info # Define the Augmented lagrangian if model_lagrangian=='Standard': def lagrangian(main_params,lagrange_params,**kargs): - main_loss = jnp.linalg.norm(loss(main_params,**kargs)) #The norm here is to ensure we have a scalr from the loss which should be a vector + main_loss = jnp.linalg.norm(loss(main_params,**kargs)) mdmm_loss, inf = constraints.loss(lagrange_params, main_params) return main_loss+mdmm_loss, (main_loss,main_loss+mdmm_loss, inf) - # Augmented Lagrangian - def lagrangian_lbfgs(main_params,lagrange_params,**kargs): - main_loss = jnp.linalg.norm(loss(main_params,**kargs)) - mdmm_loss, _ = constraints.loss(lagrange_params, main_params) - return main_loss+mdmm_loss - elif model_lagrangian=='Squared': def lagrangian(main_params,lagrange_params,**kargs): main_loss = jnp.square(jnp.linalg.norm(loss(main_params,**kargs))) - #Here we take the square because the term appearing in this Lagrangian mdmm_loss, inf = constraints.loss(lagrange_params, main_params) return main_loss+mdmm_loss, (main_loss,main_loss+mdmm_loss, inf) - # Augmented Lagrangian - def lagrangian_lbfgs(main_params,lagrange_params,**kargs): - #Here we take the square because the term appearing in this Lagrangian - main_loss = jnp.square(jnp.linalg.norm(loss(main_params,**kargs))) - mdmm_loss, _ = constraints.loss(lagrange_params, main_params) - return main_loss+mdmm_loss - - if model_mu=='Mu_Conditional': - # Do the optimization step - @partial(jit, static_argnums=(6,7,8,9,10,11,12,13)) - def update_fn(params, opt_state,grad,info,eta,omega,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol,**kargs): + if model_mu=='Mu_Adaptative_2': + @partial(jit, static_argnums=(4,5,6,7,8,9,10)) + def update_fn(params, opt_state,grad,info,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol,**kargs): main_state,lag_state=opt_state main_params,lagrange_params=params - main_updates, main_state = optimizer.update(grad[0], main_state) - main_params = optax.apply_updates(main_params, main_updates) - params=main_params,lagrange_params - grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - true_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Conditional_True',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - false_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Conditional_False',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - lag_updates, lag_state = jax.lax.cond(norm_constraints(info[2]) omega + return jnp.linalg.norm(grad[0])> omega_min def minimization_loop(state): params,main_state,grad,info=state main_params,lagrange_params=params - #jax.debug.print('Loop omega: {omega}', omega=omega) - #jax.debug.print('Loop grad: {grad}', grad=jnp.linalg.norm(grad[0])) main_updates, main_state = optimizer.update(grad[0], main_state) main_params = optax.apply_updates(main_params, main_updates) params=main_params,lagrange_params @@ -389,74 +1023,98 @@ def minimization_loop(state): params,main_state,grad,info=jax.lax.while_loop(condition,minimization_loop,state) main_params,lagrange_params=params - true_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Tolerance_True',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - false_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model='Mu_Tolerance_False',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - lag_updates, lag_state = jax.lax.cond(norm_constraints(info[2]) omega + _,_,grad,_=state + return jnp.linalg.norm(grad[0])> omega_min def minimization_loop(state): - params,main_state,grad,value,info=state + params,main_state,grad,info=state main_params,lagrange_params=params - #jax.debug.print('Loop omega: {omega}', omega=omega) - #jax.debug.print('Loop grad: {grad}', grad=jnp.linalg.norm(grad[0])) - main_updates, main_state = optimizer.update(grad[0], main_state,params=main_params,value=value,grad=grad[0],value_fn=lagrangian_lbfgs,lagrange_params=lagrange_params) + main_updates, main_state = optimizer.update(grad[0], main_state) main_params = optax.apply_updates(main_params, main_updates) params=main_params,lagrange_params - value,grad = jax.value_and_grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - #Here info is in value[1] - state=params,main_state,grad,value[0],value[1] + grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) + state=params,main_state,grad,info return state - params,main_state,grad,value,info=jax.lax.while_loop(condition,minimization_loop,state) + params,main_state,grad,info=jax.lax.while_loop(condition,minimization_loop,state) main_params,lagrange_params=params - true_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Tolerance_True',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - false_func=partial(lagrange_update(model_lagrangian=model_lagrangian).update,model_mu='Mu_Tolerance_False',beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) - lag_updates, lag_state = jax.lax.cond(norm_constraints(info[2]) Date: Mon, 29 Jun 2026 17:02:51 +0100 Subject: [PATCH 02/10] Removing commented section --- essos/augmented_lagrangian.py | 207 ---------------------------------- 1 file changed, 207 deletions(-) diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index 7ab42e8..da13e66 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -22,213 +22,6 @@ class LagrangeMultiplier(NamedTuple): -# #This is used for the usual augmented lagrangian form -# def update_method(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): -# """Different methods for updating multipliers and penalties -# """ - - -# pred = lambda x: isinstance(x, LagrangeMultiplier) -# if model_mu=='Constant': -# #jax.debug.print('{m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Monotonic': -# #jax.debug.print('{m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Conditional_True': -# #jax.debug.print('True {m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Conditional_False': -# #jax.debug.print('False {m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Tolerance_True': -# #jax.debug.print('Standard True {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=eta/mu_average**(0.1) -# #omega=omega/mu_average -# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) -# omega=jnp.maximum(omega/mu_average,omega_tol) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_False': -# #jax.debug.print('Standard False {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=1./mu_average**(0.1) -# #omega=1./mu_average -# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) -# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) -# #jax.debug.print('HMMMMMM eta {m}', m=eta) -# omega=jnp.maximum(1./mu_average,omega_tol) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_True': -# #jax.debug.print('Standard True {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=eta/mu_average**(0.1) -# #omega=omega/mu_average -# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) -# omega=jnp.maximum(omega/mu_average,omega_tol) -# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_False': -# #jax.debug.print('Standard False {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=1./mu_average**(0.1) -# #omega=1./mu_average -# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) -# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) -# #jax.debug.print('HMMMMMM eta {m}', m=eta) -# omega=jnp.maximum(1./mu_average,omega_tol) -# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.maximum(x.penalty,gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_Linear_True': -# #jax.debug.print('Standard True {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=eta/mu_average**(0.1) -# #omega=omega/mu_average -# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) -# omega=jnp.maximum(omega/mu_average,omega_tol) -# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*jnp.abs(y.value)),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_Linear_False': -# #jax.debug.print('Standard False {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=1./mu_average**(0.1) -# #omega=1./mu_average -# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) -# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) -# #jax.debug.print('HMMMMMM eta {m}', m=eta) -# omega=jnp.maximum(1./mu_average,omega_tol) -# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.maximum(x.penalty,x.penalty*(1.+gamma*(alpha*x.sq_grad+(1.-alpha)*jnp.abs(y.value)))),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*jnp.abs(y.value)),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_SOTA_True': -# #jax.debug.print('Standard True {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=eta/mu_average**(0.1) -# #omega=omega/mu_average -# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) -# omega=jnp.maximum(omega/mu_average,omega_tol) -# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_SOTA_False': -# #jax.debug.print('Standard False {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=1./mu_average**(0.1) -# #omega=1./mu_average -# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) -# #jax.debug.print('HMMMMMM mu_av {m}', m=mu_average) -# #jax.debug.print('HMMMMMM eta {m}', m=eta) -# omega=jnp.maximum(1./mu_average,omega_tol) -# ##return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #util_tree=jax.jax.tree_util.tree_map(lambda x: jax.nn.softmax(jnp.abs(x)),updates,is_leaf=pred) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(x.penalty*(1.+jnp.abs(y.value)),mu_max),0.0*x.sq_grad),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Adaptative': -# #jax.debug.print('True {m}', m=model_mu) -# #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*y.value,-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty*2.),params,updates,is_leaf=pred) -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty**2.)+epsilon)*y.value,-x.penalty+gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty**2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*y.penalty**2.),params,updates,is_leaf=pred) -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*y.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0),params,updates,is_leaf=pred) - - - -# #This is used for the squared form of the augmented Lagrangioan -# def update_method_squared(params,updates,eta,omega,model_mu='Constant',beta=2.0,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): -# """Different methods for updating multipliers and penalties) -# """ - - -# pred = lambda x: isinstance(x, LagrangeMultiplier) -# if model_mu=='Constant': -# #jax.debug.print('{m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier((y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Monotonic': -# #jax.debug.print('{m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Conditional_True': -# #jax.debug.print('True {m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Conditional_False': -# #jax.debug.print('False {m}', m=model_mu) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred) -# elif model_mu=='Mu_Tolerance_True': -# #jax.debug.print('Squared True {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=eta/mu_average**(0.1) -# #omega=omega/mu_average -# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) -# omega=jnp.maximum(omega/mu_average,omega_tol) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_False': -# #jax.debug.print('Squared False {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=1./mu_average**(0.1) -# #omega=1./mu_average -# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) -# omega=jnp.maximum(1./mu_average,omega_tol) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_True': -# #jax.debug.print('Squared True {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=eta/mu_average**(0.1) -# #omega=omega/mu_average -# eta=jnp.maximum(eta/mu_average**(0.1),eta_tol) -# omega=jnp.maximum(omega/mu_average,omega_tol) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(x.penalty*(y.value-x.value/x.penalty),0.0*x.value,0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Tolerance_Adaptative_False': -# #jax.debug.print('Squared False {m}', m=model_mu) -# mu_average=penalty_average(params) -# #eta=1./mu_average**(0.1) -# #omega=1./mu_average -# eta=jnp.maximum(1./mu_average**(0.1),eta_tol) -# omega=jnp.maximum(1./mu_average,omega_tol) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(0.0*x.value,-x.penalty+jnp.minimum(beta*x.penalty,mu_max),0.0*x.value),params,updates,is_leaf=pred),eta,omega -# elif model_mu=='Mu_Adaptative': -# #jax.debug.print('True {m}', m=model_mu) -# #Note that y.penalty is the derivative with respect to mu and so it is 0.5*C(x)**2, like the derivative with respect to lambda is C(x) -# #return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*(y.value-x.value/x.penalty),-x.penalty+gamma/(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*(y.penalty*2.+(x.value/x.penalty)**2)),params,updates,is_leaf=pred) -# return jax.jax.tree_util.tree_map(lambda x,y: LagrangeMultiplier(gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon)*(y.value-x.value/x.penalty),-x.penalty+gamma*(jnp.sqrt(alpha*x.sq_grad+(1.-alpha)*y.penalty*2.)+epsilon),-x.sq_grad+alpha*x.sq_grad+(1.-alpha)*(y.penalty*2.+(x.value/x.penalty)**2)),params,updates,is_leaf=pred) - - - - -# def lagrange_update(model_lagrangian='Standard'): -# """A gradient transformation for Optax that prepares an MDMM gradient -# descent ascent update from a normal gradient descent update. - -# It should be used like this with a base optimizer: -# optimizer = optax.chain( -# optax.sgd(1e-3), -# mdmm_jax.optax_prepare_update(), -# ) - -# Returns: -# An Optax gradient transformation that converts a gradient descent update -# into a gradient descent ascent update. -# """ -# def init_fn(params): -# del params -# return optax.EmptyState() - -# # def update_fn(lagrange_params,updates, state,eta,omega, params=None,model_mu='Constant',beta=2.,mu_max=1.e4,alpha=0.99,gamma=1.e-2,epsilon=1.e-8,eta_tol=1.e-4,omega_tol=1.e-6): -# # del params -# # if model_lagrangian=='Standard' : -# # return update_method(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state -# # elif model_lagrangian=='Squared' : -# # return update_method_squared(lagrange_params,updates,eta,omega,model_mu=model_mu,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol), state -# # else: -# # print('Lagrangian model not available please select Standard or Squared ') -# # os._exit(0) - -# return optax.GradientTransformation(init_fn, update_fn) - - - - - class BaseConstraint: """A minimal mutable container holding `init` and `loss` callables for a constraint. From 4c24857cfd53dc54336d4abd24485245c1eaffa4 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Mon, 29 Jun 2026 17:07:35 +0100 Subject: [PATCH 03/10] Adding example of coils-vmec surface optimization with augmented lagrangian compatible with custom_losses structure --- ...surface_augmented_lagrangian_comparison.py | 249 ++++++++++++++++++ 1 file changed, 249 insertions(+) create mode 100644 examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py new file mode 100644 index 0000000..3a338b6 --- /dev/null +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py @@ -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}' +from time import time +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +from essos.optimization import optimize_loss_function +from essos.coils import Coils, CreateEquallySpacedCurves +from essos.fields import BiotSavart +from essos.surfaces import SurfaceRZFourier, BdotN_over_B +from essos.losses import custom_loss, base_loss + +import essos.augmented_lagrangian as alm +from functools import partial +# In this exmple, `scipy.optimize.least_squares` is used for the normal optimization, but any other optimizer, e.g. from +# `scipy.optimize.minimize` or `jaxopt`, can be used as well and may even be preferable. +from scipy.optimize import least_squares + +# Optimization parameters +maximum_function_evaluations=100 + + +input_filepath = os.path.join(os.path.dirname(__name__), "input_files") +vmec_input = os.path.join(input_filepath, 'wout_LandremanPaul2021_QA_reactorScale_lowres.nc') +surface = SurfaceRZFourier.from_wout_file(vmec_input, s=1, ntheta=32, nphi=32, range_torus='half period') + + +""" Creating starting coils and surface """ +N_COILS = 4; FOURIER_ORDER = 3; LARGE_R = 10.; SMALL_R = 5.7; NFP = 2; N_SEGMENTS = 30; STELLSYM = True # Curve parameters +COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) + +init_curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) +init_coils = Coils(curves=init_curves, currents=[COIL_CURRENT]*N_COILS) +init_field = BiotSavart(init_coils) +init_surface=surface + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1.; LENGTH_TARGET = 40. +CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 0.5 +NORMAL_FIELD_WEIGHT = 1. +BdotN_Target_tol=1.e-6 + +""" Creating the loss functions """ +def loss(field, surface): + return jnp.sum(jnp.abs(BdotN_over_B(surface, field))) + +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))) + return bdotn_over_b_loss + +def loss_length_constraint(field): + return jnp.maximum(0, field.coils.length - LENGTH_TARGET) + #return jnp.maximum(0, (field.coils.length - LENGTH_TARGET)/LENGTH_TARGET) + +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)) + +def loss_length(field): + return jnp.mean(jnp.maximum(0, field.coils.length - LENGTH_TARGET)) + +def loss_curvature(field): + return jnp.mean(jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET)) + + +""" Defining custom losses """ +L_normal_field = custom_loss(loss, "field", surface=surface) +L_normal_field_constraint = custom_loss(BdotN_constraint, "field", surface=surface) +L_length_constraint = custom_loss(loss_length_constraint, "field") +L_curvature_constraint = custom_loss(loss_curvature_contraint, "field") +L_length = custom_loss(loss_length, "field") +L_curvature = custom_loss(loss_curvature, "field") + +""" Defining total loss + setting dependencies """ +L_normal_field.dependencies = {"field": init_field} +L_length_constraint.dependencies = {"field": init_field} +L_curvature_constraint.dependencies = {"field": init_field} +L_normal_field_constraint.dependencies = {"field": init_field} +L_length.dependencies = {"field": init_field} +L_curvature.dependencies = {"field": init_field} + + + + + + + +# Create the constraints +penalty = 1.0 #Intial penalty values +multiplier=0. #Initial lagrange multiplier values +#Initializing first tolerances for the inner minimisation loop iteration +omega=1./penalty +eta=1./penalty**0.1 +sq_grad=0.0 #Initial square gradient parameter value for Mu adaptative +model_lagrangian='Standard' #Use standard augmented lagragian suitable for bounded optimizers +#Since we are using LBFGS-B from jaxopt, model_mu will be updated with tolerances so we do not need to difinte the model +model_mu='Tolerance' + + + + +beta=2. #penalty update parameter +mu_max=1.e4 #Maximum penalty parameter allowed +alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu +gamma=1.e-2 +epsilon=1.e-8 +omega_tol=1.e-7 #desired grad_tolerance, associated with grad of lagrangian to main parameters +eta_tol=1.e-7 #desired contraint tolerance, associated with variation of contraints + +#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)) +curvature_constraint=alm.eq(loss_curvature_contraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) +length_constraint=alm.eq(loss_length_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) +field_constraint=alm.eq(BdotN_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) + + + + + + + +C_normal_field_constraint = alm.SelectiveConstraint(field_constraint, "field", surface=surface, target_tol=BdotN_Target_tol) +C_length_constraint = alm.SelectiveConstraint(length_constraint, "field") +C_curvature_constraint = alm.SelectiveConstraint(curvature_constraint, "field") + + + +C_Total_constraint = alm.combine(C_normal_field_constraint, C_length_constraint, C_curvature_constraint) + + + + + + +C_normal_field_constraint.dependencies = {"field": init_field} +C_length_constraint.dependencies = {"field": init_field} +C_curvature_constraint.dependencies = {"field": init_field} +C_Total_constraint.dependencies = {"field": init_field} + + + + + +#If loss=cost_function(x) is not prescribed, f(x)=0 is considered, uncomment second line to use B dot N as a loss and not a constraint +#ALM=alm.ALM_model_jaxopt_lbfgsb(constraints,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) +ALM=alm.ALM_model_jaxopt_lbfgsb(constraints=C_Total_constraint,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) + +#Initializing lagrange multipliers +lagrange_params=C_Total_constraint.init(init_field.dofs) +#parameters are a tuple of the primal/main optimisation parameters and the lagrange multipliers +params = init_field.dofs, lagrange_params +#This is just to initialize an empty state for the lagrange multiplier update and get some information +lag_state,grad,info=ALM.init(params) + + + + + + + + +""" Optimizing with alm""" +t_start = time() + +i=0 +while i<=maximum_function_evaluations and (jnp.linalg.norm(grad[0])>omega_tol or alm.norm_constraints(info[2])>eta_tol): + #One step of ALM optimization + params, lag_state,grad,info = ALM.update(params,lag_state,grad,info) + #if i % 5 == 0: + #print(f'i: {i}, loss f: {info[0]:g}, infeasibility: {alm.total_infeasibility(info[1]):g}') + print(f'i: {i}, loss f: {info[0]:g},loss L: {info[1]:g}, infeasibility: {alm.total_infeasibility(info[2]):g}') + #print('lagrange',params[1]) + i=i+1 + +t_end = time() + + +opt_field_alm = C_normal_field_constraint.dofs_to_pytree(params[0])[0] +opt_coils_alm = opt_field_alm.coils + + + + +""" Defining total loss for nornmal optimization""" +L_total = NORMAL_FIELD_WEIGHT*L_normal_field+ LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature +L_total.dependencies = {"field": init_field} + +""" Optimizing the total loss """ +t_start = time() +res = least_squares(L_total, L_total.starting_dofs, L_total.grad, verbose=2, ftol=1e-5, gtol=1e-5, xtol=1e-14, max_nfev=maximum_function_evaluations) +t_end = time() + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial loss:", L_total(L_total.starting_dofs)) +print("Loss after optimization:", L_total(res.x)) + +opt_field = L_total.dofs_to_pytree(res.x)["field"] +opt_coils = opt_field.coils + + +print(f"\nOptimization took {t_end - t_start:.2f} seconds") +print("Initial B dot N:", jnp.max(BdotN_over_B(surface, init_field))) +print("B dot N after optimization:", jnp.max(BdotN_over_B(surface, opt_field))) +print("B dot N after optimization alm:", jnp.max(BdotN_over_B(surface, opt_field_alm))) +print("Initial curvature :", jnp.average(init_field.coils.curvature,axis=0)) +print("Curvature after optimization:",jnp.average(opt_field.coils.curvature,axis=0)) +print("Curvature after optimization alm:",jnp.average(opt_field_alm.coils.curvature,axis=0)) +print("Curvature target:",CURVATURE_TARGET) +print("Initial length :", init_field.coils.length) +print("Length after optimization:",opt_field.coils.length) +print("Length after optimization alm:",opt_field_alm.coils.length) +print("Length target:",LENGTH_TARGET) + + + + +fig = plt.figure(figsize=(8, 4)) + +ax1 = fig.add_subplot(131, projection='3d') +init_coils.plot(ax=ax1, show=False,label='Initial coils') +surface.plot(ax=ax1, show=False) +ax2 = fig.add_subplot(132, projection='3d') +opt_coils.plot(ax=ax2, show=False,label='Standard optimized coils') +surface.plot(ax=ax2, show=False) +ax3 = fig.add_subplot(133, projection='3d') +opt_coils_alm.plot(ax=ax3, show=False,label='ALM optimized coils') +surface.plot(ax=ax3, show=False) +plt.legend() +plt.tight_layout() +plt.show() + +EXPORT = False +if EXPORT: + output_filepath = os.path.join(os.path.dirname(__file__), "output") + + """ Save the coils to a json file """ + init_coils.to_json(os.path.join(output_filepath, "init_coils_vmec_surface.json")) + opt_coils.to_json(os.path.join(output_filepath, "opt_coils_vmec_surface.json")) + + """ Save results in vtk format to analyze in Paraview """ + surface.to_vtk(os.path.join(output_filepath, "init_surface_vmec_surface.json"), field=init_field) + surface.to_vtk(os.path.join(output_filepath, "final_surface_vmec_surface.json"), field=opt_field) + init_coils.to_vtk(os.path.join(output_filepath, "init_coils_vmec_surface.json")) + opt_coils.to_vtk(os.path.join(output_filepath, "opt_coils_vmec_surface.json")) + opt_coils_alm.to_vtk(os.path.join(output_filepath, "opt_coils_alm_vmec_surface.json")) \ No newline at end of file From 378a8186e253da68cc690d96cf45d04771c586d3 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Wed, 1 Jul 2026 02:49:21 +0100 Subject: [PATCH 04/10] Making cahce behaviour of constraints custom class closer to custom_losses, adding tests for it --- essos/augmented_lagrangian.py | 23 +++++++------- tests/test_augmented_lagrangian.py | 48 +++++++++++++++++++++++++++++- 2 files changed, 57 insertions(+), 14 deletions(-) diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index da13e66..7146b0b 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -51,6 +51,10 @@ def __init__(self, init_fn: Callable, loss_fn: Callable, selective_map=None, arg self._starting_dofs = None self._dofs_to_pytree = None + def clear_cache(self): + self._starting_dofs = None + self._dofs_to_pytree = None + @property def dependencies(self): return self._dependencies @@ -59,6 +63,7 @@ def dependencies(self): def dependencies(self, value): if not isinstance(value, dict): raise TypeError("dependencies must be a dictionary mapping names to arrays") + self.clear_cache() self._dependencies = value for selective in self.selective_map.values(): selective.dependencies = value @@ -133,6 +138,10 @@ def __init__(self, constraint: BaseConstraint, *arg_names, **kwargs): self._dependencies = {} self._starting_dofs = None self._dofs_to_pytree = None + + def clear_cache(self): + self._starting_dofs = None + self._dofs_to_pytree = None @property def dependencies(self): @@ -144,6 +153,7 @@ def dependencies(self, value): """Set dependencies (mapping of arg_names to their values).""" if not isinstance(value, dict): raise TypeError("dependencies must be a dictionary mapping names to arrays") + self.clear_cache() self._dependencies = value def _get_filtered_args(self): @@ -522,23 +532,10 @@ def loss_fn(params, *args, **kwargs): combined = CompositeConstraint(init_fn, loss_fn, selective_map=selective_map, arg_names=all_arg_names) - - # keep legacy-style attributes for compatibility - combined._dependencies = {} - combined.selective_map = selective_map - # Precompute mapping from composite arg_names -> indices for each selective for sel in selective_map.values(): sel._composite_index_map = tuple(combined.arg_names.index(n) for n in sel.arg_names) - # Add property for dependency management (keeps previous API) - def set_dependencies(deps): - combined._dependencies = deps - for selective in selective_map.values(): - selective.dependencies = deps - - combined.set_dependencies = set_dependencies - return combined diff --git a/tests/test_augmented_lagrangian.py b/tests/test_augmented_lagrangian.py index 6be5c41..9f2d329 100644 --- a/tests/test_augmented_lagrangian.py +++ b/tests/test_augmented_lagrangian.py @@ -11,6 +11,7 @@ eq, ineq, combine, + SelectiveConstraint, total_infeasibility, norm_constraints, infty_norm_constraints, @@ -103,6 +104,51 @@ def fun3(x): return x * 2 params = combined.init(jnp.array([2.])) combined.loss(params, jnp.array([2.])) + def test_selective_constraint_dependencies_reset_cached_dofs(self): + selective = SelectiveConstraint(eq(lambda field: field - 1), 'field') + selective.dependencies = {'field': jnp.array([1.0, 2.0])} + first = selective.starting_dofs + + selective.dependencies = {'field': jnp.array([3.0, 4.0, 5.0])} + second = selective.starting_dofs + + self.assertEqual(first.shape[0], 2) + self.assertEqual(second.shape[0], 3) + self.assertTrue(jnp.allclose(second, jnp.array([3.0, 4.0, 5.0]))) + + def test_composite_constraint_dependencies_reset_cached_dofs(self): + c1 = SelectiveConstraint(eq(lambda field: field - 1), 'field') + c2 = SelectiveConstraint(eq(lambda surface: surface + 1), 'surface') + combined = combine(c1, c2) + combined.dependencies = { + 'field': jnp.array([1.0, 2.0]), + 'surface': jnp.array([10.0]), + } + first = combined.starting_dofs + + combined.dependencies = { + 'field': jnp.array([3.0]), + 'surface': jnp.array([20.0, 30.0]), + } + second = combined.starting_dofs + + self.assertEqual(first.shape[0], 3) + self.assertEqual(second.shape[0], 3) + self.assertTrue(jnp.allclose(second, jnp.array([3.0, 20.0, 30.0]))) + + def test_composite_constraint_set_dependencies_resets_cached_dofs(self): + c1 = SelectiveConstraint(eq(lambda field: field - 1), 'field') + combined = combine(c1) + combined.set_dependencies({'field': jnp.array([1.0, 2.0])}) + first = combined.starting_dofs + + combined.set_dependencies({'field': jnp.array([7.0])}) + second = combined.starting_dofs + + self.assertEqual(first.shape[0], 2) + self.assertEqual(second.shape[0], 1) + self.assertTrue(jnp.allclose(second, jnp.array([7.0]))) + def test_total_infeasibility(self): tree = {'a': jnp.array([1.0, -2.0]), 'b': jnp.array([3.0])} result = total_infeasibility(tree) @@ -244,4 +290,4 @@ def fun(x): return x - 1 if __name__ == "__main__": - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__]) From 6df0a09f6de43650d1a79c746bdd4a96443c8b80 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Wed, 1 Jul 2026 03:20:11 +0100 Subject: [PATCH 05/10] Correcting number of fields for LagrangeMultiplier object in tests, reducing number of zero_like intialization when initializing equality and inequality constraints. --- essos/augmented_lagrangian.py | 22 ++++++++++++++-------- tests/test_augmented_lagrangian.py | 18 ++++++++++-------- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index 7146b0b..6639460 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -20,6 +20,17 @@ class LagrangeMultiplier(NamedTuple): sq_grad: Any #For updating squared gradient in case of adaptative penalty and multiplier evolution +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, + ) + + class BaseConstraint: @@ -323,10 +334,8 @@ def eq(fun,model_lagrangian='Standard', multiplier=0.0,penalty=1.,omega=1.0,eta= """ def init_fn(*args, **kwargs): - return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)), - eta+jnp.zeros_like(fun(*args, **kwargs)), - omega+jnp.zeros_like(fun(*args, **kwargs)), - sq_grad+jnp.zeros_like(fun(*args, **kwargs)))} + out = fun(*args, **kwargs) + return {'lambda': _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad)} #return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)),sq_grad+jnp.maximum(jnp.square(fun(*args, **kwargs)),1.e-4))} if model_lagrangian=='Standard': @@ -362,10 +371,7 @@ def ineq(fun, model_lagrangian='Standard', multiplier=0.,penalty=1.,omega=1.0,et def init_fn(*args, **kwargs): out = fun(*args, **kwargs) - return {'lambda': LagrangeMultiplier(multiplier+jnp.zeros_like(fun(*args, **kwargs)),penalty+jnp.zeros_like(fun(*args, **kwargs)), - eta+jnp.zeros_like(fun(*args, **kwargs)), - omega+jnp.zeros_like(fun(*args, **kwargs)), - sq_grad+jnp.zeros_like(fun(*args, **kwargs))), + return {'lambda': _multiplier_like(out, multiplier, penalty, omega, eta, sq_grad), 'slack': jax.nn.relu(out) ** 0.5} if model_lagrangian=='Standard': diff --git a/tests/test_augmented_lagrangian.py b/tests/test_augmented_lagrangian.py index 9f2d329..b91658d 100644 --- a/tests/test_augmented_lagrangian.py +++ b/tests/test_augmented_lagrangian.py @@ -29,14 +29,16 @@ class TestAugmentedLagrangian(unittest.TestCase): def test_lagrange_multiplier(self): - lm = LagrangeMultiplier(value=1.0, penalty=2.0, sq_grad=3.0) + lm = LagrangeMultiplier(value=1.0, penalty=2.0, omega=4.0, eta=5.0, sq_grad=3.0) self.assertEqual(lm.value, 1.0) self.assertEqual(lm.penalty, 2.0) + self.assertEqual(lm.omega, 4.0) + self.assertEqual(lm.eta, 5.0) self.assertEqual(lm.sq_grad, 3.0) def test_update_method_all_modes(self): - params = LagrangeMultiplier(jnp.array([1.]), jnp.array([2.]), jnp.array([0.])) - updates = LagrangeMultiplier(jnp.array([0.5]), jnp.array([0.]), jnp.array([0.])) + params = LagrangeMultiplier(value=jnp.array([1.]), penalty=jnp.array([2.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) + updates = LagrangeMultiplier(value=jnp.array([0.5]), penalty=jnp.array([0.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) for mode in [ 'Constant', 'Mu_Monotonic', 'Mu_Conditional_True', 'Mu_Conditional_False', 'Mu_Tolerance_True', 'Mu_Tolerance_False', 'Mu_Adaptative' @@ -51,8 +53,8 @@ def test_update_method_all_modes(self): self.assertIsInstance(result, LagrangeMultiplier) def test_update_method_squared_all_modes(self): - params = LagrangeMultiplier(jnp.array([1.]), jnp.array([2.]), jnp.array([0.])) - updates = LagrangeMultiplier(jnp.array([0.5]), jnp.array([0.]), jnp.array([0.])) + params = LagrangeMultiplier(value=jnp.array([1.]), penalty=jnp.array([2.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) + updates = LagrangeMultiplier(value=jnp.array([0.5]), penalty=jnp.array([0.]), omega=jnp.array([0.]), eta=jnp.array([0.]), sq_grad=jnp.array([0.])) for mode in [ 'Constant', 'Mu_Monotonic', 'Mu_Conditional_True', 'Mu_Conditional_False', 'Mu_Tolerance_True', 'Mu_Tolerance_False', 'Mu_Adaptative' @@ -165,7 +167,7 @@ def test_infty_norm_constraints(self): self.assertAlmostEqual(float(result), 3.0) def test_penalty_average(self): - tree = {'a': LagrangeMultiplier(jnp.array([1.0]), jnp.array([2.0]), jnp.array([0.0]))} + tree = {'a': LagrangeMultiplier(value=jnp.array([1.0]), penalty=jnp.array([2.0]), omega=jnp.array([0.0]), eta=jnp.array([0.0]), sq_grad=jnp.array([0.0]))} result = penalty_average(tree) self.assertAlmostEqual(float(result), 2.0) @@ -190,8 +192,8 @@ def test_lagrange_update_gradient_transformation_and_update(self): self.assertTrue(hasattr(gt, 'update')) # Call init and update with dummy data params = {'x': jnp.array([1.0])} - lagrange_params = LagrangeMultiplier(jnp.array([0.0]), jnp.array([1.0]), jnp.array([0.0])) - updates = LagrangeMultiplier(jnp.array([-0.5]), jnp.array([1.0]), jnp.array([1.0])) + lagrange_params = LagrangeMultiplier(value=jnp.array([0.0]), penalty=jnp.array([1.0]), omega=jnp.array([0.0]), eta=jnp.array([0.0]), sq_grad=jnp.array([0.0])) + updates = LagrangeMultiplier(value=jnp.array([-0.5]), penalty=jnp.array([1.0]), omega=jnp.array([1.0]), eta=jnp.array([1.0]), sq_grad=jnp.array([1.0])) state = gt.init(params) # eta, omega, etc. are required by update_fn signature eta = {'x': jnp.array([0.0])} From 2c7e480751b84eac712b0124718d614ff5f0457c Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Wed, 1 Jul 2026 03:49:49 +0100 Subject: [PATCH 06/10] Correcting example optimizing_coils_vmec_surface_augmented_lagrangian_comparison.py --- ...surface_augmented_lagrangian_comparison.py | 37 ++++++++++--------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py index 3a338b6..9385516 100644 --- a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py @@ -1,6 +1,4 @@ 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}' from time import time import jax import jax.numpy as jnp @@ -27,19 +25,33 @@ """ Creating starting coils and surface """ -N_COILS = 4; FOURIER_ORDER = 3; LARGE_R = 10.; SMALL_R = 5.7; NFP = 2; N_SEGMENTS = 30; STELLSYM = True # Curve parameters +N_COILS = 4 +FOURIER_ORDER = 3 +LARGE_R = 10. +SMALL_R = 5.7 +NFP = 2 +N_SEGMENTS = 30 +STELLSYM = True # Curve parameters COIL_CURRENT = 1. # Amperes (optimization does not depend on current magnitude) + +""" Setting the losses weights and targets """ +LENGTH_WEIGHT = 1. +LENGTH_TARGET = 40. +CURVATURE_WEIGHT = 1. +CURVATURE_TARGET = 0.5 +NORMAL_FIELD_WEIGHT = 1. +BdotN_Target_tol=1.e-6 + +EXPORT = False + + init_curves = CreateEquallySpacedCurves(N_COILS, FOURIER_ORDER, LARGE_R, SMALL_R, n_segments=N_SEGMENTS, nfp=NFP, stellsym=STELLSYM) init_coils = Coils(curves=init_curves, currents=[COIL_CURRENT]*N_COILS) init_field = BiotSavart(init_coils) init_surface=surface -""" Setting the losses weights and targets """ -LENGTH_WEIGHT = 1.; LENGTH_TARGET = 40. -CURVATURE_WEIGHT = 1.; CURVATURE_TARGET = 0.5 -NORMAL_FIELD_WEIGHT = 1. -BdotN_Target_tol=1.e-6 + """ Creating the loss functions """ def loss(field, surface): @@ -157,12 +169,6 @@ def loss_curvature(field): lag_state,grad,info=ALM.init(params) - - - - - - """ Optimizing with alm""" t_start = time() @@ -183,8 +189,6 @@ def loss_curvature(field): opt_coils_alm = opt_field_alm.coils - - """ Defining total loss for nornmal optimization""" L_total = NORMAL_FIELD_WEIGHT*L_normal_field+ LENGTH_WEIGHT*L_length + CURVATURE_WEIGHT*L_curvature L_total.dependencies = {"field": init_field} @@ -233,7 +237,6 @@ def loss_curvature(field): plt.tight_layout() plt.show() -EXPORT = False if EXPORT: output_filepath = os.path.join(os.path.dirname(__file__), "output") From 3760c570022bf0cc94aff65c439e691e5cbc4e04 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Wed, 1 Jul 2026 18:23:03 +0100 Subject: [PATCH 07/10] Changes --- ...ptimize_coils_vmec_surface_augmented_lagrangian_comparison.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py index 9385516..2633791 100644 --- a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py @@ -158,7 +158,6 @@ def loss_curvature(field): #If loss=cost_function(x) is not prescribed, f(x)=0 is considered, uncomment second line to use B dot N as a loss and not a constraint -#ALM=alm.ALM_model_jaxopt_lbfgsb(constraints,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) ALM=alm.ALM_model_jaxopt_lbfgsb(constraints=C_Total_constraint,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) #Initializing lagrange multipliers From 972e16eb556d4beaee236fc8f698f6bac9cd668e Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Fri, 3 Jul 2026 16:47:21 +0100 Subject: [PATCH 08/10] Addressed bug with model_mu in LavembergMarquardt method in augmented_lagrangian.py --- essos/augmented_lagrangian.py | 55 ++++++++++++++++++++++++----------- 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index 6639460..260323b 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -1070,29 +1070,50 @@ def update_fn(params, lag_state,grad,info,beta=beta,mu_max=mu_max,alpha=alpha,ga main_params=state.params grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) - if model_mu == 'Mu_Tolerance': - lag_updates = apply_mu_tolerance_all_constraints( - lagrange_params, grad_dicts_map=grad[1], constraint_infos_map=info[2], - model_lagrangian=model_lagrangian, beta=beta, mu_max=mu_max, - alpha=alpha, gamma=gamma, epsilon=epsilon, eta_tol=eta_tol, omega_tol=omega_tol - ) - else: - lag_updates = jax.tree_util.tree_map( - lambda lag_p, grad_c, c_info, c_old_info: apply_mu_tolerance_per_constraint( - lag_p, grad_c, c_info, constraint_info_prev=c_old_info, - model_lagrangian=model_lagrangian, model_mu=model_mu, beta=beta, mu_max=mu_max, - alpha=alpha, gamma=gamma, epsilon=epsilon, eta_tol=eta_tol, omega_tol=omega_tol, decrease_tol=0.75 - ), - lagrange_params, grad[1], info[2], old_info, - is_leaf=lambda x: isinstance(x, dict) - ) + lag_updates = jax.tree_util.tree_map( + lambda lag_p, grad_c, c_info, c_old_info: apply_mu_tolerance_per_constraint( + lag_p, grad_c, c_info, constraint_info_prev=c_old_info, + model_lagrangian=model_lagrangian, model_mu=model_mu, beta=beta, mu_max=mu_max, + alpha=alpha, gamma=gamma, epsilon=epsilon, eta_tol=eta_tol, omega_tol=omega_tol, decrease_tol=0.75 + ), + lagrange_params, grad[1], info[2], old_info, + is_leaf=lambda x: isinstance(x, dict) + ) lagrange_params = optax.apply_updates(lagrange_params, lag_updates) params=main_params,lagrange_params grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) return params,lag_state,grad,info - + elif model_mu=='Mu_Tolerance' or model_mu=='Mu_Adaptative_1': + jax.debug.print('omega {omega}:', omega=model_mu) + @partial(jit, static_argnums=(4,5,6,7,8,9,10)) + def update_fn(params, lag_state,grad,info,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol,**kargs): + main_params,lagrange_params=params + pred = lambda x: isinstance(x, LagrangeMultiplier) + omega_vals = jax.tree_util.tree_map(lambda x: x.omega, lagrange_params, is_leaf=pred) + omega_flat = jax.flatten_util.ravel_pytree(omega_vals)[0] + omega_min = jnp.min(omega_flat) + + minimization_loop=jaxopt.LevenbergMarquardt(residual_fun=lagrangian_least_residual,has_aux=True,implicit_diff=False,xtol=1.e-14,gtol=omega_min) + state=minimization_loop.run(main_params,lagrange_params=lagrange_params,**kargs) + main_params=state.params + grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) + + lag_updates = jax.tree_util.tree_map( + lambda lag_p, grad_c, c_info: apply_mu_tolerance_per_constraint( + lag_p, grad_c, c_info, constraint_info_prev=None, + model_lagrangian=model_lagrangian, model_mu=model_mu, beta=beta, mu_max=mu_max, + alpha=alpha, gamma=gamma, epsilon=epsilon, eta_tol=eta_tol, omega_tol=omega_tol, decrease_tol=0.75 + ), + lagrange_params, grad[1], info[2], + is_leaf=lambda x: isinstance(x, dict) + ) + + lagrange_params = optax.apply_updates(lagrange_params, lag_updates) + params=main_params,lagrange_params + grad,info = jax.grad(lagrangian,has_aux=True,argnums=(0,1))(main_params,lagrange_params,**kargs) + return params,lag_state,grad,info return ALM(init_fn,partial(update_fn,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol)) #return optax.GradientTransformationExtraArgs(init_fn, update_fn) From 5e5bd94a73361f02ba8eb1bae8cc9d28f310580e Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Sun, 5 Jul 2026 17:08:50 +0100 Subject: [PATCH 09/10] Removing breaklines on examples --- ...surface_augmented_lagrangian_comparison.py | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py index 2633791..2588fe0 100644 --- a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py @@ -95,11 +95,6 @@ def loss_curvature(field): L_curvature.dependencies = {"field": init_field} - - - - - # Create the constraints penalty = 1.0 #Intial penalty values multiplier=0. #Initial lagrange multiplier values @@ -112,8 +107,6 @@ def loss_curvature(field): model_mu='Tolerance' - - beta=2. #penalty update parameter mu_max=1.e4 #Maximum penalty parameter allowed alpha=0.99 #These are parameters only used if gradient descent and adaaptative mu @@ -130,11 +123,6 @@ def loss_curvature(field): field_constraint=alm.eq(BdotN_constraint,model_lagrangian=model_lagrangian, multiplier=multiplier,penalty=penalty,omega=omega,sq_grad=sq_grad) - - - - - C_normal_field_constraint = alm.SelectiveConstraint(field_constraint, "field", surface=surface, target_tol=BdotN_Target_tol) C_length_constraint = alm.SelectiveConstraint(length_constraint, "field") C_curvature_constraint = alm.SelectiveConstraint(curvature_constraint, "field") @@ -144,19 +132,12 @@ def loss_curvature(field): C_Total_constraint = alm.combine(C_normal_field_constraint, C_length_constraint, C_curvature_constraint) - - - - C_normal_field_constraint.dependencies = {"field": init_field} C_length_constraint.dependencies = {"field": init_field} C_curvature_constraint.dependencies = {"field": init_field} C_Total_constraint.dependencies = {"field": init_field} - - - #If loss=cost_function(x) is not prescribed, f(x)=0 is considered, uncomment second line to use B dot N as a loss and not a constraint ALM=alm.ALM_model_jaxopt_lbfgsb(constraints=C_Total_constraint,model_lagrangian=model_lagrangian,beta=beta,mu_max=mu_max,alpha=alpha,gamma=gamma,epsilon=epsilon,eta_tol=eta_tol,omega_tol=omega_tol) @@ -220,7 +201,6 @@ def loss_curvature(field): - fig = plt.figure(figsize=(8, 4)) ax1 = fig.add_subplot(131, projection='3d') From b02e9f3dfa7506dccf352bdcd00f465ac35e7d45 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Sun, 5 Jul 2026 17:11:59 +0100 Subject: [PATCH 10/10] Removing comments on examples --- ...timize_coils_vmec_surface_augmented_lagrangian_comparison.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py index 2588fe0..0f40631 100644 --- a/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py +++ b/examples/coil_optimization/optimize_coils_vmec_surface_augmented_lagrangian_comparison.py @@ -60,12 +60,10 @@ def loss(field, surface): 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))) return bdotn_over_b_loss def loss_length_constraint(field): return jnp.maximum(0, field.coils.length - LENGTH_TARGET) - #return jnp.maximum(0, (field.coils.length - LENGTH_TARGET)/LENGTH_TARGET) def loss_curvature_contraint(field): return jnp.maximum(0, field.coils.curvature - CURVATURE_TARGET)