ExternalObjective function to wrap external codes

desc/objectives/__init__.py

@@ -18,7 +18,12 @@
     RadialForceBalance,
 )
 from ._free_boundary import BoundaryError, VacuumBoundaryError
-from ._generic import GenericObjective, LinearObjectiveFromUser, ObjectiveFromUser
+from ._generic import (
+    ExternalObjective,
+    GenericObjective,
+    LinearObjectiveFromUser,
+    ObjectiveFromUser,
+)
 from ._geometry import (
     AspectRatio,
     BScaleLength,

desc/objectives/_generic.py

@@ -1,11 +1,12 @@
 """Generic objectives that don't belong anywhere else."""
 
+import functools
 import inspect
 import re
 
 import numpy as np
 
-from desc.backend import jnp
+from desc.backend import jax, jnp
 from desc.compute import data_index
 from desc.compute.utils import _compute as compute_fun
 from desc.compute.utils import get_profiles, get_transforms
@@ -15,6 +16,236 @@
 from .objective_funs import _Objective
 
 
+class ExternalObjective(_Objective):
+    """Wrap an external code.
+
+    Similar to ``ObjectiveFromUser``, except derivatives of the objective function are
+    computed with finite differences instead of AD.
+
+    The user supplied function must take an Equilibrium as its only positional argument,
+    but can take additional keyword arguments.
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium that will be optimized to satisfy the Objective.
+    fun : callable
+        Custom objective function.
+    dim_f : int
+        Dimension of the output of fun.
+    target : {float, ndarray}, optional
+        Target value(s) of the objective. Only used if bounds is None.
+        Must be broadcastable to Objective.dim_f. Defaults to ``target=0``.
+    bounds : tuple of {float, ndarray}, optional
+        Lower and upper bounds on the objective. Overrides target.
+        Both bounds must be broadcastable to to Objective.dim_f.
+        Defaults to ``target=0``.
+    weight : {float, ndarray}, optional
+        Weighting to apply to the Objective, relative to other Objectives.
+        Must be broadcastable to to Objective.dim_f
+    normalize : bool, optional
+        Whether to compute the error in physical units or non-dimensionalize.
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Whether target and bounds should be normalized before comparing to computed
+        values. If `normalize` is `True` and the target is in physical units,
+        this should also be set to True.
+    loss_function : {None, 'mean', 'min', 'max'}, optional
+        Loss function to apply to the objective values once computed. This loss function
+        is called on the raw compute value, before any shifting, scaling, or
+        normalization.
+    name : str, optional
+        Name of the objective function.
+
+    # TODO: add example
+
+    """
+
+    _units = "(Unknown)"
+    _print_value_fmt = "External objective value: {:10.3e}"
+
+    def __init__(
+        self,
+        eq,
+        fun,
+        dim_f,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=False,
+        normalize_target=False,
+        loss_function=None,
+        fd_step=1e-4,  # TODO: generalize this to allow a vector of different scales
+        name="external",
+        **kwargs,
+    ):
+        if target is None and bounds is None:
+            target = 0
+        self._eq = eq.copy()
+        self._fun = fun
+        self._dim_f = dim_f
+        self._fd_step = fd_step
+        self._kwargs = kwargs
+        super().__init__(
+            things=eq,
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            loss_function=loss_function,
+            deriv_mode="fwd",
+            name=name,
+        )
+
+    def build(self, use_jit=True, verbose=1):
+        """Build constant arrays.
+
+        Parameters
+        ----------
+        use_jit : bool, optional
+            Whether to just-in-time compile the objective and derivatives.
+        verbose : int, optional
+            Level of output.
+
+        """
+        self._scalar = self._dim_f == 1
+        self._constants = {"quad_weights": 1.0}
+
+        def fun_wrapped(
+            R_lmn,
+            Z_lmn,
+            L_lmn,
+            p_l,
+            i_l,
+            c_l,
+            Psi,
+            Te_l,
+            ne_l,
+            Ti_l,
+            Zeff_l,
+            a_lmn,
+            Ra_n,
+            Za_n,
+            Rb_lmn,
+            Zb_lmn,
+            I,
+            G,
+            Phi_mn,
+        ):
+            """Wrap external function with optimiazable params arguments."""
+            for param in self._eq.optimizable_params:
+                par = eval(param)  # FIXME: how bad is it to use eval here?
+                if len(par):
+                    setattr(self._eq, param, par)
+            return self._fun(self._eq, **self._kwargs)
+
+        # check to make sure fun_wrapped has the correct signature
+        # in case we ever update Equilibrium.optimizable_params
+        args = inspect.getfullargspec(fun_wrapped).args
+        assert args == self._eq.optimizable_params
+
+        # wrap external function to work with JAX
+        abstract_eval = lambda *args, **kwargs: jnp.empty(self._dim_f)
+        self._fun_wrapped = self._jaxify(fun_wrapped, abstract_eval)
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, params, constants=None):
+        """Compute the quantity.
+
+        Parameters
+        ----------
+        params : list of dict
+            List of dictionaries of degrees of freedom, eg CoilSet.params_dict
+        constants : dict
+            Dictionary of constant data, eg transforms, profiles etc. Defaults to
+            self.constants
+
+        Returns
+        -------
+        f : ndarray
+            Computed quantity.
+
+        """
+        # ensure positional args are passed in the correct order
+        args = [params[k] for k in self._eq.optimizable_params]
+        f = self._fun_wrapped(*args)
+        return f
+
+    def _jaxify(self, func, abstract_eval):
+        """Make an external (python) function work with JAX.
+
+        Positional arguments to func can be differentiated,
+        use keyword args for static values and non-differentiable stuff.
+
+        Note: Only forward mode differentiation is supported currently.
+
+        Parameters
+        ----------
+        func : callable
+            Function to wrap. Should be a "pure" function, in that it has no side
+            effects and doesn't maintain state. Does not need to be JAX transformable.
+        abstract_eval : callable
+            Auxilliary function that computes the output shape and dtype of func.
+            **Must be JAX transformable**. Should be of the form
+
+                abstract_eval(*args, **kwargs) -> Pytree with same shape and dtype as
+                func(*args, **kwargs)
+
+            For example, if func always returns a scalar:
+
+                abstract_eval = lambda *args, **kwargs: jnp.array(1.)
+
+            Or if func takes an array of shape(n) and returns a dict of arrays of
+            shape(n-2):
+
+                abstract_eval = lambda arr, **kwargs:
+                {"out1": jnp.empty(arr.size-2), "out2": jnp.empty(arr.size-2)}
+
+        Returns
+        -------
+        func : callable
+            New function that behaves as func but works with jit/vmap/jacfwd etc.
+
+        """
+
+        def wrap_pure_callback(func):
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                result_shape_dtype = abstract_eval(*args, **kwargs)
+                return jax.pure_callback(func, result_shape_dtype, *args, **kwargs)
+
+            return wrapper
+
+        def define_fd_jvp(func):
+            func = jax.custom_jvp(func)
+
+            @func.defjvp
+            def func_jvp(primals, tangents):
+                primal_out = func(*primals)
+
+                # flatten everything into 1D vectors for easier finite differences
+                y, unflaty = jax.flatten_util.ravel_pytree(primal_out)
+                x, unflatx = jax.flatten_util.ravel_pytree(primals)
+                v, _______ = jax.flatten_util.ravel_pytree(tangents)
+                # scale to unit norm if nonzero
+                normv = jnp.linalg.norm(v)
+                vh = jnp.where(normv == 0, v, v / normv)
+
+                def f(x):
+                    return jax.flatten_util.ravel_pytree(func(*unflatx(x)))[0]
+
+                tangent_out = (f(x + self._fd_step * vh) - y) / self._fd_step * normv
+                tangent_out = unflaty(tangent_out)
+
+                return primal_out, tangent_out
+
+            return func
+
+        return define_fd_jvp(wrap_pure_callback(func))
+
+
 class GenericObjective(_Objective):
     """A generic objective that can compute any quantity from the `data_index`.
 
@@ -58,7 +289,7 @@
 
     """
 
-    _print_value_fmt = "GenericObjective value: {:10.3e} "
+    _print_value_fmt = "Generic objective value: {:10.3e} "
 
     def __init__(
         self,
@@ -207,7 +438,7 @@
     _linear = True
     _fixed = True
     _units = "(Unknown)"
-    _print_value_fmt = "Custom linear Objective value: {:10.3e}"
+    _print_value_fmt = "Custom linear objective value: {:10.3e}"
 
     def __init__(
         self,
@@ -345,18 +576,17 @@
         def myfun(grid, data):
             # This will compute the flux surface average of the function
             # R*B_T from the Grad-Shafranov equation
-            f = data['R']*data['B_phi']
+            f = data['R'] * data['B_phi']
             f_fsa = surface_averages(grid, f, sqrt_g=data['sqrt_g'])
-            # this has the FSA values on the full grid, but we just want
-            # the unique values:
+            # this is the FSA on the full grid, but we only want the unique values:
             return grid.compress(f_fsa)
 
         myobj = ObjectiveFromUser(myfun)
 
     """
 
     _units = "(Unknown)"
-    _print_value_fmt = "Custom Objective value: {:10.3e}"
+    _print_value_fmt = "Custom objective value: {:10.3e}"
 
     def __init__(
         self,

desc/objectives/objective_funs.py

@@ -58,6 +58,7 @@ def __init__(
 
     def _set_derivatives(self):
         """Set up derivatives of the objective functions."""
+        # TODO: does deriv_mode have to be "blocked" if there is an ExternalObjective?
         if self._deriv_mode == "auto":
             if all((obj._deriv_mode == "fwd") for obj in self.objectives):
                 self._deriv_mode = "batched"
@@ -90,9 +91,7 @@ def jac_(op, x, constants=None):
                 for obj, const in zip(self.objectives, constants):
                     # get the xs that go to that objective
                     xi = [x for x, t in zip(xs, self.things) if t in obj.things]
-                    Ji_ = getattr(obj, op)(
-                        *xi, constants=const
-                    )  # jac wrt to just those things
+                    Ji_ = getattr(obj, op)(*xi, constants=const)  # jac wrt only xi
                     Ji = []  # jac wrt all things
                     for thing in self.things:
                         if thing in obj.things: