ExternalObjective function to wrap external codes

desc/backend.py

@@ -1,6 +1,7 @@
 """Backend functions for DESC, with options for JAX or regular numpy."""
 
 import functools
+import multiprocessing as mp
 import os
 import warnings
 
@@ -11,15 +12,19 @@
 from desc import config as desc_config
 from desc import set_device
 
+# only print details in the main process, not child processes spawned by multiprocessing
+verbose = bool(mp.current_process().name == "MainProcess")
+
 if os.environ.get("DESC_BACKEND") == "numpy":
     jnp = np
     use_jax = False
     set_device(kind="cpu")
-    print(
-        "DESC version {}, using numpy backend, version={}, dtype={}".format(
-            desc.__version__, np.__version__, np.linspace(0, 1).dtype
+    if verbose:
+        print(
+            "DESC version {}, using numpy backend, version={}, dtype={}".format(
+                desc.__version__, np.__version__, np.linspace(0, 1).dtype
+            )
         )
-    )
 else:
     if desc_config.get("device") is None:
         set_device("cpu")
@@ -41,11 +46,12 @@
             x = jnp.linspace(0, 5)
             y = jnp.exp(x)
         use_jax = True
-        print(
-            f"DESC version {desc.__version__},"
-            + f"using JAX backend, jax version={jax.__version__}, "
-            + f"jaxlib version={jaxlib.__version__}, dtype={y.dtype}"
-        )
+        if verbose:
+            print(
+                f"DESC version {desc.__version__}, "
+                + f"using JAX backend, jax version={jax.__version__}, "
+                + f"jaxlib version={jaxlib.__version__}, dtype={y.dtype}"
+            )
         del x, y
     except ModuleNotFoundError:
         jnp = np
@@ -59,11 +65,13 @@
                 desc.__version__, np.__version__, y.dtype
             )
         )
-print(
-    "Using device: {}, with {:.2f} GB available memory".format(
-        desc_config.get("device"), desc_config.get("avail_mem")
+
+if verbose:
+    print(
+        "Using device: {}, with {:.2f} GB available memory".format(
+            desc_config.get("device"), desc_config.get("avail_mem")
+        )
     )
-)
 
 if use_jax:  # noqa: C901 - FIXME: simplify this, define globally and then assign?
     jit = jax.jit
@@ -515,7 +523,7 @@
             val = body_fun(i, val)
         return val
 
-    def cond(pred, true_fun, false_fun, *operand):
+    def cond(pred, true_fun, false_fun, *operands):
         """Conditionally apply true_fun or false_fun.
 
         This version is for the numpy backend, for jax backend see jax.lax.cond
@@ -528,7 +536,7 @@
             Function (A -> B), to be applied if pred is True.
         false_fun: callable
             Function (A -> B), to be applied if pred is False.
-        operand: any
+        operands: any
             input to either branch depending on pred. The type can be a scalar, array,
             or any pytree (nested Python tuple/list/dict) thereof.
 
@@ -541,9 +549,9 @@
 
         """
         if pred:
-            return true_fun(*operand)
+            return true_fun(*operands)
         else:
-            return false_fun(*operand)
+            return false_fun(*operands)
 
     def switch(index, branches, operand):
         """Apply exactly one of branches given by index.

desc/compute/utils.py

@@ -1113,7 +1113,7 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True, tol=1e-14)
             has_endpoint_dupe,
             lambda _: put(mask, jnp.array([0, -1]), mask[0] | mask[-1]),
             lambda _: mask,
-            operand=None,
+            None,
         )
     else:
         # If we don't have the idx attributes, we are forced to expand out.

desc/objectives/__init__.py

@@ -20,7 +20,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

@@ -2,6 +2,7 @@
 
 import inspect
 import re
+from abc import ABC
 
 import numpy as np
 
@@ -11,12 +12,178 @@
 from desc.compute.utils import _parse_parameterization, get_profiles, get_transforms
 from desc.grid import QuadratureGrid
 from desc.optimizable import OptimizableCollection
-from desc.utils import errorif, parse_argname_change
+from desc.utils import errorif, jaxify, parse_argname_change
 
 from .linear_objectives import _FixedObjective
 from .objective_funs import _Objective
 
 
+class ExternalObjective(_Objective, ABC):
+    """Wrap an external code.
+
+    Similar to ``ObjectiveFromUser``, except derivatives of the objective function are
+    computed with finite differences instead of AD. The function does not need not be
+    JAX transformable.
+
+    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
+        External objective function. It must take an Equilibrium as its only positional
+        argument, but can take additional kewyord arguments. It does not need to be JAX
+        transformable.
+    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.
+    vectorized : bool, optional
+        Whether or not ``fun`` is vectorized. Default = False.
+    abs_step : float, optional
+        Absolute finite difference step size. Default = 1e-4.
+        Total step size is ``abs_step + rel_step * mean(abs(x))``.
+    rel_step : float, optional
+        Relative finite difference step size. Default = 0.
+        Total step size is ``abs_step + rel_step * mean(abs(x))``.
+    name : str, optional
+        Name of the objective function.
+    kwargs : any, optional
+        Keyword arguments that are passed as inputs to ``fun``.
+
+    # 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,
+        vectorized=False,
+        abs_step=1e-4,
+        rel_step=0,
+        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._vectorized = vectorized
+        self._abs_step = abs_step
+        self._rel_step = rel_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(params):
+            """Wrap external function with possibly vectorized params."""
+            # number of equilibria for vectorized computations
+            param_shape = params["Psi"].shape
+            num_eq = param_shape[0] if len(param_shape) > 1 else 1
+
+            # convert params to list of equilibria
+            eqs = [self._eq.copy() for _ in range(num_eq)]
+            for k, eq in enumerate(eqs):
+                # update equilibria with new params
+                for param_key in self._eq.optimizable_params:
+                    param_value = np.atleast_2d(params[param_key])[k, :]
+                    if len(param_value):
+                        setattr(eq, param_key, param_value)
+
+            # call external function on equilibrium or list of equilibria
+            if not self._vectorized:
+                eqs = eqs[0]
+            return self._fun(eqs, **self._kwargs)
+
+        # wrap external function to work with JAX
+        abstract_eval = lambda *args, **kwargs: jnp.empty(self._dim_f)
+        self._fun_wrapped = jaxify(
+            fun_wrapped,
+            abstract_eval,
+            vectorized=self._vectorized,
+            abs_step=self._abs_step,
+            rel_step=self._rel_step,
+        )
+
+        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.
+
+        """
+        f = self._fun_wrapped(params)
+        return f
+
+
 class GenericObjective(_Objective):
     """A generic objective that can compute any quantity from the `data_index`.
 
@@ -352,10 +519,9 @@
         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(fun=myfun, thing=eq)
@@ -414,6 +580,8 @@
             Level of output.
 
         """
+        import jax
+
         thing = self.things[0]
         if self._grid is None:
             errorif(
@@ -444,7 +612,6 @@
                 ).squeeze()
 
         self._fun_wrapped = lambda data: self._fun(grid, data)
-        import jax
 
         self._dim_f = jax.eval_shape(self._fun_wrapped, dummy_data).size
         self._scalar = self._dim_f == 1

desc/objectives/objective_funs.py

@@ -90,9 +90,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: