diff --git a/essos/losses.py b/essos/losses.py index bccdd602..7854d249 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,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 diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index 019a0ac5..d8333762 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,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},