Dd/optimizable

desc/objectives/__init__.py

@@ -49,6 +49,7 @@
     FixAxisZ,
     FixBoundaryR,
     FixBoundaryZ,
+    FixCollectionParameters,
     FixCurrent,
     FixCurveRotation,
     FixCurveShift,

desc/objectives/_free_boundary.py

@@ -481,12 +481,11 @@ def build(self, use_jit=True, verbose=1):
         if self._source_grid is None:
             # for axisymmetry we still need to know about toroidal effects, so its
             # cheapest to pretend there are extra field periods
-            source_NFP = eq.NFP if eq.N > 0 else 64
             source_grid = LinearGrid(
                 rho=np.array([1.0]),
                 M=eq.M_grid,
                 N=eq.N_grid,
-                NFP=source_NFP,
+                NFP=int(eq.NFP) if eq.N > 0 else 64,
                 sym=False,
             )
         else:

desc/objectives/getters.py

@@ -1,5 +1,7 @@
 """Utilities for getting standard groups of objectives and constraints."""
 
+from desc.utils import flatten_list, unique_list
+
 from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance
 from .linear_objectives import (
     AxisRSelfConsistency,
@@ -213,12 +215,8 @@
         return any([isinstance(t, cls) for t in things])
 
     # Equilibrium
-    if (
-        hasattr(thing, "Ra_n")
-        and hasattr(thing, "Za_n")
-        and hasattr(thing, "Rb_lmn")
-        and hasattr(thing, "Zb_lmn")
-        and hasattr(thing, "L_lmn")
+    if {"Rb_lmn", "Zb_lmn", "L_lmn", "Ra_n", "Za_n"} <= set(
+        unique_list(flatten_list(thing.optimizable_params))[0]
     ):
         if not _is_any_instance(constraints, BoundaryRSelfConsistency):
             constraints += (BoundaryRSelfConsistency(eq=thing),)
@@ -232,7 +230,10 @@
             constraints += (AxisZSelfConsistency(eq=thing),)
 
     # Curve
-    elif hasattr(thing, "shift") and hasattr(thing, "rotmat"):
+    # FIXME: make this work for CoilSet
+    elif not hasattr(thing, "__len__") and {"shift", "rotmat"} <= set(
+        unique_list(flatten_list(thing.optimizable_params))[0]
+    ):
         if not _is_any_instance(constraints, FixCurveShift):
             constraints += (FixCurveShift(curve=thing),)
         if not _is_any_instance(constraints, FixCurveRotation):

desc/objectives/linear_objectives.py

@@ -14,7 +14,7 @@
 
 from desc.backend import jnp
 from desc.basis import zernike_radial, zernike_radial_coeffs
-from desc.utils import errorif, setdefault
+from desc.utils import errorif, flatten_list, setdefault, unique_list
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective
@@ -110,6 +110,7 @@
         normalize_target=False,
         name="Fixed parameter",
     ):
+        # TODO: assert that `thing` is not of type `OptimizableCollection`
         self._target_from_user = target
         self._params = params
         self._indices = indices
@@ -160,8 +161,8 @@
         errorif(
             len(self._params) != len(self._indices),
             ValueError,
-            f"not enough indices ({len(self._indices)}) "
-            + f"for params ({len(self._params)})",
+            f"Unequal number of indices ({len(self._indices)}) "
+            + f"and params ({len(self._params)}).",
         )
         for idx, par in zip(self._indices, self._params):
             if isinstance(idx, bool) and idx:
@@ -210,6 +211,202 @@
         )
 
 
+class FixCollectionParameters(_FixedObjective):
+    """Fix specific degrees of freedom associated with a given OptimizableCollection.
+
+    Parameters
+    ----------
+    thing : OptimizableCollection
+        Object whose degrees of freedom are being fixed.
+    params : str or list of str
+        Names of parameters to fix. Defaults to all parameters.
+    index : array-like or list of array-like
+        Indices to fix for each parameter in params. Use True to fix all indices.
+    target : dict of {float, ndarray}, optional
+        Target value(s) of the objective. Only used if bounds is None.
+        Should have the same tree structure as thing.params. Defaults to things.params.
+    bounds : tuple of dict {float, ndarray}, optional
+        Lower and upper bounds on the objective. Overrides target.
+        Should have the same tree structure as thing.params.
+    weight : dict of {float, ndarray}, optional
+        Weighting to apply to the Objective, relative to other Objectives.
+        Should be a scalar or have the same tree structure as thing.params.
+    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. Has no effect for this objective.
+    name : str, optional
+        Name of the objective function.
+
+    """
+
+    _scalar = False
+    _linear = True
+    _fixed = True
+    _units = "(~)"
+    _print_value_fmt = "Fixed parameters error: {:10.3e} "
+
+    def __init__(
+        self,
+        thing,
+        params=None,
+        indices=True,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=False,
+        normalize_target=False,
+        name="Fixed parameters",
+    ):
+        errorif(
+            not hasattr(thing, "__len__"),
+            ValueError,
+            f"Thing must be of type `OptimizableCollection`; got {thing}.",
+        )
+        self._target_from_user = target
+        self._params = params
+        self._indices = indices
+        super().__init__(
+            things=thing,
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            name=name,
+        )
+
+    def build(self, use_jit=False, 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.
+
+        """
+        thing = self.things[0]
+        thing_idx = range(len(thing))  # indices of things in OptimizableCollection
+        params = setdefault(self._params, thing.optimizable_params)
+
+        if not isinstance(params, (list, tuple)):
+            params = [params]
+        if not all([isinstance(par, (list, tuple)) for par in params]):
+            params = [
+                params if set(params) <= set(thing[k].optimizable_params) else []
+                for k in thing_idx
+            ]
+        for k in thing_idx:
+            for par in params[k]:
+                errorif(
+                    par not in unique_list(flatten_list(thing.optimizable_params))[0],
+                    ValueError,
+                    f"Parameter {par} not found in optimizable_parameters: "
+                    + f"{thing.optimizable_params}.",
+                )
+        self._params = params
+
+        # replace indices=True with actual indices
+        if isinstance(self._indices, bool) and self._indices:
+            self._indices = [
+                [
+                    np.arange(thing.dimensions[k][par])
+                    for par in self._params[k]
+                    if par in thing.dimensions[k].keys()
+                ]
+                for k in thing_idx
+            ]
+        # make sure its iterable if only a scalar was passed in
+        if not isinstance(self._indices, (list, tuple)):
+            self._indices = [self._indices]
+        # replace idx=True with array of all indices, throwing an error if the length
+        # of indices is different from number of params
+        errorif(
+            len(sum(self._params, [])) != len(sum(self._indices, [])),
+            ValueError,
+            f"Unequal number of indices ({len(sum(self._indices, []))}) "
+            + f"and params ({len(sum(self._params, []))}).",
+        )
+        indices = []
+        for k in thing_idx:
+            indices.append({})
+            for idx, par in zip(self._indices[k], self._params[k]):
+                if isinstance(idx, bool) and idx:
+                    idx = np.arange(thing.dimensions[k][par])
+                indices[k][par] = np.atleast_1d(idx)
+        self._indices = indices
+        self._dim_f = sum(
+            sum(t.size for t in self._indices[k].values()) for k in thing_idx
+        )
+
+        # FIXME: I don't think custom target/bounds works yet (default target is ok)
+        default_target = [
+            {par: thing.params_dict[k][par][self._indices[k][par]] for par in params[k]}
+            for k in thing_idx
+        ]
+        default_bounds = None
+        target, bounds = self._parse_target_from_user(
+            self._target_from_user, default_target, default_bounds, indices
+        )
+        if target:
+            self.target = jnp.concatenate(
+                [
+                    (
+                        jnp.concatenate([target[k][par] for par in params[k]])
+                        if par in target[k].keys()
+                        else jnp.array([])
+                    )
+                    for k in thing_idx
+                ]
+            )
+            self.bounds = None
+        else:
+            self.target = None
+            self.bounds = (
+                jnp.concatenate([bounds[0][par] for par in params]),
+                jnp.concatenate([bounds[1][par] for par in params]),
+            )
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, params, constants=None):
+        """Compute fixed degree of freedom errors.
+
+        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
+            Fixed degree of freedom errors.
+
+        """
+        return jnp.concatenate(
+            [
+                (
+                    jnp.concatenate(
+                        [
+                            params[k][par][self._indices[k][par]]
+                            for par in self._params[k]
+                        ]
+                    )
+                    if len(self._params[k])
+                    else jnp.array([])
+                )
+                for k in range(len(self._params))
+            ]
+        )
+
+
 class BoundaryRSelfConsistency(_Objective):
     """Ensure that the boundary and interior surfaces are self-consistent.
 

desc/optimizable.py

@@ -4,6 +4,8 @@
 import warnings
 from abc import ABC
 
+import numpy as np
+
 from desc.backend import jnp
 
 
@@ -86,6 +88,7 @@ def pack_params(self, p):
         x : ndarray
             optimizable parameters concatenated into a single array, with indices
             given by ``x_idx``
+
         """
         return jnp.concatenate(
             [jnp.atleast_1d(jnp.asarray(p[key])) for key in self.optimizable_params]
@@ -104,6 +107,7 @@ def unpack_params(self, x):
         -------
         p : dict
             Dictionary of ndarray of optimizable parameters.
+
         """
         x_idx = self.x_idx
         params = {}
@@ -115,7 +119,8 @@ def _sort_args(self, args):
         """Put arguments in a canonical order. Returns unique sorted elements.
 
         Actual order doesn't really matter as long as its consistent, though subclasses
-        may override this method to enforce a specific ordering
+        may override this method to enforce a specific ordering.
+
         """
         return sorted(set(list(args)))
 
@@ -177,6 +182,7 @@ def pack_params(self, params):
         x : ndarray
             optimizable parameters concatenated into a single array, with indices
             given by ``x_idx``
+
         """
         return jnp.concatenate([s.pack_params(p) for s, p in zip(self, params)])
 
@@ -193,8 +199,9 @@ def unpack_params(self, x):
         -------
         p : list dict
             list of dictionary of ndarray of optimizable parameters.
+
         """
-        split_idx = jnp.cumsum(jnp.array([s.dim_x for s in self]))
+        split_idx = np.cumsum([s.dim_x for s in self])  # must be np not jnp
         xs = jnp.split(x, split_idx)
         params = [s.unpack_params(xi) for s, xi in zip(self, xs)]
         return params

tests/test_examples.py

@@ -17,7 +17,12 @@
 from desc.geometry import FourierRZToroidalSurface
 from desc.grid import LinearGrid
 from desc.io import load
-from desc.magnetic_fields import OmnigenousField, SplineMagneticField
+from desc.magnetic_fields import (
+    OmnigenousField,
+    SplineMagneticField,
+    ToroidalMagneticField,
+    VerticalMagneticField,
+)
 from desc.objectives import (
     AspectRatio,
     BoundaryError,
@@ -27,6 +32,7 @@
     CurrentDensity,
     FixBoundaryR,
     FixBoundaryZ,
+    FixCollectionParameters,
     FixCurrent,
     FixIota,
     FixOmniBmax,
@@ -1344,3 +1350,23 @@ def test_single_coil_optimization():
     np.testing.assert_allclose(
         coil.compute("torsion", grid=grid)["torsion"], target, atol=1e-5
     )
+
+
+@pytest.mark.unit
+def test_second_stage_optimization():
+    """Test optimizing magnetic field for a fixed axisymmetric equilibrium."""
+    # This also tests that FixCollectionParameters works properly when fixing a
+    # parameter that does not exist for all things in the collection.
+
+    # TODO: change the objective to QuadraticFlux
+    eq = get("DSHAPE")
+    field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
+    objective = ObjectiveFunction(BoundaryError(eq=eq, field=field))
+    constraints = (FixParameter(eq), FixCollectionParameters(field, "R0"))
+    optimizer = Optimizer("lsq-exact")
+    (eq, field), _ = optimizer.optimize(
+        things=(eq, field), objective=objective, constraints=constraints
+    )
+    np.testing.assert_allclose(field[0].R0, 3.5)
+    np.testing.assert_allclose(field[0].B0, 0.218, rtol=1e-3)  # toroidal field
+    np.testing.assert_allclose(field[1].B0, -0.021, rtol=2e-2)  # vertical field