ExternalObjective function to wrap external codes

desc/backend.py

@@ -1,5 +1,6 @@
 """Backend functions for DESC, with options for JAX or regular numpy."""
 
+import multiprocessing
 import os
 import warnings
 
@@ -10,15 +11,23 @@
 from desc import config as desc_config
 from desc import set_device
 
+verbose = True
+
+# set child processes to use numpy backend and suppress print statements
+if not multiprocessing.current_process().name == "MainProcess":
+    os.environ["DESC_BACKEND"] = "numpy"
+    verbose = False
+
 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")
@@ -40,11 +49,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
@@ -58,11 +68,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
@@ -489,7 +501,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
@@ -502,7 +514,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.
 
@@ -515,9 +527,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

@@ -927,7 +927,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:
         expand_out = False

desc/objectives/__init__.py

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