En/review eg/losse val grad

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,38 +71,74 @@ 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'])
     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:
-        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:
-        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

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,48 @@ 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(field=None, coils=None, vmec=None, surface=None, x=None):
+    return jnp.sum(x)
+
+
+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)
+    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"])
+    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)
+    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):
+    optimizer = MultiObjectiveOptimizer(
+        loss_functions=[dummy_loss_fn],
         vmec=vmec,
         coils_init=None,
         function_inputs={"extra": 42},