Allow for optimizing multiple objects

desc/coils.py

@@ -16,6 +16,7 @@
 )
 from desc.grid import LinearGrid
 from desc.magnetic_fields import _MagneticField
+from desc.optimizable import Optimizable, optimizable_parameter
 from desc.utils import equals, errorif, flatten_list
 
 
@@ -110,7 +111,7 @@ def biot_savart_quad(eval_pts, coil_pts, tangents, current):
     return B
 
 
-class _Coil(_MagneticField, ABC):
+class _Coil(_MagneticField, Optimizable, ABC):
     """Base class representing a magnetic field coil.
 
     Represents coils as a combination of a Curve and current
@@ -134,6 +135,7 @@ def __init__(self, current, *args, **kwargs):
         self._current = float(current)
         super().__init__(*args, **kwargs)
 
+    @optimizable_parameter
     @property
     def current(self):
         """float: Current passing through the coil, in Amperes."""

desc/compute/__init__.py

@@ -43,7 +43,6 @@
 from .data_index import data_index
 from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec
 from .utils import (
-    arg_order,
     compute,
     get_data_deps,
     get_derivs,

desc/compute/utils.py

@@ -12,6 +12,17 @@
 
 from .data_index import data_index
 
+# map from profile name to equilibrium parameter name
+profile_names = {
+    "pressure": "p_l",
+    "iota": "i_l",
+    "current": "c_l",
+    "electron_temperature": "Te_l",
+    "electron_density": "ne_l",
+    "ion_temperature": "Ti_l",
+    "atomic_number": "Zeff_l",
+}
+
 
 def _parse_parameterization(p):
     if isinstance(p, str):
@@ -23,10 +34,6 @@ def _parse_parameterization(p):
     return module + "." + klass.__qualname__
 
 
-def _sort_args(args):
-    return [arg for arg in arg_order if arg in args]
-
-
 def compute(parameterization, names, params, transforms, profiles, data=None, **kwargs):
     """Compute the quantity given by name on grid.
 
@@ -283,10 +290,6 @@ def get_params(keys, obj, has_axis=False, **kwargs):
     params = []
     for key in deps:
         params += data_index[p][key]["dependencies"]["params"]
-    if p == "desc.equilibrium.equilibrium.Equilibrium":
-        # probably need some way to distinguish between params from different instances
-        # of the same class?
-        params = _sort_args(list(set(params)))
     if isinstance(obj, str) or inspect.isclass(obj):
         return params
     temp_params = {}

desc/continuation.py

@@ -6,7 +6,6 @@
 import numpy as np
 from termcolor import colored
 
-from desc.compute import arg_order
 from desc.equilibrium import EquilibriaFamily, Equilibrium
 from desc.objectives import get_equilibrium_objective, get_fixed_boundary_constraints
 from desc.optimize import Optimizer
@@ -75,12 +74,6 @@ def _solve_axisym(
 
     if not isinstance(optimizer, Optimizer):
         optimizer = Optimizer(optimizer)
-    constraints_i = get_fixed_boundary_constraints(
-        eq=eqi,
-        iota=objective != "vacuum" and eq.iota is not None,
-        kinetic=eq.electron_temperature is not None,
-    )
-    objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
 
     eqfam = EquilibriaFamily()
 
@@ -105,6 +98,13 @@ def _solve_axisym(
             deltas = get_deltas({"surface": surf_i}, {"surface": surf_i2})
             surf_i = surf_i2
 
+        constraints_i = get_fixed_boundary_constraints(
+            eq=eqi,
+            iota=eqi.iota,
+            kinetic=eqi.electron_temperature,
+        )
+        objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
+
         if verbose:
             _print_iteration_summary(
                 ii,
@@ -118,12 +118,6 @@ def _solve_axisym(
                 optimizer,
             )
 
-        constraints_i = get_fixed_boundary_constraints(
-            eq=eqi,
-            iota=objective != "vacuum" and eq.iota is not None,
-            kinetic=eq.electron_temperature is not None,
-        )
-        objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
         if len(deltas) > 0:
             if verbose > 0:
                 print("Perturbing equilibrium")
@@ -215,13 +209,6 @@ def _add_pressure(
     # make sure its at full radial/poloidal resolution
     eqi.change_resolution(L=eq.L, M=eq.M, L_grid=eq.L_grid, M_grid=eq.M_grid)
 
-    constraints_i = get_fixed_boundary_constraints(
-        eq=eqi,
-        iota=objective != "vacuum" and eq.iota is not None,
-        kinetic=eq.electron_temperature is not None,
-    )
-    objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
-
     pres_steps = (
         0
         if (abs(eq.pressure(np.linspace(0, 1, 20))) < 1e-14).all() or pres_step == 0
@@ -240,6 +227,13 @@ def _add_pressure(
         deltas["p_l"] *= pres_step
         pres_ratio += pres_step
 
+        constraints_i = get_fixed_boundary_constraints(
+            eq=eqi,
+            iota=eqi.iota,
+            kinetic=eqi.electron_temperature,
+        )
+        objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
+
         if verbose:
             _print_iteration_summary(
                 ii,
@@ -347,13 +341,6 @@ def _add_shaping(
     # make sure its at full resolution
     eqi.change_resolution(eq.L, eq.M, eq.N, eq.L_grid, eq.M_grid, eq.N_grid)
 
-    constraints_i = get_fixed_boundary_constraints(
-        eq=eqi,
-        iota=objective != "vacuum" and eq.iota is not None,
-        kinetic=eq.electron_temperature is not None,
-    )
-    objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
-
     bdry_steps = 0 if eq.N == 0 or bdry_step == 0 else int(np.ceil(1 / bdry_step))
     bdry_ratio = 0 if eq.N else 1
 
@@ -373,6 +360,13 @@ def _add_shaping(
             deltas["Zb_lmn"] *= bdry_step
         bdry_ratio += bdry_step
 
+        constraints_i = get_fixed_boundary_constraints(
+            eq=eqi,
+            iota=eqi.iota,
+            kinetic=eqi.electron_temperature,
+        )
+        objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
+
         if verbose:
             _print_iteration_summary(
                 ii,
@@ -578,10 +572,7 @@ def solve_continuation_automatic(  # noqa: C901
         verbose,
         checkpoint_path,
     )
-    for arg in arg_order:
-        val = np.asarray(getattr(eqfam[-1], arg))
-        if val.size:
-            setattr(eq, arg, val)
+    eq.params_dict = eqfam[-1].params_dict
     eqfam[-1] = eq
     timer.stop("Total time")
     if verbose > 0:
@@ -670,9 +661,9 @@ def solve_continuation(  # noqa: C901
 
     if not isinstance(optimizer, Optimizer):
         optimizer = Optimizer(optimizer)
-    objective_i = get_equilibrium_objective(eq=eqfam[-1], mode=objective)
+    objective_i = get_equilibrium_objective(eq=eqfam[0], mode=objective)
     constraints_i = get_fixed_boundary_constraints(
-        eq=eqfam[-1],
+        eq=eqfam[0],
         iota=objective != "vacuum" and eqfam[0].iota is not None,
         kinetic=eqfam[0].electron_temperature is not None,
     )
@@ -716,20 +707,17 @@ def solve_continuation(  # noqa: C901
             }
             deltas = get_deltas(things1, things2)
 
-            # maybe rebuild objective if resolution changed.
-            if eqfam[ii - 1].resolution != eqi.resolution:
-                objective_i = get_equilibrium_objective(eq=eqfam[ii], mode=objective)
-                constraints_i = get_fixed_boundary_constraints(
-                    eq=eqfam[ii],
-                    iota=objective != "vacuum" and eqfam[ii].iota is not None,
-                    kinetic=eqfam[ii].electron_temperature is not None,
-                )
-
         if len(deltas) > 0:
             if verbose > 0:
                 print("Perturbing equilibrium")
             # TODO: pass Jx if available
             eqp = eqfam[ii - 1].copy()
+            objective_i = get_equilibrium_objective(eq=eqp, mode=objective)
+            constraints_i = get_fixed_boundary_constraints(
+                eq=eqp,
+                iota=eqp.iota,
+                kinetic=eqp.electron_temperature,
+            )
             eqp.change_resolution(**eqi.resolution)
             eqp.perturb(
                 objective=objective_i,
@@ -739,17 +727,21 @@ def solve_continuation(  # noqa: C901
                 verbose=verbose,
                 copy=False,
             )
-            for arg in arg_order:
-                val = np.asarray(getattr(eqp, arg))
-                if val.size:
-                    setattr(eqi, arg, val)
+            eqi.params_dict = eqp.params_dict
             deltas = {}
             del eqp
 
         if not eqi.is_nested(msg="manual"):
             stop = True
 
         if not stop:
+            # TODO: add ability to rebind objectives
+            objective_i = get_equilibrium_objective(eq=eqi, mode=objective)
+            constraints_i = get_fixed_boundary_constraints(
+                eq=eqi,
+                iota=eqi.iota,
+                kinetic=eqi.electron_temperature,
+            )
             eqi.solve(
                 optimizer=optimizer,
                 objective=objective_i,