From fc64159876c22d27b2204dafc5af270d23e1f8f2 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Sat, 27 Jun 2026 05:58:25 +0100 Subject: [PATCH 1/5] Modifying losses.py to allow custom_losses to have a val_and_grad function instead of only a grad function. This is necessary for the augmented lagragian modifications in the parallel pull request, since in lbfgsb the val_and_grad option is more efficient and the only viable option. --- essos/losses.py | 42 +++++++++++++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 7 deletions(-) diff --git a/essos/losses.py b/essos/losses.py index bccdd60..c999015 100644 --- a/essos/losses.py +++ b/essos/losses.py @@ -1,6 +1,6 @@ from functools import partial import jax.numpy as jnp -from jax import tree_util, jit, grad as jax_grad +from jax import tree_util, jit, grad as jax_grad, value_and_grad as jax_value_and_grad from jax.flatten_util import ravel_pytree class base_loss: @@ -71,23 +71,40 @@ def __init__(self, fun, *args_names, **kwargs): self.fun = fun self.args_names = args_names self.kwargs = kwargs + self._dofs_to_args = None + + def clear_cache(self): + super().clear_cache() + self._dofs_to_args = None + + def _ensure_unravelers(self): + if self._starting_dofs is None or self._dofs_to_args is None or self._dofs_to_pytree is None: + self._starting_dofs, tuple_unraveler = ravel_pytree( + tuple(self.dependencies[arg] for arg in self.args_names) + ) + self._dofs_to_args = tuple_unraveler + + def _named_unraveler(dofs): + args = tuple_unraveler(dofs) + return {name: value for name, value in zip(self.args_names, args)} + + self._dofs_to_pytree = _named_unraveler # The dofs of a custom loss are the dofs of its arguments @property def starting_dofs(self): - if self._starting_dofs is None: - self._starting_dofs, self._dofs_to_pytree = ravel_pytree(tuple(self.dependencies[arg] for arg in self.args_names)) + self._ensure_unravelers() return self._starting_dofs @property def dofs_to_pytree(self): - if self._dofs_to_pytree is None: - self._starting_dofs, self._dofs_to_pytree = ravel_pytree(tuple(self.dependencies[arg] for arg in self.args_names)) + self._ensure_unravelers() return self._dofs_to_pytree @partial(jit, static_argnames=['self']) def __call__(self, dofs: jnp.ndarray) -> float: - args = self.dofs_to_pytree(dofs) + self._ensure_unravelers() + args = self._dofs_to_args(dofs) return self.fun(*args, **self.kwargs) @partial(jit, static_argnames=['self']) @@ -96,9 +113,20 @@ def call_pytree(self, dofs_pytree) -> float: @partial(jit, static_argnames=['self']) def grad(self, dofs: jnp.ndarray) -> jnp.ndarray: - args = self.dofs_to_pytree(dofs) + self._ensure_unravelers() + args = self._dofs_to_args(dofs) gradient = jax_grad(self.fun, argnums=tuple(range(len(args))))(*args, **self.kwargs) return ravel_pytree(gradient)[0] + + @partial(jit, static_argnames=['self']) + def value_and_grad(self, dofs: jnp.ndarray): + self._ensure_unravelers() + args = self._dofs_to_args(dofs) + value, gradient = jax_value_and_grad( + self.fun, + argnums=tuple(range(len(args))), + )(*args, **self.kwargs) + return value, ravel_pytree(gradient)[0] @partial(jit, static_argnames=['self']) def grad_pytree(self, dofs_pytree) -> dict: From 1fa359b706920535142d0fb59dab0bfc6645df14 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Wed, 1 Jul 2026 21:01:34 +0100 Subject: [PATCH 2/5] Adding tests for ensure_unravelers --- tests/test_multiobjectives.py | 51 ++++++++++++++++++++++++----------- 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index 019a0ac..860369a 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -1,9 +1,10 @@ -import pytest -from unittest.mock import MagicMock, patch -from essos.multiobjectiveoptimizer import MultiObjectiveOptimizer -from essos.coils import Coils,Curves -from essos.fields import BiotSavart -from essos.objective_functions import loss_bdotn_over_b, loss_coil_length, loss_coil_curvature, loss_normB_axis +import pytest +from unittest.mock import MagicMock, patch +from essos.losses import custom_loss +from essos.multiobjectiveoptimizer import MultiObjectiveOptimizer +from essos.coils import Coils,Curves +from essos.fields import BiotSavart +from essos.objective_functions import loss_bdotn_over_b, loss_coil_length, loss_coil_curvature, loss_normB_axis # test_multiobjectiveoptimizer.py @@ -28,15 +29,35 @@ def mock_vmec(): -def dummy_loss_fn(): - def loss_fn(field=None, coils=None, vmec=None, surface=None, x=None): - return jnp.sum(x) - return loss_fn - - -def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn()): - optimizer = MultiObjectiveOptimizer( - loss_functions=[dummy_loss_fn], +def dummy_loss_fn(): + def loss_fn(field=None, coils=None, vmec=None, surface=None, x=None): + return jnp.sum(x) + return loss_fn + + +def test_custom_loss_named_unraveler(): + def loss_fn(curve_dofs, current): + return jnp.sum(curve_dofs**2) + jnp.sum(current) + + loss = custom_loss(loss_fn, "curve_dofs", "current") + loss.dependencies = { + "curve_dofs": jnp.array([[1.0, 2.0], [3.0, 4.0]]), + "current": jnp.array([5.0]), + "unused": jnp.array([99.0]), + } + + dofs = loss.starting_dofs + named_args = loss.dofs_to_pytree(dofs) + + assert set(named_args) == {"curve_dofs", "current"} + assert jnp.array_equal(named_args["curve_dofs"], loss.dependencies["curve_dofs"]) + assert jnp.array_equal(named_args["current"], loss.dependencies["current"]) + assert loss(dofs) == loss_fn(named_args["curve_dofs"], named_args["current"]) + + +def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn()): + optimizer = MultiObjectiveOptimizer( + loss_functions=[dummy_loss_fn], vmec=vmec, coils_init=None, function_inputs={"extra": 42}, From a9107104fd57704edf2d8eff1b1da20cebb80d53 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Wed, 1 Jul 2026 22:01:20 +0100 Subject: [PATCH 3/5] Solving the inconsistency between call_pytree and dofs_to_pytree and adding tests --- essos/losses.py | 12 ++++++++++-- tests/test_multiobjectives.py | 13 +++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/essos/losses.py b/essos/losses.py index c999015..7854d24 100644 --- a/essos/losses.py +++ b/essos/losses.py @@ -109,7 +109,11 @@ def __call__(self, dofs: jnp.ndarray) -> float: @partial(jit, static_argnames=['self']) def call_pytree(self, dofs_pytree) -> float: - return self.fun(*dofs_pytree, **self.kwargs) + if isinstance(dofs_pytree, dict): + args = tuple(dofs_pytree[name] for name in self.args_names) + else: + args = tuple(dofs_pytree) + return self.fun(*args, **self.kwargs) @partial(jit, static_argnames=['self']) def grad(self, dofs: jnp.ndarray) -> jnp.ndarray: @@ -130,7 +134,11 @@ def value_and_grad(self, dofs: jnp.ndarray): @partial(jit, static_argnames=['self']) def grad_pytree(self, dofs_pytree) -> dict: - gradient = jax_grad(self.fun, argnums=tuple(range(len(dofs_pytree))))(*dofs_pytree, **self.kwargs) + if isinstance(dofs_pytree, dict): + args = tuple(dofs_pytree[name] for name in self.args_names) + else: + args = tuple(dofs_pytree) + gradient = jax_grad(self.fun, argnums=tuple(range(len(args))))(*args, **self.kwargs) buffer = self.dependencies_buffer.copy() for dep, g in zip(self.args_names, gradient): buffer[dep] = g diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index 860369a..619d60e 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -48,11 +48,24 @@ def loss_fn(curve_dofs, current): dofs = loss.starting_dofs named_args = loss.dofs_to_pytree(dofs) + tuple_args = tuple(named_args[name] for name in loss.args_names) assert set(named_args) == {"curve_dofs", "current"} assert jnp.array_equal(named_args["curve_dofs"], loss.dependencies["curve_dofs"]) assert jnp.array_equal(named_args["current"], loss.dependencies["current"]) assert loss(dofs) == loss_fn(named_args["curve_dofs"], named_args["current"]) + assert loss.call_pytree(named_args) == loss_fn(named_args["curve_dofs"], named_args["current"]) + assert loss.call_pytree(tuple_args) == loss_fn(named_args["curve_dofs"], named_args["current"]) + + gradient = loss.grad_pytree(named_args) + gradient_tuple = loss.grad_pytree(tuple_args) + assert set(gradient) == {"curve_dofs", "current", "unused"} + assert jnp.array_equal(gradient["curve_dofs"], 2 * named_args["curve_dofs"]) + assert jnp.array_equal(gradient["current"], jnp.ones_like(named_args["current"])) + assert jnp.array_equal(gradient["unused"], jnp.zeros_like(loss.dependencies["unused"])) + assert jnp.array_equal(gradient_tuple["curve_dofs"], gradient["curve_dofs"]) + assert jnp.array_equal(gradient_tuple["current"], gradient["current"]) + assert jnp.array_equal(gradient_tuple["unused"], gradient["unused"]) def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn()): From f5daf56b778b2bbb8acdd50c36592d7fea028fa5 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Sat, 4 Jul 2026 04:32:27 +0100 Subject: [PATCH 4/5] Added assertion for val_and_grad on test_custom_loss_named_unraveler() --- tests/test_multiobjectives.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index 619d60e..a21f718 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -56,6 +56,8 @@ def loss_fn(curve_dofs, current): assert loss(dofs) == loss_fn(named_args["curve_dofs"], named_args["current"]) assert loss.call_pytree(named_args) == loss_fn(named_args["curve_dofs"], named_args["current"]) assert loss.call_pytree(tuple_args) == loss_fn(named_args["curve_dofs"], named_args["current"]) + value, grad = loss.value_and_grad(dofs) + assert value == loss_fn(named_args["curve_dofs"], named_args["current"]) and jnp.array_equal(grad, loss.grad(dofs)) gradient = loss.grad_pytree(named_args) gradient_tuple = loss.grad_pytree(tuple_args) From b4429c05872b2ee224788139e6946408cc3a9e53 Mon Sep 17 00:00:00 2001 From: eduardolneto Date: Sun, 5 Jul 2026 16:54:47 +0100 Subject: [PATCH 5/5] Addresing unnecessary wrapped code in test_multiobjectives.py origining on PR #21 --- tests/test_multiobjectives.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index a21f718..d833376 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -29,10 +29,8 @@ def mock_vmec(): -def dummy_loss_fn(): - def loss_fn(field=None, coils=None, vmec=None, surface=None, x=None): - return jnp.sum(x) - return loss_fn +def dummy_loss_fn(field=None, coils=None, vmec=None, surface=None, x=None): + return jnp.sum(x) def test_custom_loss_named_unraveler(): @@ -70,7 +68,7 @@ def loss_fn(curve_dofs, current): assert jnp.array_equal(gradient_tuple["unused"], gradient["unused"]) -def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn()): +def test_build_available_inputs( vmec=mock_vmec(), dummy_loss_fn=dummy_loss_fn): optimizer = MultiObjectiveOptimizer( loss_functions=[dummy_loss_fn], vmec=vmec,