From 6402ab07313cb85f957139cacdb864a57d5a167c Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 25 Mar 2024 17:31:30 -0600
Subject: [PATCH 01/50] create FixCollectionParameters objective

---
 desc/objectives/__init__.py          |   1 +
 desc/objectives/linear_objectives.py | 175 ++++++++++++++++++++++++++-
 desc/optimizable.py                  |   7 +-
 3 files changed, 180 insertions(+), 3 deletions(-)

diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py
index 25d51ce089..efaa9784d1 100644
--- a/desc/objectives/__init__.py
+++ b/desc/objectives/__init__.py
@@ -48,6 +48,7 @@
     FixAxisZ,
     FixBoundaryR,
     FixBoundaryZ,
+    FixCollectionParameters,
     FixCurrent,
     FixCurveRotation,
     FixCurveShift,
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index d57181abed..52e9fac4c6 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -110,6 +110,7 @@ def __init__(
         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 @@ def build(self, use_jit=False, verbose=1):
         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,176 @@ def compute(self, params, constants=None):
         )
 
 
+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",
+    ):
+        # TODO: assert that `thing` is of type `OptimizableCollection`
+        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]
+        params = setdefault(self._params, thing.optimizable_params)
+
+        if not isinstance(params, (list, tuple)):
+            params = [params]
+        for par in params:
+            errorif(
+                par not in thing.optimizable_params,  # FIXME: can't search nested list
+                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]]
+                for k in range(len(thing))
+            ]
+        # 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(self._indices)}) "
+            + f"and params ({len(self._params)}).",
+        )
+        indices = []
+        for k in range(len(thing)):
+            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(
+            t.size for t in self._indices[k].values() for k in range(len(thing))
+        )
+
+        # 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 range(len(thing))
+        ]
+        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]])
+                    for k in range(len(thing))
+                ]
+            )
+            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]]
+                )
+                for k in range(len(self._params))
+            ]
+        )
+
+
 class BoundaryRSelfConsistency(_Objective):
     """Ensure that the boundary and interior surfaces are self-consistent.
 
diff --git a/desc/optimizable.py b/desc/optimizable.py
index f27ff53239..b2691e25a4 100644
--- a/desc/optimizable.py
+++ b/desc/optimizable.py
@@ -86,6 +86,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 +105,7 @@ def unpack_params(self, x):
         -------
         p : dict
             Dictionary of ndarray of optimizable parameters.
+
         """
         x_idx = self.x_idx
         params = {}
@@ -115,7 +117,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 +180,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,6 +197,7 @@ 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]))
         xs = jnp.split(x, split_idx)

From 9dc9a38260508d2bc27ce5745cf5129ee73485b0 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 26 Mar 2024 11:27:28 -0600
Subject: [PATCH 02/50] allow fixing of non-default params

---
 desc/objectives/linear_objectives.py | 38 ++++++++++++++++------------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 52e9fac4c6..9b6260701e 100644
--- a/desc/objectives/linear_objectives.py
+++ b/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
@@ -261,7 +261,11 @@ def __init__(
         normalize_target=False,
         name="Fixed parameters",
     ):
-        # TODO: assert that `thing` is of type `OptimizableCollection`
+        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
@@ -287,24 +291,28 @@ def build(self, use_jit=False, verbose=1):
 
         """
         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]
-        for par in params:
-            errorif(
-                par not in thing.optimizable_params,  # FIXME: can't search nested list
-                ValueError,
-                f"parameter {par} not found in optimizable_parameters: "
-                + f"{thing.optimizable_params}",
-            )
+        if not all([isinstance(par, (list, tuple)) for par in params]):
+            params = [params for _ 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]]
-                for k in range(len(thing))
+                for k in thing_idx
             ]
         # make sure its iterable if only a scalar was passed in
         if not isinstance(self._indices, (list, tuple)):
@@ -318,21 +326,19 @@ def build(self, use_jit=False, verbose=1):
             + f"and params ({len(self._params)}).",
         )
         indices = []
-        for k in range(len(thing)):
+        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(
-            t.size for t in self._indices[k].values() for k in range(len(thing))
-        )
+        self._dim_f = 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 range(len(thing))
+            for k in thing_idx
         ]
         default_bounds = None
         target, bounds = self._parse_target_from_user(
@@ -342,7 +348,7 @@ def build(self, use_jit=False, verbose=1):
             self.target = jnp.concatenate(
                 [
                     jnp.concatenate([target[k][par] for par in params[k]])
-                    for k in range(len(thing))
+                    for k in thing_idx
                 ]
             )
             self.bounds = None

From 9a78600edba03e6f9c184a2b289f5af2b7518a54 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 26 Mar 2024 13:20:24 -0600
Subject: [PATCH 03/50] bug fix: use np instead of jnp

---
 desc/optimizable.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/desc/optimizable.py b/desc/optimizable.py
index b2691e25a4..032eb267c1 100644
--- a/desc/optimizable.py
+++ b/desc/optimizable.py
@@ -4,6 +4,8 @@
 import warnings
 from abc import ABC
 
+import numpy as np
+
 from desc.backend import jnp
 
 
@@ -199,7 +201,7 @@ def unpack_params(self, x):
             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(np.array([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

From b0df4e422920c516f44b656ed8eacd317298c2f3 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 26 Mar 2024 13:58:22 -0600
Subject: [PATCH 04/50] update maybe_add_self_consistency

---
 desc/objectives/getters.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index 3b7ccd556a..dcea86e7c0 100644
--- a/desc/objectives/getters.py
+++ b/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 @@ def _is_any_instance(things, cls):
         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 @@ def _is_any_instance(things, cls):
             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):

From bd1eedc40b3229453187cac88b926a00ff8f3aaa Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 26 Mar 2024 15:57:11 -0600
Subject: [PATCH 05/50] fix NFP int bug

---
 desc/objectives/_free_boundary.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py
index d6f7ad6c19..274fd3ee8e 100644
--- a/desc/objectives/_free_boundary.py
+++ b/desc/objectives/_free_boundary.py
@@ -462,12 +462,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:

From 03dfe05eeb709e647b9c7dcc8d496b2beb1deae9 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 29 Mar 2024 13:22:34 -0600
Subject: [PATCH 06/50] allow fixing params for only some things in collection

---
 desc/objectives/linear_objectives.py | 36 +++++++++++++++++++++-------
 desc/optimizable.py                  |  2 +-
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 9b6260701e..9fcabdaa97 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -297,7 +297,10 @@ def build(self, use_jit=False, verbose=1):
         if not isinstance(params, (list, tuple)):
             params = [params]
         if not all([isinstance(par, (list, tuple)) for par in params]):
-            params = [params for _ in thing_idx]
+            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(
@@ -311,7 +314,11 @@ def build(self, use_jit=False, verbose=1):
         # 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]]
+                [
+                    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
@@ -322,8 +329,8 @@ def build(self, use_jit=False, verbose=1):
         errorif(
             len(sum(self._params, [])) != len(sum(self._indices, [])),
             ValueError,
-            f"Unequal number of indices ({len(self._indices)}) "
-            + f"and params ({len(self._params)}).",
+            f"Unequal number of indices ({len(sum(self._indices, []))}) "
+            + f"and params ({len(sum(self._params, []))}).",
         )
         indices = []
         for k in thing_idx:
@@ -333,7 +340,9 @@ def build(self, use_jit=False, verbose=1):
                     idx = np.arange(thing.dimensions[k][par])
                 indices[k][par] = np.atleast_1d(idx)
         self._indices = indices
-        self._dim_f = sum(t.size for t in self._indices[k].values() for k in thing_idx)
+        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 = [
@@ -347,7 +356,11 @@ def build(self, use_jit=False, verbose=1):
         if target:
             self.target = jnp.concatenate(
                 [
-                    jnp.concatenate([target[k][par] for par in params[k]])
+                    (
+                        jnp.concatenate([target[k][par] for par in params[k]])
+                        if par in target[k].keys()
+                        else jnp.array([])
+                    )
                     for k in thing_idx
                 ]
             )
@@ -379,8 +392,15 @@ def compute(self, params, constants=None):
         """
         return jnp.concatenate(
             [
-                jnp.concatenate(
-                    [params[k][par][self._indices[k][par]] for par in self._params[k]]
+                (
+                    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))
             ]
diff --git a/desc/optimizable.py b/desc/optimizable.py
index 032eb267c1..4427f00518 100644
--- a/desc/optimizable.py
+++ b/desc/optimizable.py
@@ -201,7 +201,7 @@ def unpack_params(self, x):
             list of dictionary of ndarray of optimizable parameters.
 
         """
-        split_idx = np.cumsum(np.array([s.dim_x for s in self]))  # must be np not jnp
+        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

From 19a895560bc48b7e76ce4908c108fe4c1b8f8a8b Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 29 Mar 2024 13:34:39 -0600
Subject: [PATCH 07/50] add test for second stage optimization

---
 tests/test_examples.py | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index 9acdbe597c..51ed5aba92 100644
--- a/tests/test_examples.py
+++ b/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

From bfb5f610efbe8303dd22a16e8a5aa2db5ed6dac4 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 12 Apr 2024 11:51:47 -0500
Subject: [PATCH 08/50] update tests

---
 tests/test_examples.py          | 14 +++++++-------
 tests/test_linear_objectives.py | 23 +++++++++++++++++++++++
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index d343f83c7d..2a962c26df 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1286,15 +1286,15 @@ def test_second_stage_optimization():
     # 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"))
+    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
+    constraints = FixCollectionParameters(field, "R0")
     optimizer = Optimizer("lsq-exact")
-    (eq, field), _ = optimizer.optimize(
-        things=(eq, field), objective=objective, constraints=constraints
+    (field,), _ = optimizer.optimize(
+        things=field, objective=objective, constraints=constraints, verbose=2
     )
+    # TODO: could make this more robust instead of assuming only vertical field changes
     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
+    np.testing.assert_allclose(field[0].B0, 1)  # toroidal field (does not change)
+    np.testing.assert_allclose(field[1].B0, -0.022, rtol=1e-2)  # vertical field
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index cecbf2bbe4..0ec081f2fc 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -6,6 +6,7 @@
 from qsc import Qsc
 
 import desc.examples
+from desc.coils import CoilSet, FourierPlanarCoil, FourierRZCoil, MixedCoilSet
 from desc.equilibrium import Equilibrium
 from desc.geometry import FourierRZToroidalSurface
 from desc.grid import LinearGrid
@@ -22,6 +23,7 @@
     FixAxisZ,
     FixBoundaryR,
     FixBoundaryZ,
+    FixCollectionParameters,
     FixCurrent,
     FixElectronDensity,
     FixElectronTemperature,
@@ -918,3 +920,24 @@ def test_fix_omni_indices():
     constraint = FixOmniMap(field=field, indices=indices)
     constraint.build()
     assert constraint._idx.size == indices.size
+
+
+@pytest.mark.unit
+def test_fix_subset_of_params_in_collection(self):
+    """Tests FixCollectionParameters fixing a subset of things in the collection."""
+    tf_coil = FourierPlanarCoil(center=[2, 0, 0], normal=[0, 1, 0], r_n=[1])
+    tf_coilset = CoilSet.linspaced_angular(tf_coil, n=4)
+    vf_coil = FourierRZCoil(R_n=3, Z_n=-1)
+    vf_coilset = CoilSet.linspaced_linear(
+        vf_coil, displacement=[0, 0, 2], n=3, endpoint=True
+    )
+    full_coilset = MixedCoilSet((tf_coilset, vf_coilset))
+
+    params = [
+        [["current"], ["center", "normal", "r_n"], ["current", "shift", "rotmat"], []],
+        [["current", "Z_n"], ["R_n", "Z_n"], ["Z_n"]],
+    ]
+    obj = FixCollectionParameters(full_coilset, params=params)
+    if False:  # FIXME: get this test working
+        obj.build()
+    pass

From 8db2c3ff8cc3bb5b40f8e8335a7b9dfab63a0a8a Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 19 Apr 2024 10:47:55 -0500
Subject: [PATCH 09/50] broadcast_tree util function

---
 desc/backend.py     |  5 +++++
 desc/utils.py       | 46 +++++++++++++++++++++++++++++++++++++++++-
 tests/test_utils.py | 49 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/desc/backend.py b/desc/backend.py
index 721920190c..77cf6f090d 100644
--- a/desc/backend.py
+++ b/desc/backend.py
@@ -84,6 +84,7 @@
         tree_map,
         tree_structure,
         tree_unflatten,
+        treedef_is_leaf,
     )
 
     def put(arr, inds, vals):
@@ -412,6 +413,10 @@ def tree_leaves(*args, **kwargs):
         """Get leaves of pytree for numpy backend."""
         raise NotImplementedError
 
+    def treedef_is_leaf(*args, **kwargs):
+        """Check is leaf of pytree for numpy backend."""
+        raise NotImplementedError
+
     def register_pytree_node(foo, *args):
         """Dummy decorator for non-jax pytrees."""
         return foo
diff --git a/desc/utils.py b/desc/utils.py
index 710532d432..d6d72a726e 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -8,7 +8,7 @@
 from scipy.special import factorial
 from termcolor import colored
 
-from desc.backend import fori_loop, jit, jnp
+from desc.backend import fori_loop, jit, jnp, tree_structure, treedef_is_leaf
 
 
 class Timer:
@@ -608,3 +608,47 @@ def unique_list(thelist):
 def is_any_instance(things, cls):
     """Check if any of things is an instance of cls."""
     return any([isinstance(t, cls) for t in things])
+
+
+def broadcast_tree(tree_in, tree_out):
+    """Broadcast tree_in to the same pytree structure as tree_out."""
+    errorif(
+        isinstance(tree_in, (tuple, dict)),
+        ValueError,
+        "tree_in must be a pytree of nested lists, not tuples or dicts",
+    )
+    errorif(
+        isinstance(tree_out, (tuple, dict)),
+        ValueError,
+        "tree_out must be a pytree of nested lists, not tuples or dicts",
+    )
+    if not isinstance(tree_in, list):
+        tree_in = [tree_in]
+    if not isinstance(tree_out, list):
+        tree_out = [tree_out]
+    where_leaves_in = [treedef_is_leaf(tree_structure(branch)) for branch in tree_in]
+    where_leaves_out = [treedef_is_leaf(tree_structure(branch)) for branch in tree_out]
+    all_leaves_in = all(where_leaves_in)
+    all_leaves_out = all(where_leaves_out)
+    if any(where_leaves_in) and not all_leaves_in:
+        raise ValueError("base layer of tree_in must be all leaves and no branches")
+    if any(where_leaves_out) and not all_leaves_out:
+        raise ValueError("base layer of tree_out must be all leaves and no branches")
+    if all_leaves_in and all_leaves_out:  # both trees at leaf layer
+        if len(tree_in) <= len(tree_out):
+            return tree_in + [[] for _ in range(len(tree_out) - len(tree_in))]
+        else:
+            raise ValueError("tree_in cannot have more leaves than tree_out")
+    elif all_leaves_in and not all_leaves_out:  # tree_out is deeper than tree_in
+        return [broadcast_tree(tree_in, branch) for branch in tree_out]
+    elif all_leaves_out and not all_leaves_in:
+        raise ValueError("tree_in cannot have a deeper structure than tree_out")
+    else:  # both trees at branch layers
+        if len(tree_in) == len(tree_out):
+            return [
+                broadcast_tree(tree_in[k], tree_out[k]) for k in range(len(tree_out))
+            ]
+        else:
+            raise ValueError(
+                "tree_in must have the same number of branches as tree_out"
+            )
diff --git a/tests/test_utils.py b/tests/test_utils.py
index ff421875b6..79d6ce51fc 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,8 +3,9 @@
 import numpy as np
 import pytest
 
+from desc.backend import tree_map
 from desc.grid import LinearGrid
-from desc.utils import isalmostequal, islinspaced
+from desc.utils import broadcast_tree, isalmostequal, islinspaced
 
 
 @pytest.mark.unit
@@ -53,3 +54,49 @@ def test_islinspaced():
 
     # 0D arrays will return True
     assert islinspaced(np.array(0))
+
+
+@pytest.mark.unit
+def test_broadcast_tree():
+    """Test that broadcast_tree works on various pytree structures."""
+    tree_out = [[1, 2, 3], [[4], [[5, 6], [7]]]]
+
+    # tree of tuples, not lists
+    tree_in = [(0, 1), [2]]
+    with pytest.raises(ValueError):
+        _ = broadcast_tree(tree_in, tree_out)
+
+    # tree with too many leaves per branch
+    tree_in = [1, 2]
+    with pytest.raises(ValueError):
+        _ = broadcast_tree(tree_in, tree_out)
+
+    # tree with a mix of leaves and branches at the same layer
+    tree_in = [[0, 1], 2, [3]]
+    with pytest.raises(ValueError):
+        _ = broadcast_tree(tree_in, tree_out)
+
+    # tree_in is deeper than tree_out
+    tree_in = [[[1], [2, 3]], [[4], [[[5], [6]], [7]]]]
+    with pytest.raises(ValueError):
+        _ = broadcast_tree(tree_in, tree_out)
+
+    # tree with proper structure already does not change
+    tree_in = tree_map(lambda x: x * 2, tree_out)
+    tree = broadcast_tree(tree_in, tree_out)
+    assert tree == tree_in
+
+    # broadcast single leaf to full tree
+    tree_in = 0
+    tree = broadcast_tree(tree_in, tree_out)
+    assert tree == [[0, [], []], [[0], [[0, []], [0]]]]
+
+    # broadcast from only major branches
+    tree_in = [[1, 2], [3]]
+    tree = broadcast_tree(tree_in, tree_out)
+    assert tree == [[1, 2, []], [[3], [[3, []], [3]]]]
+
+    # more complicated example
+    tree_in = [[1, 2], [[3], [4]]]
+    tree = broadcast_tree(tree_in, tree_out)
+    assert tree == [[1, 2, []], [[3], [[4, []], [4]]]]

From 135c0a956e34f1e5ff44c73fd1a729fdfa6d7f06 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 19 Apr 2024 23:51:37 -0500
Subject: [PATCH 10/50] broadcast_tree until function + tests

---
 desc/utils.py       | 60 +++++++++++++++++++++++++++++++++------------
 tests/test_utils.py | 10 +++++---
 2 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/desc/utils.py b/desc/utils.py
index d6d72a726e..40f31e96c4 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -610,23 +610,41 @@ def is_any_instance(things, cls):
     return any([isinstance(t, cls) for t in things])
 
 
-def broadcast_tree(tree_in, tree_out):
-    """Broadcast tree_in to the same pytree structure as tree_out."""
-    errorif(
-        isinstance(tree_in, (tuple, dict)),
-        ValueError,
-        "tree_in must be a pytree of nested lists, not tuples or dicts",
-    )
-    errorif(
-        isinstance(tree_out, (tuple, dict)),
-        ValueError,
-        "tree_out must be a pytree of nested lists, not tuples or dicts",
-    )
+def broadcast_tree(tree_in, tree_out, value=False, sort=False):
+    """Broadcast tree_in to the same pytree structure as tree_out.
+
+    Parameters
+    ----------
+    tree_in : pytree
+        Tree to broadcast.
+    tree_out : pytree
+        Tree with structure to broadcast to.
+    value : leaf, optional
+        Leaf value to pad branches with if not present in tree_in. Default = False.
+    sort : bool, optional
+        If True, checks that leaves of tree_in are also leaves of tree_out, and keeps
+        them in the same position on their branches when broadcasting. Default = False.
+
+    Returns
+    -------
+    tree : pytree
+        Tree with the leaves of tree_in broadcast to the structure of tree_out.
+
+    """
     if not isinstance(tree_in, list):
         tree_in = [tree_in]
     if not isinstance(tree_out, list):
         tree_out = [tree_out]
-    where_leaves_in = [treedef_is_leaf(tree_structure(branch)) for branch in tree_in]
+
+    def isemptylist(x):
+        if isinstance(x, list):
+            return not x
+        return False
+
+    where_leaves_in = [  # tree_in can have empty branches that are not leaves
+        treedef_is_leaf(tree_structure(branch)) and not isemptylist(branch)
+        for branch in tree_in
+    ]
     where_leaves_out = [treedef_is_leaf(tree_structure(branch)) for branch in tree_out]
     all_leaves_in = all(where_leaves_in)
     all_leaves_out = all(where_leaves_out)
@@ -636,17 +654,27 @@ def broadcast_tree(tree_in, tree_out):
         raise ValueError("base layer of tree_out must be all leaves and no branches")
     if all_leaves_in and all_leaves_out:  # both trees at leaf layer
         if len(tree_in) <= len(tree_out):
-            return tree_in + [[] for _ in range(len(tree_out) - len(tree_in))]
+            if sort:
+                if set(tree_in) <= set(tree_out):
+                    return [leaf if leaf in tree_in else value for leaf in tree_out]
+                else:
+                    raise ValueError(
+                        "leaves of tree_in must be a subset of the leaves in tree_out "
+                        + "if sort=True"
+                    )
+            else:
+                return tree_in + [value for _ in range(len(tree_out) - len(tree_in))]
         else:
             raise ValueError("tree_in cannot have more leaves than tree_out")
     elif all_leaves_in and not all_leaves_out:  # tree_out is deeper than tree_in
-        return [broadcast_tree(tree_in, branch) for branch in tree_out]
+        return [broadcast_tree(tree_in, branch, sort=sort) for branch in tree_out]
     elif all_leaves_out and not all_leaves_in:
         raise ValueError("tree_in cannot have a deeper structure than tree_out")
     else:  # both trees at branch layers
         if len(tree_in) == len(tree_out):
             return [
-                broadcast_tree(tree_in[k], tree_out[k]) for k in range(len(tree_out))
+                broadcast_tree(tree_in[k], tree_out[k], sort=sort)
+                for k in range(len(tree_out))
             ]
         else:
             raise ValueError(
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 79d6ce51fc..f54208c4c6 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -89,14 +89,18 @@ def test_broadcast_tree():
     # broadcast single leaf to full tree
     tree_in = 0
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[0, [], []], [[0], [[0, []], [0]]]]
+    assert tree == [[0, False, False], [[0], [[0, False], [0]]]]
 
     # broadcast from only major branches
     tree_in = [[1, 2], [3]]
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[1, 2, []], [[3], [[3, []], [3]]]]
+    assert tree == [[1, 2, False], [[3], [[3, False], [3]]]]
 
     # more complicated example
     tree_in = [[1, 2], [[3], [4]]]
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[1, 2, []], [[3], [[4, []], [4]]]]
+    assert tree == [[1, 2, False], [[3], [[4, False], [4]]]]
+
+    # TODO: add test for sort=True
+    # TODO: add test for value
+    # TODO: add test with empy branches

From fafa8cac89e5d52f426fb8aa2cd66e6cd7839a33 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 19 Apr 2024 23:53:24 -0500
Subject: [PATCH 11/50] FixCollectionParameters working with custom params
 input

---
 desc/objectives/linear_objectives.py | 137 +++++++++------------------
 tests/test_linear_objectives.py      |  19 ++--
 2 files changed, 60 insertions(+), 96 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 8528eb65ad..8d711af3a4 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -12,14 +12,15 @@
 import numpy as np
 from termcolor import colored
 
-from desc.backend import jnp
+from desc.backend import jnp, tree_leaves, tree_map, tree_structure, tree_unflatten
 from desc.basis import zernike_radial, zernike_radial_coeffs
-from desc.utils import errorif, flatten_list, setdefault, unique_list
+from desc.utils import broadcast_tree, errorif, setdefault
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective
 
 
+# TODO: get rid of this class and inherit from FixParameter instead?
 class _FixedObjective(_Objective):
     _fixed = True
     _linear = True
@@ -44,6 +45,8 @@ def update_target(self, thing):
     def _parse_target_from_user(
         self, target_from_user, default_target, default_bounds, idx
     ):
+        # FIXME: add logic here to deal with `target_from_user` as a pytree?
+        # FIXME: does this actually need idx?
         if target_from_user is None:
             target = default_target
             bounds = default_bounds
@@ -70,7 +73,7 @@ class FixParameter(_FixedObjective):
         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 : 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.
@@ -113,7 +116,12 @@ def __init__(
         normalize_target=False,
         name=None,
     ):
-        # TODO: assert that `thing` is not of type `OptimizableCollection`
+        errorif(
+            hasattr(thing, "__len__"),
+            ValueError,
+            "Thing must be of type `Optimizable` and not `OptimizableCollection`; "
+            + f"got {thing}.",
+        )
         self._target_from_user = target
         self._params = params = setdefault(params, thing.optimizable_params)
         self._indices = indices
@@ -227,7 +235,7 @@ class FixCollectionParameters(_FixedObjective):
         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 : 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.
@@ -266,7 +274,7 @@ def __init__(
         weight=1,
         normalize=False,
         normalize_target=False,
-        name="Fixed parameters",
+        name="Fixed parameters",  # TODO: add params
     ):
         errorif(
             not hasattr(thing, "__len__"),
@@ -298,86 +306,47 @@ def build(self, use_jit=False, verbose=1):
 
         """
         thing = self.things[0]
-        thing_idx = range(len(thing))  # indices of things in OptimizableCollection
-        params = setdefault(self._params, thing.optimizable_params)
+        structure = tree_structure(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
+        # set default params
+        self._params = setdefault(self._params, thing.optimizable_params)
+        self._params = broadcast_tree(self._params, thing.optimizable_params, sort=True)
+        assert tree_structure(self._params) == structure
+        params_leaves = tree_leaves(self._params)
 
-        # replace indices=True with actual indices
+        # set default 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
+            indices = tree_map(lambda dim: np.arange(dim, dtype=int), thing.dimensions)
+            indices = [
+                idx if param else False
+                for idx, param in zip(tree_leaves(indices), params_leaves)
             ]
-        # 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
-        )
+            self._indices = tree_unflatten(structure, indices)
+        self._indices = broadcast_tree(self._indices, thing.optimizable_params)
+        assert tree_structure(self._indices) == structure
+        self._indices_leaves = tree_leaves(self._indices)
+        self._indices_leaves = [  # replace False leaves with empy arrays
+            indices if not isinstance(indices, bool) else np.array([], dtype=int)
+            for indices in self._indices_leaves
+        ]
 
-        # 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
+        # TODO: check that params & indices are self-consistent
+
+        # set default target
+        targets = [
+            target
+            for target, param in zip(tree_leaves(thing.params_dict), params_leaves)
+            if param
         ]
+        default_target = jnp.concatenate(targets)
         default_bounds = None
-        target, bounds = self._parse_target_from_user(
+
+        self._dim_f = sum(indices.size for indices in self._indices_leaves)
+
+        self.target, self.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):
@@ -398,19 +367,7 @@ def compute(self, params, constants=None):
 
         """
         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))
-            ]
+            [par[idx] for par, idx in zip(tree_leaves(params), self._indices_leaves)]
         )
 
 
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index 0ec081f2fc..a80df1c015 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -6,7 +6,13 @@
 from qsc import Qsc
 
 import desc.examples
-from desc.coils import CoilSet, FourierPlanarCoil, FourierRZCoil, MixedCoilSet
+from desc.coils import (
+    CoilSet,
+    FourierPlanarCoil,
+    FourierRZCoil,
+    FourierXYZCoil,
+    MixedCoilSet,
+)
 from desc.equilibrium import Equilibrium
 from desc.geometry import FourierRZToroidalSurface
 from desc.grid import LinearGrid
@@ -923,7 +929,7 @@ def test_fix_omni_indices():
 
 
 @pytest.mark.unit
-def test_fix_subset_of_params_in_collection(self):
+def test_fix_subset_of_params_in_collection():
     """Tests FixCollectionParameters fixing a subset of things in the collection."""
     tf_coil = FourierPlanarCoil(center=[2, 0, 0], normal=[0, 1, 0], r_n=[1])
     tf_coilset = CoilSet.linspaced_angular(tf_coil, n=4)
@@ -931,13 +937,14 @@ def test_fix_subset_of_params_in_collection(self):
     vf_coilset = CoilSet.linspaced_linear(
         vf_coil, displacement=[0, 0, 2], n=3, endpoint=True
     )
-    full_coilset = MixedCoilSet((tf_coilset, vf_coilset))
+    xy_coil = FourierXYZCoil()
+    full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xy_coil))
 
     params = [
         [["current"], ["center", "normal", "r_n"], ["current", "shift", "rotmat"], []],
         [["current", "Z_n"], ["R_n", "Z_n"], ["Z_n"]],
+        ["Z_n", "shift", "rotmat"],
     ]
     obj = FixCollectionParameters(full_coilset, params=params)
-    if False:  # FIXME: get this test working
-        obj.build()
-    pass
+    obj.build()
+    assert obj.dim_f == 41

From 566c81031e5119d9415cedae37bea8f4d766d58c Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Fri, 19 Apr 2024 23:57:38 -0500
Subject: [PATCH 12/50] tiny typos

---
 desc/objectives/_free_boundary.py | 2 +-
 tests/test_utils.py               | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py
index a141e550cf..ebcdf71ab2 100644
--- a/desc/objectives/_free_boundary.py
+++ b/desc/objectives/_free_boundary.py
@@ -489,7 +489,7 @@ def build(self, use_jit=True, verbose=1):
                 rho=np.array([1.0]),
                 M=eq.M_grid,
                 N=eq.N_grid,
-                NFP=int(eq.NFP) if eq.N > 0 else 64,
+                NFP=eq.NFP if eq.N > 0 else 64,
                 sym=False,
             )
         else:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index f54208c4c6..550eac6724 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -103,4 +103,4 @@ def test_broadcast_tree():
 
     # TODO: add test for sort=True
     # TODO: add test for value
-    # TODO: add test with empy branches
+    # TODO: add test with empty branches

From c39faa20cfad21a3b33728d1998e44c04d9fdad8 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sat, 20 Apr 2024 13:24:33 -0500
Subject: [PATCH 13/50] repair 2nd stage opt test

---
 tests/test_examples.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index 2a962c26df..a755ddd4a7 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1289,7 +1289,7 @@ def test_second_stage_optimization():
     eq = get("DSHAPE")
     field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
     objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
-    constraints = FixCollectionParameters(field, "R0")
+    constraints = FixCollectionParameters(field, [["R0"], []])
     optimizer = Optimizer("lsq-exact")
     (field,), _ = optimizer.optimize(
         things=field, objective=objective, constraints=constraints, verbose=2

From 80e7dbe38061f8418a004da6c6487d1149ff1250 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sat, 20 Apr 2024 13:37:59 -0500
Subject: [PATCH 14/50] more test cases for broadcast_tree

---
 tests/test_utils.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 550eac6724..51b905d57a 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -72,7 +72,7 @@ def test_broadcast_tree():
         _ = broadcast_tree(tree_in, tree_out)
 
     # tree with a mix of leaves and branches at the same layer
-    tree_in = [[0, 1], 2, [3]]
+    tree_in = [[1, 2], 3, [4]]
     with pytest.raises(ValueError):
         _ = broadcast_tree(tree_in, tree_out)
 
@@ -81,6 +81,11 @@ def test_broadcast_tree():
     with pytest.raises(ValueError):
         _ = broadcast_tree(tree_in, tree_out)
 
+    # tree_in leaves not in tree_out
+    tree_in = [[1, 2], [[3], [4]]]
+    with pytest.raises(ValueError):
+        tree = broadcast_tree(tree_in, tree_out, sort=True)
+
     # tree with proper structure already does not change
     tree_in = tree_map(lambda x: x * 2, tree_out)
     tree = broadcast_tree(tree_in, tree_out)
@@ -96,11 +101,22 @@ def test_broadcast_tree():
     tree = broadcast_tree(tree_in, tree_out)
     assert tree == [[1, 2, False], [[3], [[3, False], [3]]]]
 
-    # more complicated example
+    # broadcast from minor branches
     tree_in = [[1, 2], [[3], [4]]]
     tree = broadcast_tree(tree_in, tree_out)
     assert tree == [[1, 2, False], [[3], [[4, False], [4]]]]
 
-    # TODO: add test for sort=True
-    # TODO: add test for value
-    # TODO: add test with empty branches
+    # sort order of leaves
+    tree_in = [[3, 1], [[4], [[6], []]]]
+    tree = broadcast_tree(tree_in, tree_out, sort=True)
+    assert tree == [[1, False, 3], [[4], [[False, 6], [False]]]]
+
+    # tree_in with empty branches
+    tree_in = [[], [[1], [[2], []]]]
+    tree = broadcast_tree(tree_in, tree_out)
+    assert tree == [[False, False, False], [[1], [[2, False], [False]]]]
+
+    # custom fill value
+    tree_in = [[1, 2], [[3], [4]]]
+    tree = broadcast_tree(tree_in, tree_out, value=0)
+    assert tree == [[1, 2, 0], [[3], [[4, 0], [4]]]]

From 65f6f6e5d876ffccef534fc9e5601947925ddc3c Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sat, 20 Apr 2024 17:10:04 -0500
Subject: [PATCH 15/50] combine FixCollectionParameters into FixParameter

---
 desc/objectives/__init__.py          |   1 -
 desc/objectives/linear_objectives.py | 237 +++++----------------------
 tests/test_examples.py               |  12 +-
 tests/test_linear_objectives.py      |  14 +-
 4 files changed, 55 insertions(+), 209 deletions(-)

diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py
index eed8c4cd74..e0ae573e60 100644
--- a/desc/objectives/__init__.py
+++ b/desc/objectives/__init__.py
@@ -49,7 +49,6 @@
     FixAxisZ,
     FixBoundaryR,
     FixBoundaryZ,
-    FixCollectionParameters,
     FixCurrent,
     FixCurveRotation,
     FixCurveShift,
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 8d711af3a4..ae67a394e0 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -64,179 +64,18 @@ def _parse_target_from_user(
         return target, bounds
 
 
-class FixParameter(_FixedObjective):
-    """Fix specific degrees of freedom associated with a given Optimizable object.
+class FixParameter(_Objective):
+    """Fix specific degrees of freedom associated with a given Optimizable thing.
 
     Parameters
     ----------
     thing : Optimizable
         Object whose degrees of freedom are being fixed.
-    params : str or list of str
+    params : pytree with leaves of type str
         Names of parameters to fix. Defaults to all parameters.
-    indices : 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 ``target=thing.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.
-        Defaults to ``target=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 parameter error: {:10.3e} "
-
-    def __init__(
-        self,
-        thing,
-        params=None,
-        indices=True,
-        target=None,
-        bounds=None,
-        weight=1,
-        normalize=False,
-        normalize_target=False,
-        name=None,
-    ):
-        errorif(
-            hasattr(thing, "__len__"),
-            ValueError,
-            "Thing must be of type `Optimizable` and not `OptimizableCollection`; "
-            + f"got {thing}.",
-        )
-        self._target_from_user = target
-        self._params = params = setdefault(params, thing.optimizable_params)
-        self._indices = indices
-        self._print_value_fmt = (
-            f"Fixed parameter ({self._params}) error: " + "{:10.3e} "
-        )
-        name = setdefault(name, f"Fixed parameter ({self._params})")
-        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]
-        params = setdefault(self._params, thing.optimizable_params)
-
-        if not isinstance(params, (list, tuple)):
-            params = [params]
-        for par in params:
-            errorif(
-                par not in thing.optimizable_params,
-                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[par]) for par in self._params]
-        # 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
-        indices = {}
-        errorif(
-            len(self._params) != len(self._indices),
-            ValueError,
-            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:
-                idx = np.arange(thing.dimensions[par])
-            indices[par] = np.atleast_1d(idx)
-        self._indices = indices
-        self._dim_f = sum(t.size for t in self._indices.values())
-
-        default_target = {
-            par: thing.params_dict[par][self._indices[par]] for par in params
-        }
-        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([target[par] for par in params])
-            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 : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.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(
-            [params[par][self._indices[par]] for par in self._params]
-        )
-
-
-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.
-    indices : array-like or list of array-like
-        Indices to fix for each parameter in params. Use True to fix all indices.
+    indices : pytree with leaves of type array
+        Indices to fix for each parameter in params.
+        The default value indices=True fixes 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.
@@ -276,12 +115,6 @@ def __init__(
         normalize_target=False,
         name="Fixed parameters",  # TODO: add params
     ):
-        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__(
@@ -308,20 +141,21 @@ def build(self, use_jit=False, verbose=1):
         thing = self.things[0]
         structure = tree_structure(thing.optimizable_params)
 
-        # set default params
+        # set params
         self._params = setdefault(self._params, thing.optimizable_params)
         self._params = broadcast_tree(self._params, thing.optimizable_params, sort=True)
         assert tree_structure(self._params) == structure
         params_leaves = tree_leaves(self._params)
 
-        # set default indices
-        if isinstance(self._indices, bool) and self._indices:
+        # set indices
+        if isinstance(self._indices, bool) and self._indices:  # default to all params
             indices = tree_map(lambda dim: np.arange(dim, dtype=int), thing.dimensions)
-            indices = [
+            indices_leaves = [
                 idx if param else False
                 for idx, param in zip(tree_leaves(indices), params_leaves)
             ]
-            self._indices = tree_unflatten(structure, indices)
+            self._indices = tree_unflatten(structure, indices_leaves)
+        # caution: cannot sort the indices!
         self._indices = broadcast_tree(self._indices, thing.optimizable_params)
         assert tree_structure(self._indices) == structure
         self._indices_leaves = tree_leaves(self._indices)
@@ -330,22 +164,21 @@ def build(self, use_jit=False, verbose=1):
             for indices in self._indices_leaves
         ]
 
-        # TODO: check that params & indices are self-consistent
+        self._dim_f = sum(idx.size for idx in self._indices_leaves)
 
         # set default target
-        targets = [
-            target
-            for target, param in zip(tree_leaves(thing.params_dict), params_leaves)
-            if param
-        ]
-        default_target = jnp.concatenate(targets)
-        default_bounds = None
-
-        self._dim_f = sum(indices.size for indices in self._indices_leaves)
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, default_target, default_bounds, indices
-        )
+        if self.target is None and self.bounds is None:
+            self.target = jnp.concatenate(
+                [
+                    target[idx]
+                    for target, param, idx in zip(
+                        tree_leaves(thing.params_dict),
+                        params_leaves,
+                        self._indices_leaves,
+                    )
+                    if param
+                ]
+            )
 
         super().build(use_jit=use_jit, verbose=verbose)
 
@@ -367,9 +200,27 @@ def compute(self, params, constants=None):
 
         """
         return jnp.concatenate(
-            [par[idx] for par, idx in zip(tree_leaves(params), self._indices_leaves)]
+            [
+                param[idx]
+                for param, idx in zip(tree_leaves(params), self._indices_leaves)
+            ]
         )
 
+    def update_target(self, thing):
+        """Update target values using an Optimizable object.
+
+        Parameters
+        ----------
+        thing : Optimizable
+            Optimizable object that will be optimized to satisfy the Objective.
+
+        """
+        new_target = self.compute(thing.params_dict)
+        assert len(new_target) == len(self.target)
+        self.target = new_target
+        if self._use_jit:
+            self.jit()
+
 
 class BoundaryRSelfConsistency(_Objective):
     """Ensure that the boundary and interior surfaces are self-consistent.
diff --git a/tests/test_examples.py b/tests/test_examples.py
index a755ddd4a7..d08e4b3679 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -32,7 +32,6 @@
     CurrentDensity,
     FixBoundaryR,
     FixBoundaryZ,
-    FixCollectionParameters,
     FixCurrent,
     FixIota,
     FixOmniBmax,
@@ -641,7 +640,7 @@ def test_multiobject_optimization_al():
     constraints = (
         ForceBalance(eq=eq, bounds=(-1e-4, 1e-4), normalize_target=False),
         FixPressure(eq=eq),
-        FixParameter(surf, ["Z_lmn", "R_lmn"], [[-1], [0]]),
+        FixParameter(surf, ["R_lmn", "Z_lmn"], [np.array([0]), np.array([-1])]),
         FixParameter(eq, ["Psi", "i_l"]),
         FixBoundaryR(eq, modes=[[0, 0, 0]]),
         PlasmaVesselDistance(surface=surf, eq=eq, target=1),
@@ -680,7 +679,7 @@ def test_multiobject_optimization_prox():
     constraints = (
         ForceBalance(eq=eq),
         FixPressure(eq=eq),
-        FixParameter(surf, ["Z_lmn", "R_lmn"], [[-1], [0]]),
+        FixParameter(surf, ["R_lmn", "Z_lmn"], [np.array([0]), np.array([-1])]),
         FixParameter(eq, ["Psi", "i_l"]),
         FixBoundaryR(eq, modes=[[0, 0, 0]]),
     )
@@ -1257,7 +1256,7 @@ def test_quadratic_flux_optimization_with_analytic_field():
 
     optimizer = Optimizer("lsq-exact")
 
-    constraints = (FixParameter(field, ["R0"]),)
+    constraints = (FixParameter(field, "R0"),)
     quadflux_obj = QuadraticFlux(
         eq=eq,
         field=field,
@@ -1283,13 +1282,10 @@ def test_quadratic_flux_optimization_with_analytic_field():
 @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.
-
     eq = get("DSHAPE")
     field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
     objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
-    constraints = FixCollectionParameters(field, [["R0"], []])
+    constraints = FixParameter(field, [["R0"], []])
     optimizer = Optimizer("lsq-exact")
     (field,), _ = optimizer.optimize(
         things=field, objective=objective, constraints=constraints, verbose=2
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index a80df1c015..e14124c2e1 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -29,7 +29,6 @@
     FixAxisZ,
     FixBoundaryR,
     FixBoundaryZ,
-    FixCollectionParameters,
     FixCurrent,
     FixElectronDensity,
     FixElectronTemperature,
@@ -414,7 +413,7 @@ def test_kinetic_constraints():
 
 @pytest.mark.unit
 def test_correct_indexing_passed_modes():
-    """Test Indexing when passing in specified modes, related to gh issue #380."""
+    """Test indexing when passing in specified modes, related to gh issue #380."""
     n = 1
     eq = desc.examples.get("W7-X")
     eq.change_resolution(3, 3, 3, 6, 6, 6)
@@ -467,7 +466,7 @@ def test_correct_indexing_passed_modes():
 
 @pytest.mark.unit
 def test_correct_indexing_passed_modes_and_passed_target():
-    """Test Indexing when passing in specified modes, related to gh issue #380."""
+    """Test indexing when passing in specified modes, related to gh issue #380."""
     n = 1
     eq = desc.examples.get("W7-X")
     eq.change_resolution(3, 3, 3, 6, 6, 6)
@@ -529,7 +528,7 @@ def test_correct_indexing_passed_modes_and_passed_target():
 
 @pytest.mark.unit
 def test_correct_indexing_passed_modes_axis():
-    """Test Indexing when passing in specified axis modes, related to gh issue #380."""
+    """Test indexing when passing in specified axis modes, related to gh issue #380."""
     n = 1
     eq = desc.examples.get("W7-X")
     eq.change_resolution(3, 3, 3, 6, 6, 6)
@@ -588,7 +587,7 @@ def test_correct_indexing_passed_modes_axis():
 
 @pytest.mark.unit
 def test_correct_indexing_passed_modes_and_passed_target_axis():
-    """Test Indexing when passing in specified axis modes, related to gh issue #380."""
+    """Test indexing when passing in specified axis modes, related to gh issue #380."""
     n = 1
 
     eq = desc.examples.get("W7-X")
@@ -930,7 +929,7 @@ def test_fix_omni_indices():
 
 @pytest.mark.unit
 def test_fix_subset_of_params_in_collection():
-    """Tests FixCollectionParameters fixing a subset of things in the collection."""
+    """Tests FixParameter fixing a subset of things in the collection."""
     tf_coil = FourierPlanarCoil(center=[2, 0, 0], normal=[0, 1, 0], r_n=[1])
     tf_coilset = CoilSet.linspaced_angular(tf_coil, n=4)
     vf_coil = FourierRZCoil(R_n=3, Z_n=-1)
@@ -940,11 +939,12 @@ def test_fix_subset_of_params_in_collection():
     xy_coil = FourierXYZCoil()
     full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xy_coil))
 
+    # TODO: use a better example to test broadcasting for a whole coilset
     params = [
         [["current"], ["center", "normal", "r_n"], ["current", "shift", "rotmat"], []],
         [["current", "Z_n"], ["R_n", "Z_n"], ["Z_n"]],
         ["Z_n", "shift", "rotmat"],
     ]
-    obj = FixCollectionParameters(full_coilset, params=params)
+    obj = FixParameter(full_coilset, params)
     obj.build()
     assert obj.dim_f == 41

From 3477a960aac125303149041eda80aaece65778d6 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sat, 20 Apr 2024 18:11:11 -0500
Subject: [PATCH 16/50] fix params_leaves sorting issue

---
 desc/objectives/linear_objectives.py | 49 +++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 4 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index ae67a394e0..5cceafa266 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -12,9 +12,16 @@
 import numpy as np
 from termcolor import colored
 
-from desc.backend import jnp, tree_leaves, tree_map, tree_structure, tree_unflatten
+from desc.backend import (
+    jnp,
+    tree_leaves,
+    tree_map,
+    tree_structure,
+    tree_unflatten,
+    treedef_is_leaf,
+)
 from desc.basis import zernike_radial, zernike_radial_coeffs
-from desc.utils import broadcast_tree, errorif, setdefault
+from desc.utils import broadcast_tree, errorif, flatten_list, setdefault
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective
@@ -141,12 +148,45 @@ def build(self, use_jit=False, verbose=1):
         thing = self.things[0]
         structure = tree_structure(thing.optimizable_params)
 
+        def leaves_per_branch(tree_tuple):
+            """Get the number of leaves per branch in a pytree, in flattened order."""
+            tree, lengths = tree_tuple
+            if all([treedef_is_leaf(tree_structure(x)) for x in tree]):
+                lengths.append(len(tree))
+                return [], lengths
+            else:
+                return (
+                    [leaves_per_branch((branch, lengths)) for branch in tree],
+                    lengths,
+                )
+
+        # indices of flattened tree leaves
+        _, lengths = leaves_per_branch((thing.optimizable_params, []))
+        starts = np.insert(np.cumsum(lengths), 0, 0)
+        leaf_indices = [
+            np.arange(start, start + length, dtype=int)
+            for start, length in zip(starts, lengths)
+        ]
+
         # set params
         self._params = setdefault(self._params, thing.optimizable_params)
         self._params = broadcast_tree(self._params, thing.optimizable_params, sort=True)
         assert tree_structure(self._params) == structure
         params_leaves = tree_leaves(self._params)
 
+        # sort params to the same flattend order as indices
+        # this is necessary because of JAX GitHub Issue #4085
+        params_leaves = flatten_list(
+            [
+                list(
+                    np.array(params_leaves)[idx][
+                        np.argsort(np.array(thing.optimizable_params)[idx])
+                    ]
+                )
+                for idx in leaf_indices
+            ]
+        )
+
         # set indices
         if isinstance(self._indices, bool) and self._indices:  # default to all params
             indices = tree_map(lambda dim: np.arange(dim, dtype=int), thing.dimensions)
@@ -155,8 +195,9 @@ def build(self, use_jit=False, verbose=1):
                 for idx, param in zip(tree_leaves(indices), params_leaves)
             ]
             self._indices = tree_unflatten(structure, indices_leaves)
-        # caution: cannot sort the indices!
-        self._indices = broadcast_tree(self._indices, thing.optimizable_params)
+        self._indices = broadcast_tree(  # indices cannot be sorted like the params!
+            self._indices, thing.optimizable_params, sort=False
+        )
         assert tree_structure(self._indices) == structure
         self._indices_leaves = tree_leaves(self._indices)
         self._indices_leaves = [  # replace False leaves with empy arrays

From b41e761b99a764eea9619e1f536fdc55a761ecd7 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sun, 21 Apr 2024 13:12:36 -0500
Subject: [PATCH 17/50] fix list vs array bug

---
 desc/basis.py                        |  2 +-
 desc/objectives/linear_objectives.py | 22 +++++++++-------------
 tests/test_examples.py               |  4 +++-
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/desc/basis.py b/desc/basis.py
index b394fcfa90..f5fde6f7a6 100644
--- a/desc/basis.py
+++ b/desc/basis.py
@@ -118,7 +118,7 @@ def get_idx(self, L=0, M=0, N=0, error=True):
 
         Returns
         -------
-        idx : ndarray of int
+        idx : int
             Index of given mode numbers.
 
         """
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 5cceafa266..d512b16731 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -178,11 +178,10 @@ def leaves_per_branch(tree_tuple):
         # this is necessary because of JAX GitHub Issue #4085
         params_leaves = flatten_list(
             [
-                list(
-                    np.array(params_leaves)[idx][
-                        np.argsort(np.array(thing.optimizable_params)[idx])
-                    ]
-                )
+                [
+                    [params_leaves[i] for i in idx][j]
+                    for j in np.argsort(np.array(thing.optimizable_params)[idx])
+                ]
                 for idx in leaf_indices
             ]
         )
@@ -209,15 +208,12 @@ def leaves_per_branch(tree_tuple):
 
         # set default target
         if self.target is None and self.bounds is None:
-            self.target = jnp.concatenate(
+            self.target = np.concatenate(
                 [
-                    target[idx]
-                    for target, param, idx in zip(
-                        tree_leaves(thing.params_dict),
-                        params_leaves,
-                        self._indices_leaves,
+                    np.atleast_1d(param[idx])
+                    for param, idx in zip(
+                        tree_leaves(thing.params_dict), self._indices_leaves
                     )
-                    if param
                 ]
             )
 
@@ -242,7 +238,7 @@ def compute(self, params, constants=None):
         """
         return jnp.concatenate(
             [
-                param[idx]
+                jnp.atleast_1d(param[idx])
                 for param, idx in zip(tree_leaves(params), self._indices_leaves)
             ]
         )
diff --git a/tests/test_examples.py b/tests/test_examples.py
index d08e4b3679..a6a449f290 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1002,7 +1002,9 @@ def test_only_non_eq_optimization():
 
     surf.change_resolution(M=eq.M, N=eq.N)
     constraints = (
-        FixParameter(surf, params="R_lmn", indices=surf.R_basis.get_idx(0, 0, 0)),
+        FixParameter(
+            surf, params="R_lmn", indices=np.array(surf.R_basis.get_idx(0, 0, 0))
+        ),
     )
 
     obj = PrincipalCurvature(surf, target=1)

From 31c94c672c919162502ff0777d37fc70f2e04bc3 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sun, 21 Apr 2024 13:23:34 -0500
Subject: [PATCH 18/50] clean up a few lines

---
 desc/objectives/linear_objectives.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index d512b16731..02bcc39eda 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -187,7 +187,8 @@ def leaves_per_branch(tree_tuple):
         )
 
         # set indices
-        if isinstance(self._indices, bool) and self._indices:  # default to all params
+        if isinstance(self._indices, bool):  # default to all indices
+            errorif(not self._indices, ValueError, "indices cannot be False")
             indices = tree_map(lambda dim: np.arange(dim, dtype=int), thing.dimensions)
             indices_leaves = [
                 idx if param else False
@@ -252,9 +253,7 @@ def update_target(self, thing):
             Optimizable object that will be optimized to satisfy the Objective.
 
         """
-        new_target = self.compute(thing.params_dict)
-        assert len(new_target) == len(self.target)
-        self.target = new_target
+        self.target = self.compute(thing.params_dict)
         if self._use_jit:
             self.jit()
 

From d41a1fca4e01a8e0091ea20d409021081f0a937e Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sun, 21 Apr 2024 14:10:53 -0500
Subject: [PATCH 19/50] fix missing tree_leaves call

---
 desc/objectives/linear_objectives.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 02bcc39eda..4293f97455 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -71,6 +71,7 @@ def _parse_target_from_user(
         return target, bounds
 
 
+# Rename to `FixParameters` (plural) instead?
 class FixParameter(_Objective):
     """Fix specific degrees of freedom associated with a given Optimizable thing.
 
@@ -174,13 +175,15 @@ def leaves_per_branch(tree_tuple):
         assert tree_structure(self._params) == structure
         params_leaves = tree_leaves(self._params)
 
-        # sort params to the same flattend order as indices
+        # sort params to the same flattend leaf order as indices
         # this is necessary because of JAX GitHub Issue #4085
         params_leaves = flatten_list(
             [
                 [
                     [params_leaves[i] for i in idx][j]
-                    for j in np.argsort(np.array(thing.optimizable_params)[idx])
+                    for j in np.argsort(
+                        np.atleast_1d(tree_leaves(thing.optimizable_params))[idx]
+                    )
                 ]
                 for idx in leaf_indices
             ]

From f0dbd83ddc06593fa681a6ea7fc7c5fe8030d612 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Sun, 21 Apr 2024 20:54:03 -0500
Subject: [PATCH 20/50] add assert statement to fix later

---
 desc/objectives/linear_objectives.py | 6 +++++-
 desc/objectives/utils.py             | 2 ++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 4293f97455..c582bf4025 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -240,10 +240,14 @@ def compute(self, params, constants=None):
             Fixed degree of freedom errors.
 
         """
+        params_leaves = tree_leaves(params)
+        # FIXME: should still work if only a subset of params get passed in
+        # TODO: need to cast to full dict
+        assert len(params_leaves) == len(self._indices_leaves)
         return jnp.concatenate(
             [
                 jnp.atleast_1d(param[idx])
-                for param, idx in zip(tree_leaves(params), self._indices_leaves)
+                for param, idx in zip(params_leaves, self._indices_leaves)
             ]
         )
 
diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py
index 949b3cc09f..8c3f780b7e 100644
--- a/desc/objectives/utils.py
+++ b/desc/objectives/utils.py
@@ -144,7 +144,9 @@ def recover(x_reduced):
 
     # check that all constraints are actually satisfiable
     params = objective.unpack_state(xp, False)
+    print(params)
     for con in constraint.objectives:
+        # FIXME: xpi is not always correct -- missing params?
         xpi = [params[i] for i, t in enumerate(objective.things) if t in con.things]
         y1 = con.compute_unscaled(*xpi)
         y2 = con.target

From be981f681a8cbfd9889e62f087435b4d6496cc1a Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Wed, 24 Apr 2024 10:26:38 -0400
Subject: [PATCH 21/50] remove debugging print statement

---
 desc/objectives/utils.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py
index 8c3f780b7e..109246587b 100644
--- a/desc/objectives/utils.py
+++ b/desc/objectives/utils.py
@@ -144,7 +144,6 @@ def recover(x_reduced):
 
     # check that all constraints are actually satisfiable
     params = objective.unpack_state(xp, False)
-    print(params)
     for con in constraint.objectives:
         # FIXME: xpi is not always correct -- missing params?
         xpi = [params[i] for i, t in enumerate(objective.things) if t in con.things]

From 34cb43d63f3dcd41a702e52741433fc7e74dbab0 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Wed, 24 Apr 2024 13:53:25 -0400
Subject: [PATCH 22/50] proximal projection hack

---
 desc/optimize/_constraint_wrappers.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py
index 1c6453d1a2..93bf5385ae 100644
--- a/desc/optimize/_constraint_wrappers.py
+++ b/desc/optimize/_constraint_wrappers.py
@@ -681,6 +681,13 @@ def unpack_state(self, x, per_objective=True):
             if t is self._eq:
                 xi_splits = np.cumsum([self._eq.dimensions[arg] for arg in self._args])
                 p = {arg: xis for arg, xis in zip(self._args, jnp.split(xi, xi_splits))}
+                p.update(  # add in dummy values for missing parameters
+                    {
+                        arg: jnp.zeros_like(xis)
+                        for arg, xis in t.params_dict.items()
+                        if arg not in self._args  # R_lmn, Z_lmn, L_lmn, Ra_n, Za_n
+                    }
+                )
                 params += [p]
             else:
                 params += [t.unpack_params(xi)]

From 2818ecb85d658d6da2da1b553cc0f110d501d70d Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 10:29:11 -0500
Subject: [PATCH 23/50] cast indices to array in test

---
 tests/test_optimizer.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 64f71ee2b3..318f023b92 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1010,7 +1010,9 @@ def test_optimize_multiple_things_different_order():
         # don't let eq vary
         FixParameter(eq),
         # only let the minor radius of the surface vary
-        FixParameter(surf, params=["R_lmn"], indices=surf.R_basis.get_idx(M=0, N=0)),
+        FixParameter(
+            surf, params=["R_lmn"], indices=np.array(surf.R_basis.get_idx(M=0, N=0))
+        ),
     )
 
     target_dist = 1
@@ -1029,7 +1031,7 @@ def test_optimize_multiple_things_different_order():
     optimizer = Optimizer("lsq-exact")
 
     # ensure it runs when (eq,surf) are passed
-    (eq1, surf1), result = optimizer.optimize(
+    (eq1, surf1), _ = optimizer.optimize(
         (eq, surf), objective, constraints, verbose=3, maxiter=15, copy=True
     )
     # ensure surface changed correctly
@@ -1050,7 +1052,9 @@ def test_optimize_multiple_things_different_order():
         # don't let eq vary
         FixParameter(eq),
         # only let the minor radius of the surface vary
-        FixParameter(surf, params=["R_lmn"], indices=surf.R_basis.get_idx(M=0, N=0)),
+        FixParameter(
+            surf, params=["R_lmn"], indices=np.array(surf.R_basis.get_idx(M=0, N=0))
+        ),
     )
     obj = PlasmaVesselDistance(
         surface=surf,
@@ -1063,7 +1067,7 @@ def test_optimize_multiple_things_different_order():
     objective = ObjectiveFunction((obj,))
     # ensure it runs when (surf,eq) are passed which is opposite
     # the order of objective.things
-    (surf2, eq2), result = optimizer.optimize(
+    (surf2, eq2), _ = optimizer.optimize(
         (surf, eq), objective, constraints, verbose=3, maxiter=15, copy=True
     )
 

From 414f296d3adb7d2b6bf9d11a2cba4f1cf75bfc83 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 12:54:02 -0500
Subject: [PATCH 24/50] refactor broadcast_tree

---
 desc/utils.py       | 100 ++++++++++++++----------------
 tests/test_utils.py | 145 +++++++++++++++++++++++++++++++++-----------
 2 files changed, 155 insertions(+), 90 deletions(-)

diff --git a/desc/utils.py b/desc/utils.py
index 40f31e96c4..de1ecc5c6d 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -8,7 +8,7 @@
 from scipy.special import factorial
 from termcolor import colored
 
-from desc.backend import fori_loop, jit, jnp, tree_structure, treedef_is_leaf
+from desc.backend import fori_loop, jit, jnp
 
 
 class Timer:
@@ -610,20 +610,19 @@ def is_any_instance(things, cls):
     return any([isinstance(t, cls) for t in things])
 
 
-def broadcast_tree(tree_in, tree_out, value=False, sort=False):
+def broadcast_tree(tree_in, tree_out):
     """Broadcast tree_in to the same pytree structure as tree_out.
 
+    Both trees must be nested lists of dicts with string keys and in array values.
+    Or the values can be bools, where False broadcasts to an empty int array and True
+    broadcasts to the corresponding int array from tree_out.
+
     Parameters
     ----------
     tree_in : pytree
         Tree to broadcast.
     tree_out : pytree
         Tree with structure to broadcast to.
-    value : leaf, optional
-        Leaf value to pad branches with if not present in tree_in. Default = False.
-    sort : bool, optional
-        If True, checks that leaves of tree_in are also leaves of tree_out, and keeps
-        them in the same position on their branches when broadcasting. Default = False.
 
     Returns
     -------
@@ -631,52 +630,43 @@ def broadcast_tree(tree_in, tree_out, value=False, sort=False):
         Tree with the leaves of tree_in broadcast to the structure of tree_out.
 
     """
-    if not isinstance(tree_in, list):
-        tree_in = [tree_in]
-    if not isinstance(tree_out, list):
-        tree_out = [tree_out]
-
-    def isemptylist(x):
-        if isinstance(x, list):
-            return not x
-        return False
-
-    where_leaves_in = [  # tree_in can have empty branches that are not leaves
-        treedef_is_leaf(tree_structure(branch)) and not isemptylist(branch)
-        for branch in tree_in
-    ]
-    where_leaves_out = [treedef_is_leaf(tree_structure(branch)) for branch in tree_out]
-    all_leaves_in = all(where_leaves_in)
-    all_leaves_out = all(where_leaves_out)
-    if any(where_leaves_in) and not all_leaves_in:
-        raise ValueError("base layer of tree_in must be all leaves and no branches")
-    if any(where_leaves_out) and not all_leaves_out:
-        raise ValueError("base layer of tree_out must be all leaves and no branches")
-    if all_leaves_in and all_leaves_out:  # both trees at leaf layer
-        if len(tree_in) <= len(tree_out):
-            if sort:
-                if set(tree_in) <= set(tree_out):
-                    return [leaf if leaf in tree_in else value for leaf in tree_out]
-                else:
-                    raise ValueError(
-                        "leaves of tree_in must be a subset of the leaves in tree_out "
-                        + "if sort=True"
-                    )
-            else:
-                return tree_in + [value for _ in range(len(tree_out) - len(tree_in))]
-        else:
-            raise ValueError("tree_in cannot have more leaves than tree_out")
-    elif all_leaves_in and not all_leaves_out:  # tree_out is deeper than tree_in
-        return [broadcast_tree(tree_in, branch, sort=sort) for branch in tree_out]
-    elif all_leaves_out and not all_leaves_in:
-        raise ValueError("tree_in cannot have a deeper structure than tree_out")
-    else:  # both trees at branch layers
-        if len(tree_in) == len(tree_out):
-            return [
-                broadcast_tree(tree_in[k], tree_out[k], sort=sort)
-                for k in range(len(tree_out))
-            ]
-        else:
-            raise ValueError(
-                "tree_in must have the same number of branches as tree_out"
+    # both trees at leaf layer
+    if isinstance(tree_in, dict) and isinstance(tree_out, dict):
+        tree_new = tree_in.copy()
+        for key, value in tree_in.items():
+            errorif(
+                key not in tree_out.keys(),
+                ValueError,
+                "dict keys of tree_in must be a subset of those in tree_out",
+            )
+            if isinstance(value, bool) and value:
+                tree_new[key] = np.atleast_1d(tree_out[key]).astype(dtype=int)
+            if isinstance(value, bool) and not value:
+                tree_new[key] = np.array([], dtype=int)
+            value = np.atleast_1d(value).astype(dtype=int)
+        for key, value in tree_out.items():
+            if key not in tree_new.keys():
+                tree_new[key] = np.array([], dtype=int)
+            errorif(
+                not np.all(np.isin(tree_new[key], value)),
+                ValueError,
+                "dict values of tree_in must be a subset of those in tree_out",
             )
+        return tree_new
+    # tree_out is deeper than tree_in
+    elif isinstance(tree_in, dict) and isinstance(tree_out, list):
+        return [broadcast_tree(tree_in.copy(), branch) for branch in tree_out]
+    # both trees at branch layers
+    elif isinstance(tree_in, list) and isinstance(tree_out, list):
+        errorif(
+            len(tree_in) != len(tree_out),
+            ValueError,
+            "tree_in must have the same number of branches as tree_out",
+        )
+        return [broadcast_tree(tree_in[k], tree_out[k]) for k in range(len(tree_out))]
+    # tree_in is deeper than tree_out
+    elif isinstance(tree_in, list) and isinstance(tree_out, dict):
+        raise ValueError("tree_in cannot have a deeper structure than tree_out")
+    # invalid tree structure
+    else:
+        raise ValueError("trees must be nested lists of dicts")
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 51b905d57a..b83850022d 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from desc.backend import tree_map
+from desc.backend import tree_leaves, tree_structure
 from desc.grid import LinearGrid
 from desc.utils import broadcast_tree, isalmostequal, islinspaced
 
@@ -59,64 +59,139 @@ def test_islinspaced():
 @pytest.mark.unit
 def test_broadcast_tree():
     """Test that broadcast_tree works on various pytree structures."""
-    tree_out = [[1, 2, 3], [[4], [[5, 6], [7]]]]
-
-    # tree of tuples, not lists
-    tree_in = [(0, 1), [2]]
+    tree_out = [
+        {"a": np.arange(1), "b": np.arange(2), "c": np.arange(3)},
+        [
+            {"a": np.arange(2)},
+            [{"a": np.arange(1), "d": np.arange(3)}, {"a": np.arange(2)}],
+        ],
+    ]
+
+    # tree with tuples, not lists
+    tree_in = [{}, ({}, [{}, {}])]
     with pytest.raises(ValueError):
         _ = broadcast_tree(tree_in, tree_out)
 
-    # tree with too many leaves per branch
-    tree_in = [1, 2]
+    # tree_in is deeper than tree_out
+    tree_in = [
+        [{"a": np.arange(1)}, {"b": np.arange(2), "c": np.arange(3)}],
+        [{}, [{}, {"a": np.arange(2)}]],
+    ]
     with pytest.raises(ValueError):
         _ = broadcast_tree(tree_in, tree_out)
 
-    # tree with a mix of leaves and branches at the same layer
-    tree_in = [[1, 2], 3, [4]]
+    # tree_in has different number of branches as tree_out
+    tree_in = [{}, [{}, [{}]]]
     with pytest.raises(ValueError):
         _ = broadcast_tree(tree_in, tree_out)
 
-    # tree_in is deeper than tree_out
-    tree_in = [[[1], [2, 3]], [[4], [[[5], [6]], [7]]]]
+    # tree with incorrect keys
+    tree_in = [{"a": np.arange(1), "b": np.arange(2)}, {"d": np.arange(2)}]
     with pytest.raises(ValueError):
         _ = broadcast_tree(tree_in, tree_out)
 
-    # tree_in leaves not in tree_out
-    tree_in = [[1, 2], [[3], [4]]]
+    # tree with incorrect values
+    tree_in = [{"a": np.arange(1), "b": np.arange(2)}, {"a": np.arange(2)}]
     with pytest.raises(ValueError):
-        tree = broadcast_tree(tree_in, tree_out, sort=True)
+        _ = broadcast_tree(tree_in, tree_out)
 
     # tree with proper structure already does not change
-    tree_in = tree_map(lambda x: x * 2, tree_out)
+    tree_in = tree_out.copy()
     tree = broadcast_tree(tree_in, tree_out)
     assert tree == tree_in
 
     # broadcast single leaf to full tree
-    tree_in = 0
+    tree_in = {"a": np.arange(1)}
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[0, False, False], [[0], [[0, False], [0]]]]
+    assert tree_structure(tree) == tree_structure(tree_out)
+    tree_correct = [
+        {"a": np.arange(1), "b": np.array([], dtype=int), "c": np.array([], dtype=int)},
+        [
+            {"a": np.arange(1)},
+            [{"a": np.arange(1), "d": np.array([], dtype=int)}, {"a": np.arange(1)}],
+        ],
+    ]
+    for leaf, leaf_correct in zip(tree_leaves(tree), tree_leaves(tree_correct)):
+        np.testing.assert_allclose(leaf, leaf_correct)
 
     # broadcast from only major branches
-    tree_in = [[1, 2], [3]]
+    tree_in = [{"b": np.arange(2), "c": np.arange(1, 3)}, {"a": np.arange(1)}]
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[1, 2, False], [[3], [[3, False], [3]]]]
+    assert tree_structure(tree) == tree_structure(tree_out)
+    tree_correct = [
+        {"a": np.array([], dtype=int), "b": np.arange(2), "c": np.arange(1, 3)},
+        [
+            {"a": np.arange(1)},
+            [{"a": np.arange(1), "d": np.array([], dtype=int)}, {"a": np.arange(1)}],
+        ],
+    ]
+    for leaf, leaf_correct in zip(tree_leaves(tree), tree_leaves(tree_correct)):
+        np.testing.assert_allclose(leaf, leaf_correct)
 
     # broadcast from minor branches
-    tree_in = [[1, 2], [[3], [4]]]
+    tree_in = [
+        {"b": np.arange(2), "c": np.arange(1, 3)},
+        [{"a": np.arange(2)}, {"a": np.arange(1)}],
+    ]
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[1, 2, False], [[3], [[4, False], [4]]]]
-
-    # sort order of leaves
-    tree_in = [[3, 1], [[4], [[6], []]]]
-    tree = broadcast_tree(tree_in, tree_out, sort=True)
-    assert tree == [[1, False, 3], [[4], [[False, 6], [False]]]]
-
-    # tree_in with empty branches
-    tree_in = [[], [[1], [[2], []]]]
+    assert tree_structure(tree) == tree_structure(tree_out)
+    tree_correct = [
+        {"a": np.array([], dtype=int), "b": np.arange(2), "c": np.arange(1, 3)},
+        [
+            {"a": np.arange(2)},
+            [{"a": np.arange(1), "d": np.array([], dtype=int)}, {"a": np.arange(1)}],
+        ],
+    ]
+    for leaf, leaf_correct in zip(tree_leaves(tree), tree_leaves(tree_correct)):
+        np.testing.assert_allclose(leaf, leaf_correct)
+
+    # tree_in with empty dicts and arrays
+    tree_in = [
+        {},
+        [
+            {"a": np.array([], dtype=int)},
+            [{"a": np.arange(1), "d": np.array([0, 2], dtype=int)}, {}],
+        ],
+    ]
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == [[False, False, False], [[1], [[2, False], [False]]]]
-
-    # custom fill value
-    tree_in = [[1, 2], [[3], [4]]]
-    tree = broadcast_tree(tree_in, tree_out, value=0)
-    assert tree == [[1, 2, 0], [[3], [[4, 0], [4]]]]
+    assert tree_structure(tree) == tree_structure(tree_out)
+    tree_correct = [
+        {
+            "a": np.array([], dtype=int),
+            "b": np.array([], dtype=int),
+            "c": np.array([], dtype=int),
+        },
+        [
+            {"a": np.array([], dtype=int)},
+            [
+                {"a": np.arange(1), "d": np.array([0, 2], dtype=int)},
+                {"a": np.array([], dtype=int)},
+            ],
+        ],
+    ]
+    for leaf, leaf_correct in zip(tree_leaves(tree), tree_leaves(tree_correct)):
+        np.testing.assert_allclose(leaf, leaf_correct)
+
+    # tree_in with bool values
+    tree_in = [
+        {"a": False, "b": True, "c": np.array([0, 2], dtype=int)},
+        [
+            {"a": True},
+            [{"a": False, "d": np.arange(2)}, {"a": True}],
+        ],
+    ]
+    tree = broadcast_tree(tree_in, tree_out)
+    assert tree_structure(tree) == tree_structure(tree_out)
+    tree_correct = [
+        {
+            "a": np.array([], dtype=int),
+            "b": np.arange(2),
+            "c": np.array([0, 2], dtype=int),
+        },
+        [
+            {"a": np.arange(2)},
+            [{"a": np.array([], dtype=int), "d": np.arange(2)}, {"a": np.arange(2)}],
+        ],
+    ]
+    for leaf, leaf_correct in zip(tree_leaves(tree), tree_leaves(tree_correct)):
+        np.testing.assert_allclose(leaf, leaf_correct)

From 1a5ae5122b5548a877ab92ab7c82440c1ac750e9 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 16:16:09 -0500
Subject: [PATCH 25/50] refactor FixParameter objective

---
 desc/objectives/linear_objectives.py | 102 +++++----------------------
 tests/test_examples.py               |  21 +++---
 tests/test_linear_objectives.py      |  27 +++++--
 tests/test_optimizer.py              |  34 +++++----
 4 files changed, 66 insertions(+), 118 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index c582bf4025..47b17d1658 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -12,16 +12,9 @@
 import numpy as np
 from termcolor import colored
 
-from desc.backend import (
-    jnp,
-    tree_leaves,
-    tree_map,
-    tree_structure,
-    tree_unflatten,
-    treedef_is_leaf,
-)
+from desc.backend import jnp, tree_leaves, tree_map, tree_structure
 from desc.basis import zernike_radial, zernike_radial_coeffs
-from desc.utils import broadcast_tree, errorif, flatten_list, setdefault
+from desc.utils import broadcast_tree, errorif, setdefault
 
 from .normalization import compute_scaling_factors
 from .objective_funs import _Objective
@@ -79,11 +72,11 @@ class FixParameter(_Objective):
     ----------
     thing : Optimizable
         Object whose degrees of freedom are being fixed.
-    params : pytree with leaves of type str
-        Names of parameters to fix. Defaults to all parameters.
-    indices : pytree with leaves of type array
-        Indices to fix for each parameter in params.
-        The default value indices=True fixes all indices.
+    params : nested list of dicts
+        Dict keys are the names of parameters to fix (str), and dict values are the
+        indices to fix for each corresonding parameter (int array).
+        Must have the same pytree structure as thing.params_dict.
+        The default is to fix all indices of all parameters.
     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.
@@ -115,7 +108,6 @@ def __init__(
         self,
         thing,
         params=None,
-        indices=True,
         target=None,
         bounds=None,
         weight=1,
@@ -124,7 +116,6 @@ def __init__(
         name="Fixed parameters",  # TODO: add params
     ):
         self._params = params
-        self._indices = indices
         super().__init__(
             things=thing,
             target=target,
@@ -147,77 +138,22 @@ def build(self, use_jit=False, verbose=1):
 
         """
         thing = self.things[0]
-        structure = tree_structure(thing.optimizable_params)
-
-        def leaves_per_branch(tree_tuple):
-            """Get the number of leaves per branch in a pytree, in flattened order."""
-            tree, lengths = tree_tuple
-            if all([treedef_is_leaf(tree_structure(x)) for x in tree]):
-                lengths.append(len(tree))
-                return [], lengths
-            else:
-                return (
-                    [leaves_per_branch((branch, lengths)) for branch in tree],
-                    lengths,
-                )
-
-        # indices of flattened tree leaves
-        _, lengths = leaves_per_branch((thing.optimizable_params, []))
-        starts = np.insert(np.cumsum(lengths), 0, 0)
-        leaf_indices = [
-            np.arange(start, start + length, dtype=int)
-            for start, length in zip(starts, lengths)
-        ]
-
-        # set params
-        self._params = setdefault(self._params, thing.optimizable_params)
-        self._params = broadcast_tree(self._params, thing.optimizable_params, sort=True)
-        assert tree_structure(self._params) == structure
-        params_leaves = tree_leaves(self._params)
-
-        # sort params to the same flattend leaf order as indices
-        # this is necessary because of JAX GitHub Issue #4085
-        params_leaves = flatten_list(
-            [
-                [
-                    [params_leaves[i] for i in idx][j]
-                    for j in np.argsort(
-                        np.atleast_1d(tree_leaves(thing.optimizable_params))[idx]
-                    )
-                ]
-                for idx in leaf_indices
-            ]
-        )
 
-        # set indices
-        if isinstance(self._indices, bool):  # default to all indices
-            errorif(not self._indices, ValueError, "indices cannot be False")
-            indices = tree_map(lambda dim: np.arange(dim, dtype=int), thing.dimensions)
-            indices_leaves = [
-                idx if param else False
-                for idx, param in zip(tree_leaves(indices), params_leaves)
-            ]
-            self._indices = tree_unflatten(structure, indices_leaves)
-        self._indices = broadcast_tree(  # indices cannot be sorted like the params!
-            self._indices, thing.optimizable_params, sort=False
-        )
-        assert tree_structure(self._indices) == structure
-        self._indices_leaves = tree_leaves(self._indices)
-        self._indices_leaves = [  # replace False leaves with empy arrays
-            indices if not isinstance(indices, bool) else np.array([], dtype=int)
-            for indices in self._indices_leaves
-        ]
+        # default params
+        default_params = tree_map(lambda dim: np.arange(dim), thing.dimensions)
+        self._params = setdefault(self._params, default_params)
+        self._params = broadcast_tree(self._params, default_params)
+        self._indices = tree_leaves(self._params)
+        assert tree_structure(self._params) == tree_structure(default_params)
 
-        self._dim_f = sum(idx.size for idx in self._indices_leaves)
+        self._dim_f = sum(idx.size for idx in self._indices)
 
-        # set default target
+        # default target
         if self.target is None and self.bounds is None:
             self.target = np.concatenate(
                 [
                     np.atleast_1d(param[idx])
-                    for param, idx in zip(
-                        tree_leaves(thing.params_dict), self._indices_leaves
-                    )
+                    for param, idx in zip(tree_leaves(thing.params_dict), self._indices)
                 ]
             )
 
@@ -240,14 +176,10 @@ def compute(self, params, constants=None):
             Fixed degree of freedom errors.
 
         """
-        params_leaves = tree_leaves(params)
-        # FIXME: should still work if only a subset of params get passed in
-        # TODO: need to cast to full dict
-        assert len(params_leaves) == len(self._indices_leaves)
         return jnp.concatenate(
             [
                 jnp.atleast_1d(param[idx])
-                for param, idx in zip(params_leaves, self._indices_leaves)
+                for param, idx in zip(tree_leaves(params), self._indices)
             ]
         )
 
diff --git a/tests/test_examples.py b/tests/test_examples.py
index a6a449f290..65e2c2d524 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -640,8 +640,8 @@ def test_multiobject_optimization_al():
     constraints = (
         ForceBalance(eq=eq, bounds=(-1e-4, 1e-4), normalize_target=False),
         FixPressure(eq=eq),
-        FixParameter(surf, ["R_lmn", "Z_lmn"], [np.array([0]), np.array([-1])]),
-        FixParameter(eq, ["Psi", "i_l"]),
+        FixParameter(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}),
+        FixParameter(eq, {"Psi": True, "i_l": True}),
         FixBoundaryR(eq, modes=[[0, 0, 0]]),
         PlasmaVesselDistance(surface=surf, eq=eq, target=1),
     )
@@ -679,8 +679,8 @@ def test_multiobject_optimization_prox():
     constraints = (
         ForceBalance(eq=eq),
         FixPressure(eq=eq),
-        FixParameter(surf, ["R_lmn", "Z_lmn"], [np.array([0]), np.array([-1])]),
-        FixParameter(eq, ["Psi", "i_l"]),
+        FixParameter(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}),
+        FixParameter(eq, {"Psi": True, "i_l": True}),
         FixBoundaryR(eq, modes=[[0, 0, 0]]),
     )
 
@@ -999,19 +999,14 @@ def test_only_non_eq_optimization():
     """Test for optimizing only a non-eq object."""
     eq = get("DSHAPE")
     surf = eq.surface
-
     surf.change_resolution(M=eq.M, N=eq.N)
     constraints = (
-        FixParameter(
-            surf, params="R_lmn", indices=np.array(surf.R_basis.get_idx(0, 0, 0))
-        ),
+        FixParameter(surf, {"R_lmn": np.array(surf.R_basis.get_idx(0, 0, 0))}),
     )
-
     obj = PrincipalCurvature(surf, target=1)
-
     objective = ObjectiveFunction((obj,))
     optimizer = Optimizer("lsq-exact")
-    (surf), result = optimizer.optimize(
+    (surf), _ = optimizer.optimize(
         (surf), objective, constraints, verbose=3, maxiter=100
     )
     surf = surf[0]
@@ -1258,7 +1253,7 @@ def test_quadratic_flux_optimization_with_analytic_field():
 
     optimizer = Optimizer("lsq-exact")
 
-    constraints = (FixParameter(field, "R0"),)
+    constraints = (FixParameter(field, {"R0": True}),)
     quadflux_obj = QuadraticFlux(
         eq=eq,
         field=field,
@@ -1287,7 +1282,7 @@ def test_second_stage_optimization():
     eq = get("DSHAPE")
     field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
     objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
-    constraints = FixParameter(field, [["R0"], []])
+    constraints = FixParameter(field, [{"R0": True}, {}])
     optimizer = Optimizer("lsq-exact")
     (field,), _ = optimizer.optimize(
         things=field, objective=objective, constraints=constraints, verbose=2
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index e14124c2e1..fa523a2b59 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -939,12 +939,29 @@ def test_fix_subset_of_params_in_collection():
     xy_coil = FourierXYZCoil()
     full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xy_coil))
 
-    # TODO: use a better example to test broadcasting for a whole coilset
     params = [
-        [["current"], ["center", "normal", "r_n"], ["current", "shift", "rotmat"], []],
-        [["current", "Z_n"], ["R_n", "Z_n"], ["Z_n"]],
-        ["Z_n", "shift", "rotmat"],
+        [
+            {"current": True},
+            {"center": True, "normal": np.array([1])},
+            {"r_n": True},
+            {},
+        ],
+        {"shift": True, "rotmat": True},
+        {"X_n": np.array([1, 2]), "Y_n": False, "Z_n": np.array([0])},
     ]
+    target = np.concatenate(
+        (
+            np.array([1, 2, 0, 0, 1, 1]),
+            np.eye(3).flatten(),
+            np.array([0, 0, 0]),
+            np.eye(3).flatten(),
+            np.array([0, 0, 1]),
+            np.eye(3).flatten(),
+            np.array([0, 0, 2]),
+            np.array([10, 2, -2]),
+        )
+    )
+
     obj = FixParameter(full_coilset, params)
     obj.build()
-    assert obj.dim_f == 41
+    np.testing.assert_allclose(obj.target, target)
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 318f023b92..8955008b21 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -895,9 +895,7 @@ def test_constrained_AL_lsq():
     constraints = (
         FixBoundaryR(eq=eq, modes=[0, 0, 0]),  # fix specified major axis position
         FixPressure(eq=eq),  # fix pressure profile
-        FixParameter(
-            eq, "i_l", bounds=(eq.i_l * 0.9, eq.i_l * 1.1)
-        ),  # linear inequality
+        FixIota(eq, bounds=(eq.i_l * 0.9, eq.i_l * 1.1)),  # linear inequality
         FixPsi(eq=eq, bounds=(eq.Psi * 0.99, eq.Psi * 1.01)),  # linear inequality
     )
     # some random constraints to keep the shape from getting wacky
@@ -1007,11 +1005,9 @@ def test_optimize_multiple_things_different_order():
         NFP=eq.NFP,
     )
     constraints = (
-        # don't let eq vary
-        FixParameter(eq),
-        # only let the minor radius of the surface vary
-        FixParameter(
-            surf, params=["R_lmn"], indices=np.array(surf.R_basis.get_idx(M=0, N=0))
+        FixParameter(eq),  # don't let eq vary
+        FixParameter(  # only let the minor radius of the surface vary
+            surf, params={"R_lmn": np.array(surf.R_basis.get_idx(M=0, N=0))}
         ),
     )
 
@@ -1049,11 +1045,9 @@ def test_optimize_multiple_things_different_order():
 
     # fresh start
     constraints = (
-        # don't let eq vary
-        FixParameter(eq),
-        # only let the minor radius of the surface vary
-        FixParameter(
-            surf, params=["R_lmn"], indices=np.array(surf.R_basis.get_idx(M=0, N=0))
+        FixParameter(eq),  # don't let eq vary
+        FixParameter(  # only let the minor radius of the surface vary
+            surf, params={"R_lmn": np.array(surf.R_basis.get_idx(M=0, N=0))}
         ),
     )
     obj = PlasmaVesselDistance(
@@ -1091,8 +1085,18 @@ def test_optimize_with_single_constraint():
     eq = Equilibrium()
     optimizer = Optimizer("lsq-exact")
     objectective = ObjectiveFunction(GenericObjective("|B|", eq), use_jit=False)
-    constraints = FixParameter(  # Psi is not constrained
-        eq, ["R_lmn", "Z_lmn", "L_lmn", "Rb_lmn", "Zb_lmn", "p_l", "c_l"]
+    constraints = FixParameter(
+        eq,
+        {
+            "R_lmn": True,
+            "Z_lmn": True,
+            "L_lmn": True,
+            "Rb_lmn": True,
+            "Zb_lmn": True,
+            "p_l": True,
+            "c_l": True,
+            "Psi": False,  # Psi is not constrained
+        },
     )
 
     # test depends on verbose > 0

From 70fd03f7951c9b89c9c78f10728dfbb1ec3c19f2 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 16:23:54 -0500
Subject: [PATCH 26/50] add dtype option to broadcast_tree

---
 desc/utils.py | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/desc/utils.py b/desc/utils.py
index de1ecc5c6d..0c638f3fd2 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -610,12 +610,12 @@ def is_any_instance(things, cls):
     return any([isinstance(t, cls) for t in things])
 
 
-def broadcast_tree(tree_in, tree_out):
+def broadcast_tree(tree_in, tree_out, dtype=int):
     """Broadcast tree_in to the same pytree structure as tree_out.
 
-    Both trees must be nested lists of dicts with string keys and in array values.
-    Or the values can be bools, where False broadcasts to an empty int array and True
-    broadcasts to the corresponding int array from tree_out.
+    Both trees must be nested lists of dicts with string keys and array values.
+    Or the values can be bools, where False broadcasts to an empty array and True
+    broadcasts to the corresponding array from tree_out.
 
     Parameters
     ----------
@@ -623,6 +623,8 @@ def broadcast_tree(tree_in, tree_out):
         Tree to broadcast.
     tree_out : pytree
         Tree with structure to broadcast to.
+    dtype : optional
+        Data type of array values. Default = int.
 
     Returns
     -------
@@ -640,13 +642,13 @@ def broadcast_tree(tree_in, tree_out):
                 "dict keys of tree_in must be a subset of those in tree_out",
             )
             if isinstance(value, bool) and value:
-                tree_new[key] = np.atleast_1d(tree_out[key]).astype(dtype=int)
+                tree_new[key] = np.atleast_1d(tree_out[key]).astype(dtype=dtype)
             if isinstance(value, bool) and not value:
-                tree_new[key] = np.array([], dtype=int)
-            value = np.atleast_1d(value).astype(dtype=int)
+                tree_new[key] = np.array([], dtype=dtype)
+            value = np.atleast_1d(value).astype(dtype=dtype)
         for key, value in tree_out.items():
             if key not in tree_new.keys():
-                tree_new[key] = np.array([], dtype=int)
+                tree_new[key] = np.array([], dtype=dtype)
             errorif(
                 not np.all(np.isin(tree_new[key], value)),
                 ValueError,
@@ -656,7 +658,7 @@ def broadcast_tree(tree_in, tree_out):
     # tree_out is deeper than tree_in
     elif isinstance(tree_in, dict) and isinstance(tree_out, list):
         return [broadcast_tree(tree_in.copy(), branch) for branch in tree_out]
-    # both trees at branch layers
+    # both trees at branch layer
     elif isinstance(tree_in, list) and isinstance(tree_out, list):
         errorif(
             len(tree_in) != len(tree_out),

From b997d635c9c0477a11b9f32ce7866f0e40bf4db4 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 16:54:16 -0500
Subject: [PATCH 27/50] replace some FixedObjectives with FixParameter

---
 desc/objectives/getters.py           |   3 +-
 desc/objectives/linear_objectives.py | 258 ++-------------------------
 2 files changed, 18 insertions(+), 243 deletions(-)

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index bf73d7b7c8..edae26a828 100644
--- a/desc/objectives/getters.py
+++ b/desc/objectives/getters.py
@@ -238,8 +238,7 @@ def maybe_add_self_consistency(thing, constraints):
             constraints += (AxisZSelfConsistency(eq=thing),)
 
     # Curve
-    # FIXME: make this work for CoilSet
-    elif not hasattr(thing, "__len__") and {"shift", "rotmat"} <= set(
+    elif {"shift", "rotmat"} <= set(
         unique_list(flatten_list(thing.optimizable_params))[0]
     ):
         if not is_any_instance(constraints, FixCurveShift):
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 47b17d1658..788aa22c4f 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -113,7 +113,7 @@ def __init__(
         weight=1,
         normalize=False,
         normalize_target=False,
-        name="Fixed parameters",  # TODO: add params
+        name="Fixed parameters",
     ):
         self._params = params
         super().__init__(
@@ -3123,7 +3123,7 @@ def compute(self, params, constants=None):
         return params["Zeff_l"][self._idx]
 
 
-class FixPsi(_FixedObjective):
+class FixPsi(FixParameter):
     """Fixes total toroidal magnetic flux within the last closed flux surface.
 
     Parameters
@@ -3164,9 +3164,9 @@ def __init__(
         normalize_target=True,
         name="fixed-Psi",
     ):
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"Psi": True},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -3174,52 +3174,12 @@ def __init__(
             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.
-
-        """
-        eq = self.things[0]
-        self._dim_f = 1
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.Psi, None, np.array([0])
-        )
-
         if self._normalize:
-            scales = compute_scaling_factors(eq)
+            scales = compute_scaling_factors(self.things[0])
             self._normalization = scales["Psi"]
 
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed-Psi error.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Total toroidal magnetic flux error (Wb).
-
-        """
-        return params["Psi"]
 
-
-class FixCurveShift(_FixedObjective):
+class FixCurveShift(FixParameter):
     """Fixes Curve.shift attribute, which is redundant with other Curve params.
 
     Parameters
@@ -3259,9 +3219,9 @@ def __init__(
         normalize_target=True,
         name="fixed-shift",
     ):
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=curve,
+            thing=curve,
+            params={"shift": True},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -3270,50 +3230,8 @@ def __init__(
             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.
-
-        """
-        curve = self.things[0]
-        self._dim_f = curve.shift.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, curve.shift, None, np.arange(self._dim_f)
-        )
-
-        if self._normalize:
-            self._normalization = 1
 
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed-shift error.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of curve degrees of freedom, eg Curve.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Curve shift (m).
-
-        """
-        return params["shift"]
-
-
-class FixCurveRotation(_FixedObjective):
+class FixCurveRotation(FixParameter):
     """Fixes Curve.rotmat attribute, which is redundant with other Curve params.
 
     Parameters
@@ -3353,9 +3271,9 @@ def __init__(
         normalize_target=True,
         name="fixed-rotation",
     ):
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=curve,
+            thing=curve,
+            params={"rotmat": True},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -3364,50 +3282,8 @@ def __init__(
             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.
-
-        """
-        curve = self.things[0]
-        self._dim_f = curve.rotmat.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, curve.rotmat, None, np.arange(self._dim_f)
-        )
-
-        if self._normalize:
-            self._normalization = 1
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed-rotation error.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of curve degrees of freedom, eg Curve.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Curve rotation matrix (rad).
-
-        """
-        return params["rotmat"]
-
-
-class FixOmniWell(_FixedObjective):
+class FixOmniWell(FixParameter):
     """Fixes OmnigenousField.B_lm coefficients.
 
     Parameters
@@ -3449,11 +3325,9 @@ def __init__(
         indices=True,
         name="fixed omnigenity well",
     ):
-        self._field = field
-        self._indices = indices
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=field,
+            thing=field,
+            params={"B_lm": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -3462,56 +3336,8 @@ def __init__(
             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.
-
-        """
-        field = self.things[0]
-
-        # find indices to fix
-        if self._indices is False or self._indices is None:  # no indices to fix
-            self._idx = np.array([], dtype=int)
-        elif self._indices is True:  # all indices
-            self._idx = np.arange(np.size(self._field.B_lm))
-        else:  # specified indices
-            self._idx = np.atleast_1d(self._indices)
-
-        self._dim_f = self._idx.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, field.B_lm[self._idx], None, self._idx
-        )
-
-        super().build(use_jit=use_jit, verbose=verbose)
 
-    def compute(self, params, constants=None):
-        """Compute fixed omnigenity well error.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of field degrees of freedom, eg OmnigenousField.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed well shape error.
-
-        """
-        return params["B_lm"][self._idx]
-
-
-class FixOmniMap(_FixedObjective):
+class FixOmniMap(FixParameter):
     """Fixes OmnigenousField.x_lmn coefficients.
 
     Parameters
@@ -3553,11 +3379,9 @@ def __init__(
         indices=True,
         name="fixed omnigenity map",
     ):
-        self._field = field
-        self._indices = indices
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=field,
+            thing=field,
+            params={"x_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -3566,54 +3390,6 @@ def __init__(
             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.
-
-        """
-        field = self.things[0]
-
-        # find indices to fix
-        if self._indices is False or self._indices is None:  # no indices to fix
-            self._idx = np.array([], dtype=int)
-        elif self._indices is True:  # all indices
-            self._idx = np.arange(np.size(self._field.x_lmn))
-        else:  # specified indices
-            self._idx = np.atleast_1d(self._indices)
-
-        self._dim_f = self._idx.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, field.x_lmn[self._idx], None, self._idx
-        )
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed omnigenity map error.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of field degrees of freedom, eg OmnigenousField.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed omnigenity map error.
-
-        """
-        return params["x_lmn"][self._idx]
-
 
 class FixOmniBmax(_FixedObjective):
     """Ensures the B_max contour is straight in Boozer coordinates.

From 2d7c033cc7f82fcf3e430ddec11dce55786cfb4e Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 17:05:58 -0500
Subject: [PATCH 28/50] add FixSheetCurrent objective

---
 desc/objectives/__init__.py          |  1 +
 desc/objectives/getters.py           | 16 +++------
 desc/objectives/linear_objectives.py | 53 ++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py
index e0ae573e60..5357e78013 100644
--- a/desc/objectives/__init__.py
+++ b/desc/objectives/__init__.py
@@ -66,6 +66,7 @@
     FixParameter,
     FixPressure,
     FixPsi,
+    FixSheetCurrent,
     FixSumModesLambda,
     FixSumModesR,
     FixSumModesZ,
diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index edae26a828..489e5401b7 100644
--- a/desc/objectives/getters.py
+++ b/desc/objectives/getters.py
@@ -1,7 +1,5 @@
 """Utilities for getting standard groups of objectives and constraints."""
 
-import numpy as np
-
 from desc.utils import flatten_list, is_any_instance, unique_list
 
 from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance
@@ -24,9 +22,9 @@
     FixIonTemperature,
     FixIota,
     FixLambdaGauge,
-    FixParameter,
     FixPressure,
     FixPsi,
+    FixSheetCurrent,
 )
 from .nae_utils import calc_zeroth_order_lambda, make_RZ_cons_1st_order
 from .objective_funs import ObjectiveFunction
@@ -107,9 +105,7 @@ def get_fixed_axis_constraints(eq, profiles=True, normalize=True):
                 constraints += (
                     con(eq=eq, normalize=normalize, normalize_target=normalize),
                 )
-    for param in ["I", "G", "Phi_mn"]:
-        if np.array(getattr(eq, param, [])).size:
-            constraints += (FixParameter(eq, param),)
+    constraints += (FixSheetCurrent(eq),)
 
     return constraints
 
@@ -143,9 +139,7 @@ def get_fixed_boundary_constraints(eq, profiles=True, normalize=True):
                 constraints += (
                     con(eq=eq, normalize=normalize, normalize_target=normalize),
                 )
-    for param in ["I", "G", "Phi_mn"]:
-        if np.array(getattr(eq, param, [])).size:
-            constraints += (FixParameter(eq, param),)
+    constraints += (FixSheetCurrent(eq),)
 
     return constraints
 
@@ -201,9 +195,7 @@ def get_NAE_constraints(
                 constraints += (
                     con(eq=desc_eq, normalize=normalize, normalize_target=normalize),
                 )
-    for param in ["I", "G", "Phi_mn"]:
-        if np.array(getattr(desc_eq, param, [])).size:
-            constraints += (FixParameter(desc_eq, param),)
+    constraints += (FixSheetCurrent(desc_eq),)
 
     if fix_lambda or (fix_lambda >= 0 and type(fix_lambda) is int):
         L_axis_constraints, _, _ = calc_zeroth_order_lambda(
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 788aa22c4f..2a2cf1cc56 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -3495,3 +3495,56 @@ def compute(self, params, constants=None):
         """
         f = jnp.dot(self._A, params["x_lmn"])
         return f
+
+
+class FixSheetCurrent(FixParameter):
+    """Fixes the sheet current parameters of a free-boundary equilibrium.
+
+    Parameters
+    ----------
+    eq : Equilibrium
+        Equilibrium that will be optimized to satisfy the Objective.
+    target : {float, ndarray}, optional
+        Target value(s) of the objective. Only used if bounds is None.
+        Must be broadcastable to Objective.dim_f. Default is ``target=eq.Psi``.
+    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.
+        Default is ``target=eq.Psi``.
+    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.
+    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.
+    name : str, optional
+        Name of the objective function.
+
+    """
+
+    _units = "(~)"
+    _print_value_fmt = "Fixed sheet current error: {:10.3e} "
+
+    def __init__(
+        self,
+        eq,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        name="fixed sheet current",
+    ):
+        super().__init__(
+            thing=eq,
+            params={"I": True, "G": True, "Phi_mn": True},
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            name=name,
+        )

From f88eb4b580b856ecc676a6c713e4a1e878af5b91 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 18:20:42 -0500
Subject: [PATCH 29/50] use FixParameter for axis/boundary/etc. objectives

---
 desc/basis.py                        |   5 +-
 desc/objectives/linear_objectives.py | 692 +++++----------------------
 2 files changed, 121 insertions(+), 576 deletions(-)

diff --git a/desc/basis.py b/desc/basis.py
index f5fde6f7a6..04dc2da147 100644
--- a/desc/basis.py
+++ b/desc/basis.py
@@ -114,7 +114,8 @@ def get_idx(self, L=0, M=0, N=0, error=True):
         N : int
             Toroidal mode number.
         error : bool
-            whether to raise exception if mode is not in basis, or return empty array
+            Whether to raise exception if the mode is not in the basis (default),
+            or to return an empty array.
 
         Returns
         -------
@@ -130,7 +131,7 @@ def get_idx(self, L=0, M=0, N=0, error=True):
                     "mode ({}, {}, {}) is not in basis {}".format(L, M, N, str(self))
                 ) from e
             else:
-                return np.array([]).astype(int)
+                return np.array([], dtype=int)
 
     @abstractmethod
     def _get_modes(self):
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 2a2cf1cc56..5b20f48e39 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -567,7 +567,7 @@ def compute(self, params, constants=None):
         return f
 
 
-class FixBoundaryR(_FixedObjective):
+class FixBoundaryR(FixParameter):
     """Boundary condition on the R boundary parameters.
 
     Parameters
@@ -594,16 +594,9 @@ class FixBoundaryR(_FixedObjective):
         Basis modes numbers [l,m,n] of boundary modes to fix.
         len(target) = len(weight) = len(modes).
         If True/False uses all/none of the profile modes.
-    surface_label : float, optional
-        Surface to enforce boundary conditions on. Defaults to Equilibrium.surface.rho
     name : str, optional
         Name of the objective function.
 
-    Notes
-    -----
-    If specifying particular modes to fix, the rows of the resulting constraint `A`
-    matrix and `target` vector will be re-sorted according to the ordering of
-    `basis.modes` which may be different from the order that was passed in.
     """
 
     _units = "(m)"
@@ -618,14 +611,24 @@ def __init__(
         normalize=True,
         normalize_target=True,
         modes=True,
-        surface_label=None,
         name="lcfs R",
     ):
-        self._modes = modes
-        self._target_from_user = setdefault(bounds, target)
-        self._surface_label = surface_label
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.surface.R_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
+            )
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"Rb_lmn": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -634,86 +637,8 @@ def __init__(
             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.
-
-        """
-        eq = self.things[0]
-        if self._modes is False or self._modes is None:  # no modes
-            modes = np.array([[]], dtype=int)
-            idx = np.array([], dtype=int)
-            modes_idx = idx
-        elif self._modes is True:  # all modes
-            modes = eq.surface.R_basis.modes
-            idx = np.arange(eq.surface.R_basis.num_modes)
-            modes_idx = idx
-        else:  # specified modes
-            modes = np.atleast_2d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, idx, modes_idx = np.intersect1d(
-                eq.surface.R_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            # rearrange modes to match order of eq.surface.R_basis.modes
-            # and eq.surface.R_lmn,
-            # necessary so that the A matrix rows match up with the target b
-            modes = np.atleast_2d(eq.surface.R_basis.modes[idx, :])
-
-            if idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the surface, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = idx.size
-        # Rb_lmn -> Rb optimization space
-        self._A = np.eye(eq.surface.R_basis.num_modes)[idx, :]
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.surface.R_lmn[idx], None, modes_idx
-        )
-
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute boundary R errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            boundary R errors.
-
-        """
-        return jnp.dot(self._A, params["Rb_lmn"])
-
-
-class FixBoundaryZ(_FixedObjective):
+class FixBoundaryZ(FixParameter):
     """Boundary condition on the Z boundary parameters.
 
     Parameters
@@ -740,16 +665,9 @@ class FixBoundaryZ(_FixedObjective):
         Basis modes numbers [l,m,n] of boundary modes to fix.
         len(target) = len(weight) = len(modes).
         If True/False uses all/none of the surface modes.
-    surface_label : float, optional
-        Surface to enforce boundary conditions on. Defaults to Equilibrium.surface.rho
     name : str, optional
         Name of the objective function.
 
-    Notes
-    -----
-    If specifying particular modes to fix, the rows of the resulting constraint `A`
-    matrix and `target` vector will be re-sorted according to the ordering of
-    `basis.modes` which may be different from the order that was passed in.
     """
 
     _units = "(m)"
@@ -764,14 +682,24 @@ def __init__(
         normalize=True,
         normalize_target=True,
         modes=True,
-        surface_label=None,
         name="lcfs Z",
     ):
-        self._modes = modes
-        self._target_from_user = setdefault(bounds, target)
-        self._surface_label = surface_label
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.surface.Z_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
+            )
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"Zb_lmn": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -780,84 +708,6 @@ def __init__(
             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.
-
-        """
-        eq = self.things[0]
-        if self._modes is False or self._modes is None:  # no modes
-            modes = np.array([[]], dtype=int)
-            idx = np.array([], dtype=int)
-            modes_idx = idx
-        elif self._modes is True:  # all modes
-            modes = eq.surface.Z_basis.modes
-            idx = np.arange(eq.surface.Z_basis.num_modes)
-            modes_idx = idx
-        else:  # specified modes
-            modes = np.atleast_2d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, idx, modes_idx = np.intersect1d(
-                eq.surface.Z_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            # rearrange modes to match order of eq.surface.Z_basis.modes
-            # and eq.surface.Z_lmn,
-            # necessary so that the A matrix rows match up with the target b
-            modes = np.atleast_2d(eq.surface.Z_basis.modes[idx, :])
-
-            if idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the surface, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = idx.size
-        # Zb_lmn -> Zb optimization space
-        self._A = np.eye(eq.surface.Z_basis.num_modes)[idx, :]
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.surface.Z_lmn[idx], None, modes_idx
-        )
-
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute boundary Z errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            boundary Z errors.
-
-        """
-        return jnp.dot(self._A, params["Zb_lmn"])
-
 
 class FixLambdaGauge(_Objective):
     """Fixes gauge freedom for lambda: lambda(theta=0,zeta=0)=0.
@@ -1020,7 +870,7 @@ def compute(self, params, constants=None):
         return fixed_params
 
 
-class FixAxisR(_FixedObjective):
+class FixAxisR(FixParameter):
     """Fixes magnetic axis R coefficients.
 
     Parameters
@@ -1066,10 +916,22 @@ def __init__(
         modes=True,
         name="axis R",
     ):
-        self._modes = modes
-        self._target_from_user = setdefault(bounds, target)
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.axis.R_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
+            )
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"Ra_n": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1078,87 +940,8 @@ def __init__(
             normalize_target=normalize_target,
         )
 
-    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.
-
-        """
-        eq = self.things[0]
-
-        if self._modes is False or self._modes is None:  # no modes
-            modes = np.array([[]], dtype=int)
-            idx = np.array([], dtype=int)
-            modes_idx = idx
-        elif self._modes is True:  # all modes
-            modes = eq.axis.R_basis.modes
-            idx = np.arange(eq.axis.R_basis.num_modes)
-            modes_idx = idx
-        else:  # specified modes
-            modes = np.atleast_1d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, idx, modes_idx = np.intersect1d(
-                eq.axis.R_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            # rearrange modes to match order of eq.axis.R_basis.modes and eq.axis.R_n,
-            # necessary so that the A matrix rows match up with the target b
-            modes = np.atleast_2d(eq.axis.R_basis.modes[idx, :])
 
-            if idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the axis, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = idx.size
-        # Ra_lmn -> Ra optimization space
-        self._A = np.eye(eq.axis.R_basis.num_modes)[idx, :]
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.axis.R_n[idx], None, modes_idx
-        )
-
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute axis R errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Axis R errors.
-
-        """
-        f = jnp.dot(self._A, params["Ra_n"])
-        return f
-
-
-class FixAxisZ(_FixedObjective):
+class FixAxisZ(FixParameter):
     """Fixes magnetic axis Z coefficients.
 
     Parameters
@@ -1204,10 +987,22 @@ def __init__(
         modes=True,
         name="axis Z",
     ):
-        self._modes = modes
-        self._target_from_user = setdefault(bounds, target)
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.axis.Z_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
+            )
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"Za_n": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1216,87 +1011,8 @@ def __init__(
             normalize_target=normalize_target,
         )
 
-    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.
-
-        """
-        eq = self.things[0]
-
-        if self._modes is False or self._modes is None:  # no modes
-            modes = np.array([[]], dtype=int)
-            idx = np.array([], dtype=int)
-            modes_idx = idx
-        elif self._modes is True:  # all modes
-            modes = eq.axis.Z_basis.modes
-            idx = np.arange(eq.axis.Z_basis.num_modes)
-            modes_idx = idx
-        else:  # specified modes
-            modes = np.atleast_1d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, idx, modes_idx = np.intersect1d(
-                eq.axis.Z_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            # rearrange modes to match order of eq.axis.Z_basis.modes and eq.axis.Z_n,
-            # necessary so that the A matrix rows match up with the target b
-            modes = np.atleast_2d(eq.axis.Z_basis.modes[idx, :])
-
-            if idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the axis, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = idx.size
-        # Za_lmn -> Za optimization space
-        self._A = np.eye(eq.axis.Z_basis.num_modes)[idx, :]
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.axis.Z_n[idx], None, modes_idx
-        )
-
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute axis Z errors.
 
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Axis Z errors.
-
-        """
-        f = jnp.dot(self._A, params["Za_n"])
-        return f
-
-
-class FixModeR(_FixedObjective):
+class FixModeR(FixParameter):
     """Fixes Fourier-Zernike R coefficients.
 
     Parameters
@@ -1322,8 +1038,7 @@ class FixModeR(_FixedObjective):
     modes : ndarray, optional
         Basis modes numbers [l,m,n] of Fourier-Zernike modes to fix.
         len(target) = len(weight) = len(modes).
-        If True uses all of the Equilibrium's modes.
-        Must be either True or specified as an array
+        If True/False uses all/none of the basis modes.
     name : str, optional
         Name of the objective function.
 
@@ -1341,16 +1056,24 @@ def __init__(
         normalize=True,
         normalize_target=True,
         modes=True,
-        name="Fix Mode R",
+        name="fix mode R",
     ):
-        self._modes = modes
-        if modes is None or modes is False:
-            raise ValueError(
-                f"modes kwarg must be specified or True with FixModeR got {modes}"
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.R_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
             )
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"R_lmn": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1359,72 +1082,8 @@ def __init__(
             normalize_target=normalize_target,
         )
 
-    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.
-
-        """
-        eq = self.things[0]
-        if self._modes is True:  # all modes
-            modes = eq.R_basis.modes
-            self._idx = np.arange(eq.R_basis.num_modes)
-            modes_idx = self._idx
-        else:  # specified modes
-            modes = np.atleast_2d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, self._idx, modes_idx = np.intersect1d(
-                eq.R_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            if self._idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the basis, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = modes_idx.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.R_lmn[self._idx], None, modes_idx
-        )
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute Fixed mode R errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed mode R errors.
-
-        """
-        fixed_params = params["R_lmn"][self._idx]
-        return fixed_params
-
-
-class FixModeZ(_FixedObjective):
+class FixModeZ(FixParameter):
     """Fixes Fourier-Zernike Z coefficients.
 
     Parameters
@@ -1450,8 +1109,7 @@ class FixModeZ(_FixedObjective):
     modes : ndarray, optional
         Basis modes numbers [l,m,n] of Fourier-Zernike modes to fix.
         len(target) = len(weight) = len(modes).
-        If True uses all of the Equilibrium's modes.
-        Must be either True or specified as an array
+        If True/False uses all/none of the basis modes.
     name : str, optional
         Name of the objective function.
 
@@ -1469,16 +1127,24 @@ def __init__(
         normalize=True,
         normalize_target=True,
         modes=True,
-        name="Fix Mode Z",
+        name="fix mode Z",
     ):
-        self._modes = modes
-        if modes is None or modes is False:
-            raise ValueError(
-                f"modes kwarg must be specified or True with FixModeZ got {modes}"
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.Z_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
             )
-        self._target_from_user = setdefault(bounds, target)
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"Z_lmn": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1487,72 +1153,8 @@ def __init__(
             normalize_target=normalize_target,
         )
 
-    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.
-
-        """
-        eq = self.things[0]
-        if self._modes is True:  # all modes
-            modes = eq.Z_basis.modes
-            self._idx = np.arange(eq.Z_basis.num_modes)
-            modes_idx = self._idx
-        else:  # specified modes
-            modes = np.atleast_2d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, self._idx, modes_idx = np.intersect1d(
-                eq.Z_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            if self._idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the basis, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = modes_idx.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.Z_lmn[self._idx], None, modes_idx
-        )
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute Fixed mode Z errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed mode Z errors.
-
-        """
-        fixed_params = params["Z_lmn"][self._idx]
-        return fixed_params
-
 
-class FixModeLambda(_FixedObjective):
+class FixModeLambda(FixParameter):
     """Fixes Fourier-Zernike lambda coefficients.
 
     Parameters
@@ -1579,8 +1181,7 @@ class FixModeLambda(_FixedObjective):
     modes : ndarray, optional
         Basis modes numbers [l,m,n] of Fourier-Zernike modes to fix.
         len(target) = len(weight) = len(modes).
-        If True uses all of the Equilibrium's modes.
-        Must be either True or specified as an array
+        If True/False uses all/none of the basis modes.
     name : str
         Name of the objective function.
 
@@ -1598,17 +1199,24 @@ def __init__(
         normalize=True,
         normalize_target=True,
         modes=True,
-        name="Fix Mode lambda",
+        name="fix mode lambda",
     ):
-        self._modes = modes
-        if modes is None or modes is False:
-            raise ValueError(
-                "modes kwarg must be specified"
-                + f" or True with FixModeLambda got {modes}"
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        if isinstance(modes, bool):
+            idx = modes
+        else:
+            idx = np.concatenate(
+                [
+                    [eq.L_basis.get_idx(L=l, M=m, N=n)]
+                    for l, m, n in np.atleast_2d(modes)
+                ],
+                dtype=int,
             )
-        self._target_from_user = target
         super().__init__(
-            things=eq,
+            thing=eq,
+            params={"L_lmn": idx},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1617,70 +1225,6 @@ def __init__(
             normalize_target=normalize_target,
         )
 
-    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.
-
-        """
-        eq = self.things[0]
-        if self._modes is True:  # all modes
-            modes = eq.L_basis.modes
-            self._idx = np.arange(eq.L_basis.num_modes)
-            modes_idx = self._idx
-        else:  # specified modes
-            modes = np.atleast_2d(self._modes)
-            dtype = {
-                "names": ["f{}".format(i) for i in range(3)],
-                "formats": 3 * [modes.dtype],
-            }
-            _, self._idx, modes_idx = np.intersect1d(
-                eq.L_basis.modes.astype(modes.dtype).view(dtype),
-                modes.view(dtype),
-                return_indices=True,
-            )
-            if self._idx.size < modes.shape[0]:
-                warnings.warn(
-                    colored(
-                        "Some of the given modes are not in the basis, "
-                        + "these modes will not be fixed.",
-                        "yellow",
-                    )
-                )
-
-        self._dim_f = modes_idx.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, eq.L_lmn[self._idx], None, modes_idx
-        )
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute Fixed mode lambda errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed mode lambda errors.
-
-        """
-        fixed_params = params["L_lmn"][self._idx]
-        return fixed_params
-
 
 class FixSumModesR(_FixedObjective):
     """Fixes a linear sum of Fourier-Zernike R coefficients.
@@ -3164,6 +2708,9 @@ def __init__(
         normalize_target=True,
         name="fixed-Psi",
     ):
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["Psi"]
         super().__init__(
             thing=eq,
             params={"Psi": True},
@@ -3174,9 +2721,6 @@ def __init__(
             normalize_target=normalize_target,
             name=name,
         )
-        if self._normalize:
-            scales = compute_scaling_factors(self.things[0])
-            self._normalization = scales["Psi"]
 
 
 class FixCurveShift(FixParameter):

From 9dc6a17d94e5030ac1dc37eca842005b4b5b44a4 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 18:42:17 -0500
Subject: [PATCH 30/50] replace FixProfile with FixParameter

---
 desc/objectives/linear_objectives.py | 645 +++++----------------------
 1 file changed, 122 insertions(+), 523 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 5b20f48e39..c509cc9f22 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -7,7 +7,6 @@
 """
 
 import warnings
-from abc import ABC
 
 import numpy as np
 from termcolor import colored
@@ -617,9 +616,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.surface.R_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -628,7 +627,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"Rb_lmn": idx},
+            params={"Rb_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -688,9 +687,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.surface.Z_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -699,7 +698,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"Zb_lmn": idx},
+            params={"Zb_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -920,9 +919,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.axis.R_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -931,7 +930,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"Ra_n": idx},
+            params={"Ra_n": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -991,9 +990,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.axis.Z_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -1002,7 +1001,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"Za_n": idx},
+            params={"Za_n": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1062,9 +1061,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.R_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -1073,7 +1072,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"R_lmn": idx},
+            params={"R_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1133,9 +1132,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.Z_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -1144,7 +1143,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"Z_lmn": idx},
+            params={"Z_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1205,9 +1204,9 @@ def __init__(
             scales = compute_scaling_factors(eq)
             self._normalization = scales["a"]
         if isinstance(modes, bool):
-            idx = modes
+            indices = modes
         else:
-            idx = np.concatenate(
+            indices = np.concatenate(
                 [
                     [eq.L_basis.get_idx(L=l, M=m, N=n)]
                     for l, m, n in np.atleast_2d(modes)
@@ -1216,7 +1215,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"L_lmn": idx},
+            params={"L_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1728,105 +1727,7 @@ def compute(self, params, constants=None):
         return f
 
 
-class _FixProfile(_FixedObjective, ABC):
-    """Fixes profile coefficients (or values, for SplineProfile).
-
-    Parameters
-    ----------
-    eq : Equilibrium
-        Equilibrium that will be optimized to satisfy the Objective.
-    target : {float, ndarray}, optional
-        Target value(s) of the objective. Only used if bounds is None.
-        Must be broadcastable to Objective.dim_f.
-    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
-    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.
-    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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
-    indices : ndarray or Bool, optional
-        indices of the Profile.params array to fix.
-        (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
-        corresponding to knots for a SplineProfile).
-        Must have len(target) = len(weight) = len(indices).
-        If True/False uses all/none of the Profile.params indices.
-    name : str, optional
-        Name of the objective function.
-
-    """
-
-    _print_value_fmt = "Fix-profile error: {:10.3e} "
-
-    def __init__(
-        self,
-        eq,
-        target=None,
-        bounds=None,
-        weight=1,
-        normalize=True,
-        normalize_target=True,
-        profile=None,
-        indices=True,
-        name="",
-    ):
-        self._profile = profile
-        self._indices = indices
-        self._target_from_user = setdefault(bounds, target)
-        super().__init__(
-            things=eq,
-            target=target,
-            bounds=bounds,
-            weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
-            name=name,
-        )
-
-    def build(self, eq, profile, use_jit=False, verbose=1):
-        """Build constant arrays.
-
-        Parameters
-        ----------
-        eq : Equilibrium
-            Equilibrium that will be optimized to satisfy the Objective.
-        profile : Profile, optional
-            profile to fix
-        use_jit : bool, optional
-            Whether to just-in-time compile the objective and derivatives.
-        verbose : int, optional
-            Level of output.
-
-        """
-        eq = self.things[0]
-        if self._profile is None or self._profile.params.size != eq.L + 1:
-            self._profile = profile
-
-        # find indices to fix
-        if self._indices is False or self._indices is None:  # no indices to fix
-            self._idx = np.array([], dtype=int)
-        elif self._indices is True:  # all indices of Profile.params
-            self._idx = np.arange(np.size(self._profile.params))
-        else:  # specified indices
-            self._idx = np.atleast_1d(self._indices)
-
-        self._dim_f = self._idx.size
-
-        self.target, self.bounds = self._parse_target_from_user(
-            self._target_from_user, self._profile.params[self._idx], None, self._idx
-        )
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-
-class FixPressure(_FixProfile):
+class FixPressure(FixParameter):
     """Fixes pressure coefficients.
 
     Parameters
@@ -1849,8 +1750,6 @@ class FixPressure(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -1863,7 +1762,7 @@ class FixPressure(_FixProfile):
     """
 
     _units = "(Pa)"
-    _print_value_fmt = "Fixed-pressure profile error: {:10.3e} "
+    _print_value_fmt = "Fixed pressure profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -1873,66 +1772,30 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        profile=None,
         indices=True,
-        name="fixed-pressure",
+        name="fixed pressure",
     ):
+        if eq.pressure is None:
+            raise RuntimeError(
+                "Attempting to fix pressure on an Equilibrium with no "
+                + "pressure profile assigned."
+            )
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["p"]
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"p_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.pressure is None:
-            raise RuntimeError(
-                "Attempting to fix pressure on an equilibrium with no "
-                + "pressure profile assigned"
-            )
-        profile = eq.pressure
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["p"]
-        super().build(eq, profile, use_jit, verbose)
 
-    def compute(self, params, constants=None):
-        """Compute fixed pressure profile errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["p_l"][self._idx]
-
-
-class FixAnisotropy(_FixProfile):
+class FixAnisotropy(FixParameter):
     """Fixes anisotropic pressure coefficients.
 
     Parameters
@@ -1955,8 +1818,6 @@ class FixAnisotropy(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -1969,7 +1830,7 @@ class FixAnisotropy(_FixProfile):
     """
 
     _units = "(dimensionless)"
-    _print_value_fmt = "Fixed-anisotropy profile error: {:10.3e} "
+    _print_value_fmt = "Fixed anisotropy profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -1979,63 +1840,27 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        profile=None,
         indices=True,
-        name="fixed-anisotropy",
+        name="fixed anisotropy",
     ):
+        if eq.anisotropy is None:
+            raise RuntimeError(
+                "Attempting to fix anisotropy on an Equilibrium with no "
+                + "anisotropy profile assigned."
+            )
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"a_lmn", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.anisotropy is None:
-            raise RuntimeError(
-                "Attempting to fix anisotropy on an equilibrium with no "
-                + "anisotropy profile assigned"
-            )
-        profile = eq.anisotropy
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed pressure profile errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["a_lmn"][self._idx]
 
-
-class FixIota(_FixProfile):
+class FixIota(FixParameter):
     """Fixes rotational transform coefficients.
 
     Parameters
@@ -2059,8 +1884,6 @@ class FixIota(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices.
@@ -2073,7 +1896,7 @@ class FixIota(_FixProfile):
     """
 
     _units = "(dimensionless)"
-    _print_value_fmt = "Fixed-iota profile error: {:10.3e} "
+    _print_value_fmt = "Fixed iota profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -2083,63 +1906,27 @@ def __init__(
         weight=1,
         normalize=False,
         normalize_target=False,
-        profile=None,
         indices=True,
-        name="fixed-iota",
+        name="fixed iota",
     ):
+        if eq.iota is None:
+            raise RuntimeError(
+                "Attempting to fix iota on an Equilibrium with no "
+                + "iota profile assigned."
+            )
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"i_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.iota is None:
-            raise RuntimeError(
-                "Attempt to fix rotational transform on an equilibrium with no "
-                + "rotational transform profile assigned"
-            )
-        profile = eq.iota
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed iota errors.
 
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["i_l"][self._idx]
-
-
-class FixCurrent(_FixProfile):
+class FixCurrent(FixParameter):
     """Fixes toroidal current profile coefficients.
 
     Parameters
@@ -2162,8 +1949,6 @@ class FixCurrent(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -2176,7 +1961,7 @@ class FixCurrent(_FixProfile):
     """
 
     _units = "(A)"
-    _print_value_fmt = "Fixed-current profile error: {:10.3e} "
+    _print_value_fmt = "Fixed current profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -2186,66 +1971,30 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        profile=None,
         indices=True,
-        name="fixed-current",
+        name="fixed current",
     ):
+        if eq.current is None:
+            raise RuntimeError(
+                "Attempting to fix current on an Equilibrium with no "
+                + "current profile assigned."
+            )
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["I"]
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"c_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.current is None:
-            raise RuntimeError(
-                "Attempting to fix toroidal current on an equilibrium with no "
-                + "current profile assigned"
-            )
-        profile = eq.current
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["I"]
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed current errors.
 
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["c_l"][self._idx]
-
-
-class FixElectronTemperature(_FixProfile):
+class FixElectronTemperature(FixParameter):
     """Fixes electron temperature profile coefficients.
 
     Parameters
@@ -2268,8 +2017,6 @@ class FixElectronTemperature(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -2282,7 +2029,7 @@ class FixElectronTemperature(_FixProfile):
     """
 
     _units = "(eV)"
-    _print_value_fmt = "Fixed-electron-temperature profile error: {:10.3e} "
+    _print_value_fmt = "Fixed electron temperature profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -2292,66 +2039,30 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        profile=None,
         indices=True,
-        name="fixed-electron-temperature",
+        name="fixed electron temperature",
     ):
+        if eq.electron_temperature is None:
+            raise RuntimeError(
+                "Attempting to fix electron temperature on an Equilibrium with no "
+                + "electron temperature profile assigned."
+            )
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["T"]
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"Te_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.electron_temperature is None:
-            raise RuntimeError(
-                "Attempting to fix electron temperature on an equilibrium with no "
-                + "electron temperature profile assigned"
-            )
-        profile = eq.electron_temperature
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["T"]
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed electron temperature errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["Te_l"][self._idx]
-
-
-class FixElectronDensity(_FixProfile):
+class FixElectronDensity(FixParameter):
     """Fixes electron density profile coefficients.
 
     Parameters
@@ -2388,7 +2099,7 @@ class FixElectronDensity(_FixProfile):
     """
 
     _units = "(m^-3)"
-    _print_value_fmt = "Fixed-electron-density profile error: {:10.3e} "
+    _print_value_fmt = "Fixed electron density profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -2398,66 +2109,30 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        profile=None,
         indices=True,
-        name="fixed-electron-density",
+        name="fixed electron density",
     ):
+        if eq.electron_density is None:
+            raise RuntimeError(
+                "Attempting to fix electron density on an Equilibrium with no "
+                + "electron density profile assigned."
+            )
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["n"]
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"ne_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.electron_density is None:
-            raise RuntimeError(
-                "Attempting to fix electron density on an equilibrium with no "
-                + "electron density profile assigned"
-            )
-        profile = eq.electron_density
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["n"]
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed electron density errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["ne_l"][self._idx]
-
-
-class FixIonTemperature(_FixProfile):
+class FixIonTemperature(FixParameter):
     """Fixes ion temperature profile coefficients.
 
     Parameters
@@ -2480,8 +2155,6 @@ class FixIonTemperature(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -2494,7 +2167,7 @@ class FixIonTemperature(_FixProfile):
     """
 
     _units = "(eV)"
-    _print_value_fmt = "Fixed-ion-temperature profile error: {:10.3e} "
+    _print_value_fmt = "Fixed ion temperature profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -2504,66 +2177,30 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        profile=None,
         indices=True,
-        name="fixed-ion-temperature",
+        name="fixed ion temperature",
     ):
+        if eq.ion_temperature is None:
+            raise RuntimeError(
+                "Attempting to fix ion temperature on an Equilibrium with no "
+                + "ion temperature profile assigned."
+            )
+        if normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["T"]
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"Ti_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.ion_temperature is None:
-            raise RuntimeError(
-                "Attempting to fix ion temperature on an equilibrium with no "
-                + "ion temperature profile assigned"
-            )
-        profile = eq.ion_temperature
-        if self._normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["T"]
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed ion temperature errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["Ti_l"][self._idx]
-
 
-class FixAtomicNumber(_FixProfile):
+class FixAtomicNumber(FixParameter):
     """Fixes effective atomic number profile coefficients.
 
     Parameters
@@ -2587,8 +2224,6 @@ class FixAtomicNumber(_FixProfile):
         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.
-    profile : Profile, optional
-        Profile containing the radial modes to evaluate at.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -2601,7 +2236,7 @@ class FixAtomicNumber(_FixProfile):
     """
 
     _units = "(dimensionless)"
-    _print_value_fmt = "Fixed-atomic-number profile error: {:10.3e} "
+    _print_value_fmt = "Fixed atomic number profile error: {:10.3e} "
 
     def __init__(
         self,
@@ -2611,61 +2246,25 @@ def __init__(
         weight=1,
         normalize=False,
         normalize_target=False,
-        profile=None,
         indices=True,
-        name="fixed-atomic-number",
+        name="fixed atomic number",
     ):
+        if eq.atomic_number is None:
+            raise RuntimeError(
+                "Attempting to fix atomic number on an Equilibrium with no "
+                + "atomic_number profile assigned."
+            )
         super().__init__(
-            eq=eq,
+            thing=eq,
+            params={"Zeff_l", indices},
             target=target,
             bounds=bounds,
             weight=weight,
             normalize=normalize,
             normalize_target=normalize_target,
-            profile=profile,
-            indices=indices,
             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.
-
-        """
-        eq = self.things[0]
-        if eq.atomic_number is None:
-            raise RuntimeError(
-                "Attempting to fix atomic number on an equilibrium with no "
-                + "atomic number profile assigned"
-            )
-        profile = eq.atomic_number
-        super().build(eq, profile, use_jit, verbose)
-
-    def compute(self, params, constants=None):
-        """Compute fixed atomic number errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
-
-        Returns
-        -------
-        f : ndarray
-            Fixed profile errors.
-
-        """
-        return params["Zeff_l"][self._idx]
-
 
 class FixPsi(FixParameter):
     """Fixes total toroidal magnetic flux within the last closed flux surface.
@@ -2696,7 +2295,7 @@ class FixPsi(FixParameter):
     """
 
     _units = "(Wb)"
-    _print_value_fmt = "Fixed-Psi error: {:10.3e} "
+    _print_value_fmt = "Fixed Psi error: {:10.3e} "
 
     def __init__(
         self,
@@ -2706,7 +2305,7 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        name="fixed-Psi",
+        name="fixed Psi",
     ):
         if normalize:
             scales = compute_scaling_factors(eq)
@@ -2751,7 +2350,7 @@ class FixCurveShift(FixParameter):
     """
 
     _units = "(m)"
-    _print_value_fmt = "Fixed-shift error: {:10.3e} "
+    _print_value_fmt = "Fixed shift error: {:10.3e} "
 
     def __init__(
         self,
@@ -2761,7 +2360,7 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        name="fixed-shift",
+        name="fixed shift",
     ):
         super().__init__(
             thing=curve,
@@ -2803,7 +2402,7 @@ class FixCurveRotation(FixParameter):
     """
 
     _units = "(rad)"
-    _print_value_fmt = "Fixed-rotation error: {:10.3e} "
+    _print_value_fmt = "Fixed rotation error: {:10.3e} "
 
     def __init__(
         self,
@@ -2813,7 +2412,7 @@ def __init__(
         weight=1,
         normalize=True,
         normalize_target=True,
-        name="fixed-rotation",
+        name="fixed rotation",
     ):
         super().__init__(
             thing=curve,

From eb6dc48ffde7170816450b943526609b9d963b5d Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 18:52:23 -0500
Subject: [PATCH 31/50] FixTheatSFL inherit from FixParameter

---
 desc/objectives/linear_objectives.py | 97 ++++++++++++----------------
 1 file changed, 40 insertions(+), 57 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index c509cc9f22..b417f16cbb 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -215,7 +215,7 @@ class BoundaryRSelfConsistency(_Objective):
 
     _scalar = False
     _linear = True
-    _fixed = False
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(m)"
     _print_value_fmt = "R boundary self consistency error: {:10.3e} "
 
@@ -313,7 +313,7 @@ class BoundaryZSelfConsistency(_Objective):
 
     _scalar = False
     _linear = True
-    _fixed = False
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(m)"
     _print_value_fmt = "Z boundary self consistency error: {:10.3e} "
 
@@ -411,7 +411,7 @@ class AxisRSelfConsistency(_Objective):
 
     _scalar = False
     _linear = True
-    _fixed = False
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _print_value_fmt = "R axis self consistency error: {:10.3e} (m)"
 
     def __init__(
@@ -497,7 +497,7 @@ class AxisZSelfConsistency(_Objective):
 
     _scalar = False
     _linear = True
-    _fixed = False
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _print_value_fmt = "Z axis self consistency error: {:10.3e} (m)"
 
     def __init__(
@@ -725,7 +725,7 @@ class FixLambdaGauge(_Objective):
 
     _scalar = False
     _linear = True
-    _fixed = False
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(rad)"
     _print_value_fmt = "lambda gauge error: {:10.3e} "
 
@@ -805,68 +805,48 @@ def compute(self, params, constants=None):
         return jnp.dot(self._A, params["L_lmn"])
 
 
-class FixThetaSFL(_Objective):
+class FixThetaSFL(FixParameter):
     """Fixes lambda=0 so that poloidal angle is the SFL poloidal angle.
 
     Parameters
     ----------
     eq : Equilibrium
         Equilibrium that will be optimized to satisfy the Objective.
+    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.
+    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.
     name : str, optional
         Name of the objective function.
 
     """
 
-    _scalar = False
-    _linear = True
-    _fixed = True
     _units = "(rad)"
-    _print_value_fmt = "Theta - Theta SFL error: {:10.3e} "
-
-    def __init__(self, eq, name="Theta SFL"):
-        super().__init__(things=eq, target=0, weight=1, 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.
-
-        """
-        eq = self.things[0]
-        idx = np.arange(eq.L_basis.num_modes)
-        modes_idx = idx
-        self._idx = idx
-
-        self._dim_f = modes_idx.size
-
-        self.target = np.zeros_like(modes_idx)
-
-        super().build(use_jit=use_jit, verbose=verbose)
-
-    def compute(self, params, constants=None):
-        """Compute Theta SFL errors.
-
-        Parameters
-        ----------
-        params : dict
-            Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict
-        constants : dict
-            Dictionary of constant data, eg transforms, profiles etc. Defaults to
-            self.constants
+    _print_value_fmt = "theta - theta SFL error: {:10.3e} "
 
-        Returns
-        -------
-        f : ndarray
-            Theta - Theta SFL errors.
-
-        """
-        fixed_params = params["L_lmn"][self._idx]
-        return fixed_params
+    def __init__(
+        self,
+        eq,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        name="theta SFL",
+    ):
+        super().__init__(
+            thing=eq,
+            params={"L_lmn": True},
+            target=0,
+            bounds=None,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            name=name,
+        )
 
 
 class FixAxisR(FixParameter):
@@ -1261,7 +1241,7 @@ class FixSumModesR(_FixedObjective):
 
     """
 
-    _fixed = False  # not "diagonal", since its fixing a sum
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(m)"
     _print_value_fmt = "Fixed-R sum modes error: {:10.3e} "
 
@@ -1427,7 +1407,7 @@ class FixSumModesZ(_FixedObjective):
 
     """
 
-    _fixed = False  # not "diagonal", since its fixing a sum
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(m)"
     _print_value_fmt = "Fixed-Z sum modes error: {:10.3e} "
 
@@ -1595,7 +1575,7 @@ class FixSumModesLambda(_FixedObjective):
 
     """
 
-    _fixed = False  # not "diagonal", since its fixing a sum
+    _fixed = False  # not "diagonal", since it is fixing a sum
     _units = "(rad)"
     _print_value_fmt = "Fixed-lambda sum modes error: {:10.3e} "
 
@@ -2643,6 +2623,9 @@ def compute(self, params, constants=None):
 class FixSheetCurrent(FixParameter):
     """Fixes the sheet current parameters of a free-boundary equilibrium.
 
+    Note: this constraint is automatically applied when needed, and does not need to be
+    included by the user.
+
     Parameters
     ----------
     eq : Equilibrium

From e9d95a8c2a77a185dd0eb2987670992c2a9e6fa8 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 22:32:39 -0500
Subject: [PATCH 32/50] syntax error

---
 desc/objectives/linear_objectives.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index b417f16cbb..dd55b8c44b 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -1765,7 +1765,7 @@ def __init__(
             self._normalization = scales["p"]
         super().__init__(
             thing=eq,
-            params={"p_l", indices},
+            params={"p_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1830,7 +1830,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"a_lmn", indices},
+            params={"a_lmn": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1896,7 +1896,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"i_l", indices},
+            params={"i_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -1964,7 +1964,7 @@ def __init__(
             self._normalization = scales["I"]
         super().__init__(
             thing=eq,
-            params={"c_l", indices},
+            params={"c_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -2032,7 +2032,7 @@ def __init__(
             self._normalization = scales["T"]
         super().__init__(
             thing=eq,
-            params={"Te_l", indices},
+            params={"Te_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -2102,7 +2102,7 @@ def __init__(
             self._normalization = scales["n"]
         super().__init__(
             thing=eq,
-            params={"ne_l", indices},
+            params={"ne_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -2170,7 +2170,7 @@ def __init__(
             self._normalization = scales["T"]
         super().__init__(
             thing=eq,
-            params={"Ti_l", indices},
+            params={"Ti_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,
@@ -2236,7 +2236,7 @@ def __init__(
             )
         super().__init__(
             thing=eq,
-            params={"Zeff_l", indices},
+            params={"Zeff_l": indices},
             target=target,
             bounds=bounds,
             weight=weight,

From 900012c4e94bfd54a3e43938ad9ecbdc58498f90 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 23:35:13 -0500
Subject: [PATCH 33/50] bug fixes

---
 desc/objectives/linear_objectives.py | 70 +++++++++-------------------
 desc/utils.py                        | 14 +++---
 tests/test_linear_objectives.py      | 12 ++---
 3 files changed, 35 insertions(+), 61 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index dd55b8c44b..26a0be8941 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -618,13 +618,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.surface.R_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"Rb_lmn": indices},
@@ -689,13 +685,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.surface.Z_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"Zb_lmn": indices},
@@ -901,13 +893,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.axis.R_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"Ra_n": indices},
@@ -972,13 +960,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.axis.Z_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"Za_n": indices},
@@ -1043,13 +1027,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.R_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"R_lmn": indices},
@@ -1114,13 +1094,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.Z_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"Z_lmn": indices},
@@ -1186,13 +1162,9 @@ def __init__(
         if isinstance(modes, bool):
             indices = modes
         else:
-            indices = np.concatenate(
-                [
-                    [eq.L_basis.get_idx(L=l, M=m, N=n)]
-                    for l, m, n in np.atleast_2d(modes)
-                ],
-                dtype=int,
-            )
+            indices = np.array([], dtype=int)
+            for l, m, n in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
         super().__init__(
             thing=eq,
             params={"L_lmn": indices},
diff --git a/desc/utils.py b/desc/utils.py
index 0c638f3fd2..38d93c3c45 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -634,18 +634,20 @@ def broadcast_tree(tree_in, tree_out, dtype=int):
     """
     # both trees at leaf layer
     if isinstance(tree_in, dict) and isinstance(tree_out, dict):
-        tree_new = tree_in.copy()
+        tree_new = {}
         for key, value in tree_in.items():
             errorif(
                 key not in tree_out.keys(),
                 ValueError,
                 "dict keys of tree_in must be a subset of those in tree_out",
             )
-            if isinstance(value, bool) and value:
-                tree_new[key] = np.atleast_1d(tree_out[key]).astype(dtype=dtype)
-            if isinstance(value, bool) and not value:
-                tree_new[key] = np.array([], dtype=dtype)
-            value = np.atleast_1d(value).astype(dtype=dtype)
+            if isinstance(value, bool):
+                if value:
+                    tree_new[key] = np.atleast_1d(tree_out[key]).astype(dtype=dtype)
+                else:
+                    tree_new[key] = np.array([], dtype=dtype)
+            else:
+                tree_new[key] = np.atleast_1d(value).astype(dtype=dtype)
         for key, value in tree_out.items():
             if key not in tree_new.keys():
                 tree_new[key] = np.array([], dtype=dtype)
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index fa523a2b59..c33b229367 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -904,27 +904,27 @@ def test_fix_omni_indices():
     # no indices
     constraint = FixOmniWell(field=field, indices=False)
     constraint.build()
-    assert constraint._idx.size == 0
+    assert constraint.dim_f == 0
     constraint = FixOmniMap(field=field, indices=False)
     constraint.build()
-    assert constraint._idx.size == 0
+    assert constraint.dim_f == 0
 
     # all indices
     constraint = FixOmniWell(field=field, indices=True)
     constraint.build()
-    assert constraint._idx.size == field.B_lm.size
+    assert constraint.dim_f == field.B_lm.size
     constraint = FixOmniMap(field=field, indices=True)
     constraint.build()
-    assert constraint._idx.size == field.x_lmn.size
+    assert constraint.dim_f == field.x_lmn.size
 
     # specified indices
     indices = np.arange(3, 8)
     constraint = FixOmniWell(field=field, indices=indices)
     constraint.build()
-    assert constraint._idx.size == indices.size
+    assert constraint.dim_f == indices.size
     constraint = FixOmniMap(field=field, indices=indices)
     constraint.build()
-    assert constraint._idx.size == indices.size
+    assert constraint.dim_f == indices.size
 
 
 @pytest.mark.unit

From 5fb1614fe6455960ff34076e2145b8ffeda46ec6 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Mon, 29 Apr 2024 23:41:58 -0500
Subject: [PATCH 34/50] remove outdated FIXME comment

---
 desc/objectives/utils.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py
index 109246587b..949b3cc09f 100644
--- a/desc/objectives/utils.py
+++ b/desc/objectives/utils.py
@@ -145,7 +145,6 @@ def recover(x_reduced):
     # check that all constraints are actually satisfiable
     params = objective.unpack_state(xp, False)
     for con in constraint.objectives:
-        # FIXME: xpi is not always correct -- missing params?
         xpi = [params[i] for i, t in enumerate(objective.things) if t in con.things]
         y1 = con.compute_unscaled(*xpi)
         y2 = con.target

From 7a314b94ed269c1e552ab965639b962426a2e779 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 30 Apr 2024 11:03:34 -0500
Subject: [PATCH 35/50] repairing tests

---
 desc/objectives/linear_objectives.py | 381 +++++++++++++++++++++------
 desc/objectives/objective_funs.py    |  15 +-
 tests/test_linear_objectives.py      |  12 +-
 tests/test_utils.py                  |   4 +-
 4 files changed, 315 insertions(+), 97 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 26a0be8941..21435cf57b 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -612,15 +612,12 @@ def __init__(
         modes=True,
         name="lcfs R",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"Rb_lmn": indices},
@@ -632,6 +629,23 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixBoundaryZ(FixParameter):
     """Boundary condition on the Z boundary parameters.
@@ -679,15 +693,12 @@ def __init__(
         modes=True,
         name="lcfs Z",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"Zb_lmn": indices},
@@ -699,6 +710,23 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixLambdaGauge(_Objective):
     """Fixes gauge freedom for lambda: lambda(theta=0,zeta=0)=0.
@@ -887,15 +915,12 @@ def __init__(
         modes=True,
         name="axis R",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"Ra_n": indices},
@@ -907,6 +932,23 @@ def __init__(
             normalize_target=normalize_target,
         )
 
+    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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixAxisZ(FixParameter):
     """Fixes magnetic axis Z coefficients.
@@ -954,15 +996,12 @@ def __init__(
         modes=True,
         name="axis Z",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"Za_n": indices},
@@ -974,6 +1013,23 @@ def __init__(
             normalize_target=normalize_target,
         )
 
+    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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixModeR(FixParameter):
     """Fixes Fourier-Zernike R coefficients.
@@ -1021,15 +1077,12 @@ def __init__(
         modes=True,
         name="fix mode R",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"R_lmn": indices},
@@ -1041,6 +1094,23 @@ def __init__(
             normalize_target=normalize_target,
         )
 
+    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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixModeZ(FixParameter):
     """Fixes Fourier-Zernike Z coefficients.
@@ -1088,15 +1158,12 @@ def __init__(
         modes=True,
         name="fix mode Z",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"Z_lmn": indices},
@@ -1108,6 +1175,23 @@ def __init__(
             normalize_target=normalize_target,
         )
 
+    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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["a"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixModeLambda(FixParameter):
     """Fixes Fourier-Zernike lambda coefficients.
@@ -1143,7 +1227,7 @@ class FixModeLambda(FixParameter):
     """
 
     _units = "(rad)"
-    _print_value_fmt = "Fixed-lambda modes error: {:10.3e} "
+    _print_value_fmt = "Fixed lambda modes error: {:10.3e} "
 
     def __init__(
         self,
@@ -1156,15 +1240,12 @@ def __init__(
         modes=True,
         name="fix mode lambda",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["a"]
         if isinstance(modes, bool):
             indices = modes
         else:
             indices = np.array([], dtype=int)
             for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(L=l, M=m, N=n))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
         super().__init__(
             thing=eq,
             params={"L_lmn": indices},
@@ -1727,14 +1808,6 @@ def __init__(
         indices=True,
         name="fixed pressure",
     ):
-        if eq.pressure is None:
-            raise RuntimeError(
-                "Attempting to fix pressure on an Equilibrium with no "
-                + "pressure profile assigned."
-            )
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["p"]
         super().__init__(
             thing=eq,
             params={"p_l": indices},
@@ -1746,6 +1819,28 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.pressure is None:
+            raise RuntimeError(
+                "Attempting to fix pressure on an Equilibrium with no "
+                + "pressure profile assigned."
+            )
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["p"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixAnisotropy(FixParameter):
     """Fixes anisotropic pressure coefficients.
@@ -1795,11 +1890,6 @@ def __init__(
         indices=True,
         name="fixed anisotropy",
     ):
-        if eq.anisotropy is None:
-            raise RuntimeError(
-                "Attempting to fix anisotropy on an Equilibrium with no "
-                + "anisotropy profile assigned."
-            )
         super().__init__(
             thing=eq,
             params={"a_lmn": indices},
@@ -1811,6 +1901,25 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.anisotropy is None:
+            raise RuntimeError(
+                "Attempting to fix anisotropy on an Equilibrium with no "
+                + "anisotropy profile assigned."
+            )
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixIota(FixParameter):
     """Fixes rotational transform coefficients.
@@ -1861,11 +1970,6 @@ def __init__(
         indices=True,
         name="fixed iota",
     ):
-        if eq.iota is None:
-            raise RuntimeError(
-                "Attempting to fix iota on an Equilibrium with no "
-                + "iota profile assigned."
-            )
         super().__init__(
             thing=eq,
             params={"i_l": indices},
@@ -1877,6 +1981,25 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.iota is None:
+            raise RuntimeError(
+                "Attempting to fix iota on an Equilibrium with no "
+                + "iota profile assigned."
+            )
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixCurrent(FixParameter):
     """Fixes toroidal current profile coefficients.
@@ -1926,14 +2049,6 @@ def __init__(
         indices=True,
         name="fixed current",
     ):
-        if eq.current is None:
-            raise RuntimeError(
-                "Attempting to fix current on an Equilibrium with no "
-                + "current profile assigned."
-            )
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["I"]
         super().__init__(
             thing=eq,
             params={"c_l": indices},
@@ -1945,6 +2060,28 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.current is None:
+            raise RuntimeError(
+                "Attempting to fix current on an Equilibrium with no "
+                + "current profile assigned."
+            )
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["I"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixElectronTemperature(FixParameter):
     """Fixes electron temperature profile coefficients.
@@ -1994,14 +2131,6 @@ def __init__(
         indices=True,
         name="fixed electron temperature",
     ):
-        if eq.electron_temperature is None:
-            raise RuntimeError(
-                "Attempting to fix electron temperature on an Equilibrium with no "
-                + "electron temperature profile assigned."
-            )
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["T"]
         super().__init__(
             thing=eq,
             params={"Te_l": indices},
@@ -2013,6 +2142,28 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.electron_temperature is None:
+            raise RuntimeError(
+                "Attempting to fix electron temperature on an Equilibrium with no "
+                + "electron temperature profile assigned."
+            )
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["T"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixElectronDensity(FixParameter):
     """Fixes electron density profile coefficients.
@@ -2064,14 +2215,6 @@ def __init__(
         indices=True,
         name="fixed electron density",
     ):
-        if eq.electron_density is None:
-            raise RuntimeError(
-                "Attempting to fix electron density on an Equilibrium with no "
-                + "electron density profile assigned."
-            )
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["n"]
         super().__init__(
             thing=eq,
             params={"ne_l": indices},
@@ -2083,6 +2226,28 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.electron_density is None:
+            raise RuntimeError(
+                "Attempting to fix electron density on an Equilibrium with no "
+                + "electron density profile assigned."
+            )
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["n"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixIonTemperature(FixParameter):
     """Fixes ion temperature profile coefficients.
@@ -2132,14 +2297,6 @@ def __init__(
         indices=True,
         name="fixed ion temperature",
     ):
-        if eq.ion_temperature is None:
-            raise RuntimeError(
-                "Attempting to fix ion temperature on an Equilibrium with no "
-                + "ion temperature profile assigned."
-            )
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["T"]
         super().__init__(
             thing=eq,
             params={"Ti_l": indices},
@@ -2151,6 +2308,28 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.ion_temperature is None:
+            raise RuntimeError(
+                "Attempting to fix ion temperature on an Equilibrium with no "
+                + "ion temperature profile assigned."
+            )
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["T"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixAtomicNumber(FixParameter):
     """Fixes effective atomic number profile coefficients.
@@ -2201,11 +2380,6 @@ def __init__(
         indices=True,
         name="fixed atomic number",
     ):
-        if eq.atomic_number is None:
-            raise RuntimeError(
-                "Attempting to fix atomic number on an Equilibrium with no "
-                + "atomic_number profile assigned."
-            )
         super().__init__(
             thing=eq,
             params={"Zeff_l": indices},
@@ -2217,6 +2391,25 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if eq.atomic_number is None:
+            raise RuntimeError(
+                "Attempting to fix atomic number on an Equilibrium with no "
+                + "atomic_number profile assigned."
+            )
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixPsi(FixParameter):
     """Fixes total toroidal magnetic flux within the last closed flux surface.
@@ -2259,9 +2452,6 @@ def __init__(
         normalize_target=True,
         name="fixed Psi",
     ):
-        if normalize:
-            scales = compute_scaling_factors(eq)
-            self._normalization = scales["Psi"]
         super().__init__(
             thing=eq,
             params={"Psi": True},
@@ -2273,6 +2463,23 @@ def __init__(
             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.
+
+        """
+        eq = self.things[0]
+        if self._normalize:
+            scales = compute_scaling_factors(eq)
+            self._normalization = scales["Psi"]
+        super().build(use_jit=use_jit, verbose=verbose)
+
 
 class FixCurveShift(FixParameter):
     """Fixes Curve.shift attribute, which is redundant with other Curve params.
diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py
index 48ee1ef0f2..668c31b76e 100644
--- a/desc/objectives/objective_funs.py
+++ b/desc/objectives/objective_funs.py
@@ -865,17 +865,26 @@ def _check_dimensions(self):
         if self.bounds is not None:  # must be a tuple of length 2
             self._bounds = tuple([np.asarray(bound) for bound in self._bounds])
             for bound in self.bounds:
-                if not is_broadcastable((self.dim_f,), bound.shape):
+                if (
+                    not is_broadcastable((self.dim_f,), bound.shape)
+                    or bound.size > self.dim_f
+                ):
                     raise ValueError("len(bounds) != dim_f")
             if np.any(self.bounds[1] < self.bounds[0]):
                 raise ValueError("bounds must be: (lower bound, upper bound)")
         else:  # target only gets used if bounds is None
             self._target = np.asarray(self._target)
-            if not is_broadcastable((self.dim_f,), self.target.shape):
+            if (
+                not is_broadcastable((self.dim_f,), self.target.shape)
+                or self.target.size > self.dim_f
+            ):
                 raise ValueError("len(target) != dim_f")
 
         self._weight = np.asarray(self._weight)
-        if not is_broadcastable((self.dim_f,), self.weight.shape):
+        if (
+            not is_broadcastable((self.dim_f,), self.weight.shape)
+            or self.weight.size > self.dim_f
+        ):
             raise ValueError("len(weight) != dim_f")
 
     @abstractmethod
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index c33b229367..c07190c900 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -750,18 +750,18 @@ def test_FixBoundary_passed_target_no_passed_modes_error():
 def test_FixAxis_passed_target_no_passed_modes_error():
     """Test Fixing Axis with no passed-in modes."""
     eq = Equilibrium()
-    FixZ = FixAxisZ(eq=eq, modes=True, target=np.array([0, 0]))
-    with pytest.raises(ValueError):
-        FixZ.build()
-    FixZ = FixAxisZ(eq=eq, modes=False, target=np.array([0, 0]))
-    with pytest.raises(ValueError):
-        FixZ.build()
     FixR = FixAxisR(eq=eq, modes=True, target=np.array([0, 0]))
     with pytest.raises(ValueError):
         FixR.build()
     FixR = FixAxisR(eq=eq, modes=False, target=np.array([0, 0]))
     with pytest.raises(ValueError):
         FixR.build()
+    FixZ = FixAxisZ(eq=eq, modes=True, target=np.array([0, 0]))
+    with pytest.raises(ValueError):
+        FixZ.build()
+    FixZ = FixAxisZ(eq=eq, modes=False, target=np.array([0, 0]))
+    with pytest.raises(ValueError):
+        FixZ.build()
 
 
 @pytest.mark.unit
diff --git a/tests/test_utils.py b/tests/test_utils.py
index b83850022d..6bfadb4008 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -98,7 +98,9 @@ def test_broadcast_tree():
     # tree with proper structure already does not change
     tree_in = tree_out.copy()
     tree = broadcast_tree(tree_in, tree_out)
-    assert tree == tree_in
+    assert tree_structure(tree) == tree_structure(tree_out)
+    for leaf, leaf_correct in zip(tree_leaves(tree), tree_leaves(tree_out)):
+        np.testing.assert_allclose(leaf, leaf_correct)
 
     # broadcast single leaf to full tree
     tree_in = {"a": np.arange(1)}

From 6917eb60ffafe7a1ac3918d870526ce415c73c56 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 30 Apr 2024 12:00:58 -0500
Subject: [PATCH 36/50] copy and paste is dangerous

---
 desc/objectives/linear_objectives.py | 29 +++++++++++++--------------
 desc/objectives/objective_funs.py    | 16 +++++++--------
 desc/utils.py                        |  6 ++++--
 tests/test_linear_objectives.py      | 30 ++++++----------------------
 4 files changed, 31 insertions(+), 50 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 21435cf57b..6621c8e14b 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -616,8 +616,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"Rb_lmn": indices},
@@ -697,8 +697,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.surface.R_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"Zb_lmn": indices},
@@ -802,7 +802,6 @@ def build(self, use_jit=False, verbose=1):
             self._A = A
 
         self._dim_f = self._A.shape[0]
-
         super().build(use_jit=use_jit, verbose=verbose)
 
     def compute(self, params, constants=None):
@@ -919,8 +918,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.axis.R_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"Ra_n": indices},
@@ -1000,8 +999,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.axis.Z_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"Za_n": indices},
@@ -1081,8 +1080,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.R_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"R_lmn": indices},
@@ -1162,8 +1161,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.Z_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"Z_lmn": indices},
@@ -1244,8 +1243,8 @@ def __init__(
             indices = modes
         else:
             indices = np.array([], dtype=int)
-            for l, m, n in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(l, m, n, False))
+            for mode in np.atleast_2d(modes):
+                indices = np.append(indices, eq.L_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"L_lmn": indices},
diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py
index 668c31b76e..853e75d897 100644
--- a/desc/objectives/objective_funs.py
+++ b/desc/objectives/objective_funs.py
@@ -865,25 +865,23 @@ def _check_dimensions(self):
         if self.bounds is not None:  # must be a tuple of length 2
             self._bounds = tuple([np.asarray(bound) for bound in self._bounds])
             for bound in self.bounds:
-                if (
-                    not is_broadcastable((self.dim_f,), bound.shape)
-                    or bound.size > self.dim_f
+                if not is_broadcastable((self.dim_f,), bound.shape) or (
+                    self.dim_f == 1 and bound.size != 1
                 ):
                     raise ValueError("len(bounds) != dim_f")
             if np.any(self.bounds[1] < self.bounds[0]):
                 raise ValueError("bounds must be: (lower bound, upper bound)")
         else:  # target only gets used if bounds is None
             self._target = np.asarray(self._target)
-            if (
-                not is_broadcastable((self.dim_f,), self.target.shape)
-                or self.target.size > self.dim_f
+            print(self.target)
+            if not is_broadcastable((self.dim_f,), self.target.shape) or (
+                self.dim_f == 1 and self.target.size != 1
             ):
                 raise ValueError("len(target) != dim_f")
 
         self._weight = np.asarray(self._weight)
-        if (
-            not is_broadcastable((self.dim_f,), self.weight.shape)
-            or self.weight.size > self.dim_f
+        if not is_broadcastable((self.dim_f,), self.weight.shape) or (
+            self.dim_f == 1 and self.weight.size != 1
         ):
             raise ValueError("len(weight) != dim_f")
 
diff --git a/desc/utils.py b/desc/utils.py
index 38d93c3c45..3447f4804d 100644
--- a/desc/utils.py
+++ b/desc/utils.py
@@ -639,7 +639,8 @@ def broadcast_tree(tree_in, tree_out, dtype=int):
             errorif(
                 key not in tree_out.keys(),
                 ValueError,
-                "dict keys of tree_in must be a subset of those in tree_out",
+                f"dict key '{key}' of tree_in must be a subset of those in tree_out: "
+                + f"{list(tree_out.keys())}",
             )
             if isinstance(value, bool):
                 if value:
@@ -654,7 +655,8 @@ def broadcast_tree(tree_in, tree_out, dtype=int):
             errorif(
                 not np.all(np.isin(tree_new[key], value)),
                 ValueError,
-                "dict values of tree_in must be a subset of those in tree_out",
+                f"dict value {tree_new[key]} of tree_in must be a subset "
+                + f"of those in tree_out: {value}",
             )
         return tree_new
     # tree_out is deeper than tree_in
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index c07190c900..1620982b83 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -799,40 +799,22 @@ def test_FixSumModes_passed_target_too_long():
         )
 
 
-@pytest.mark.unit
-def test_FixMode_False_or_None_modes():
-    """Test Fixing Modes without specifying modes or All modes."""
-    eq = Equilibrium(L=3, M=4)
-    with pytest.raises(ValueError):
-        FixModeR(eq, modes=False, target=np.array([[0, 1]]))
-    with pytest.raises(ValueError):
-        FixModeR(eq, modes=None, target=np.array([[0, 1]]))
-    with pytest.raises(ValueError):
-        FixModeZ(eq, modes=False, target=np.array([[0, 1]]))
-    with pytest.raises(ValueError):
-        FixModeZ(eq, modes=None, target=np.array([[0, 1]]))
-    with pytest.raises(ValueError):
-        FixModeLambda(eq, modes=False, target=np.array([[0, 1]]))
-    with pytest.raises(ValueError):
-        FixModeLambda(eq, modes=None, target=np.array([[0, 1]]))
-
-
 @pytest.mark.unit
 def test_FixSumModes_False_or_None_modes():
     """Test Fixing Sum Modes without specifying modes or All modes."""
     eq = Equilibrium(L=3, M=4)
     with pytest.raises(ValueError):
-        FixSumModesZ(eq, modes=False, target=np.array([[0, 1]]))
+        FixSumModesR(eq, modes=False)
     with pytest.raises(ValueError):
-        FixSumModesZ(eq, modes=None, target=np.array([[0, 1]]))
+        FixSumModesR(eq, modes=None)
     with pytest.raises(ValueError):
-        FixSumModesR(eq, modes=False, target=np.array([[0, 1]]))
+        FixSumModesZ(eq, modes=False)
     with pytest.raises(ValueError):
-        FixSumModesR(eq, modes=None, target=np.array([[0, 1]]))
+        FixSumModesZ(eq, modes=None)
     with pytest.raises(ValueError):
-        FixSumModesLambda(eq, modes=False, target=np.array([[0, 1]]))
+        FixSumModesLambda(eq, modes=False)
     with pytest.raises(ValueError):
-        FixSumModesLambda(eq, modes=None, target=np.array([[0, 1]]))
+        FixSumModesLambda(eq, modes=None)
 
 
 def _is_any_instance(things, cls):

From 32731e68eca0fdbbdad4a02c3a5a30ab015e806e Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 30 Apr 2024 12:11:18 -0500
Subject: [PATCH 37/50] another copy/paste mistake

---
 desc/objectives/linear_objectives.py | 2 +-
 desc/objectives/objective_funs.py    | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 6621c8e14b..b6a9a5f0d0 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -698,7 +698,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.surface.Z_basis.get_idx(*mode, False))
         super().__init__(
             thing=eq,
             params={"Zb_lmn": indices},
diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py
index 853e75d897..0a98bb0a28 100644
--- a/desc/objectives/objective_funs.py
+++ b/desc/objectives/objective_funs.py
@@ -873,7 +873,6 @@ def _check_dimensions(self):
                 raise ValueError("bounds must be: (lower bound, upper bound)")
         else:  # target only gets used if bounds is None
             self._target = np.asarray(self._target)
-            print(self.target)
             if not is_broadcastable((self.dim_f,), self.target.shape) or (
                 self.dim_f == 1 and self.target.size != 1
             ):

From 3f3d3c2628285b8b98567a08c4938542ce47843c Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 30 Apr 2024 13:43:35 -0500
Subject: [PATCH 38/50] syntax issue in freeb test

---
 .../tutorials/free_boundary_equilibrium.ipynb |  7 +++----
 tests/test_examples.py                        | 21 +++++++++----------
 2 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb b/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb
index e1a4d7b651..a645fd1bec 100644
--- a/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb
+++ b/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb
@@ -971,10 +971,9 @@
     "for k in [2, 4]:\n",
     "\n",
     "    # get modes where |m|, |n| > k\n",
-    "    R_modes = (\n",
-    "        eq2.surface.R_basis.modes[np.max(np.abs(eq2.surface.R_basis.modes), 1) > k, :],\n",
-    "    )\n",
-    "\n",
+    "    R_modes = eq2.surface.R_basis.modes[\n",
+    "        np.max(np.abs(eq2.surface.R_basis.modes), 1) > k, :\n",
+    "    ]\n",
     "    Z_modes = eq2.surface.Z_basis.modes[\n",
     "        np.max(np.abs(eq2.surface.Z_basis.modes), 1) > k, :\n",
     "    ]\n",
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 65e2c2d524..de25985002 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1030,9 +1030,9 @@ def test_freeb_vacuum():
         modes_Z=[[-1, 0]],
         NFP=5,
     )
-
     eq = Equilibrium(M=6, N=6, Psi=-0.035, surface=surf)
     eq.solve()
+
     constraints = (
         ForceBalance(eq=eq),
         FixCurrent(eq=eq),
@@ -1042,15 +1042,15 @@ def test_freeb_vacuum():
     objective = ObjectiveFunction(
         VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True)
     )
-    eq, out = eq.optimize(
+    eq, _ = eq.optimize(
         objective,
         constraints,
         optimizer="proximal-lsq-exact",
         verbose=3,
         options={},
     )
-    rho_err, _ = area_difference_vmec(eq, "tests/inputs/wout_test_freeb.nc")
 
+    rho_err, _ = area_difference_vmec(eq, "tests/inputs/wout_test_freeb.nc")
     np.testing.assert_allclose(rho_err[:, -1], 0, atol=4e-2)  # only check rho=1
 
 
@@ -1088,9 +1088,9 @@ def test_freeb_axisym():
         modes_Z=[[-1, 0]],
         NFP=1,
     )
-
     eq = Equilibrium(M=10, N=0, Psi=1.0, surface=surf, pressure=pres, iota=iota)
     eq.solve()
+
     constraints = (
         ForceBalance(eq=eq),
         FixIota(eq=eq),
@@ -1102,27 +1102,26 @@ def test_freeb_axisym():
     )
 
     # we know this is a pretty simple shape so we'll only use |m| <= 2
-    R_modes = (
-        eq.surface.R_basis.modes[np.max(np.abs(eq.surface.R_basis.modes), 1) > 2, :],
-    )
-
+    R_modes = eq.surface.R_basis.modes[
+        np.max(np.abs(eq.surface.R_basis.modes), 1) > 2, :
+    ]
     Z_modes = eq.surface.Z_basis.modes[
         np.max(np.abs(eq.surface.Z_basis.modes), 1) > 2, :
     ]
-
     bdry_constraints = (
         FixBoundaryR(eq=eq, modes=R_modes),
         FixBoundaryZ(eq=eq, modes=Z_modes),
     )
-    eq, out = eq.optimize(
+
+    eq, _ = eq.optimize(
         objective,
         constraints + bdry_constraints,
         optimizer="proximal-lsq-exact",
         verbose=3,
         options={},
     )
-    rho_err, _ = area_difference_vmec(eq, "tests/inputs/wout_solovev_freeb.nc")
 
+    rho_err, _ = area_difference_vmec(eq, "tests/inputs/wout_solovev_freeb.nc")
     np.testing.assert_allclose(rho_err[:, -1], 0, atol=2e-2)  # only check rho=1
 
 

From 94532a6c00adf631528ba5f11bbc8fa160edaf98 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Tue, 30 Apr 2024 14:31:47 -0500
Subject: [PATCH 39/50] syntax issue in freeb notebook

---
 docs/notebooks/tutorials/free_boundary_equilibrium.ipynb | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb b/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb
index a645fd1bec..9c7929f24a 100644
--- a/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb
+++ b/docs/notebooks/tutorials/free_boundary_equilibrium.ipynb
@@ -446,12 +446,8 @@
    ],
    "source": [
     "# we know this is a pretty simple shape so we'll only use |m| <= 2\n",
-    "R_modes = (\n",
-    "    eq2.surface.R_basis.modes[np.max(np.abs(eq2.surface.R_basis.modes), 1) > 2, :],\n",
-    ")\n",
-    "\n",
+    "R_modes = eq2.surface.R_basis.modes[np.max(np.abs(eq2.surface.R_basis.modes), 1) > 2, :]\n",
     "Z_modes = eq2.surface.Z_basis.modes[np.max(np.abs(eq2.surface.Z_basis.modes), 1) > 2, :]\n",
-    "\n",
     "bdry_constraints = (\n",
     "    FixBoundaryR(eq=eq2, modes=R_modes),\n",
     "    FixBoundaryZ(eq=eq2, modes=Z_modes),\n",

From 13305b5b88a7e5b3df70bea4b6e98cc8fae4ac26 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 11:10:16 -0500
Subject: [PATCH 40/50] make Rory's suggested changes

---
 desc/objectives/__init__.py          |   2 +-
 desc/objectives/getters.py           |  42 +++----
 desc/objectives/linear_objectives.py | 159 ++++++---------------------
 docs/api.rst                         |   2 +-
 docs/api_objectives.rst              |   2 +-
 tests/test_examples.py               |  26 ++---
 tests/test_linear_objectives.py      |   8 +-
 tests/test_optimizer.py              |  12 +-
 8 files changed, 83 insertions(+), 170 deletions(-)

diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py
index 5357e78013..a99da8cfcf 100644
--- a/desc/objectives/__init__.py
+++ b/desc/objectives/__init__.py
@@ -63,7 +63,7 @@
     FixOmniBmax,
     FixOmniMap,
     FixOmniWell,
-    FixParameter,
+    FixParameters,
     FixPressure,
     FixPsi,
     FixSheetCurrent,
diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index 489e5401b7..8a33b03800 100644
--- a/desc/objectives/getters.py
+++ b/desc/objectives/getters.py
@@ -214,28 +214,32 @@ def get_NAE_constraints(
 
 def maybe_add_self_consistency(thing, constraints):
     """Add self consistency constraints if needed."""
+    params = set(unique_list(flatten_list(thing.optimizable_params))[0])
+
     # Equilibrium
-    if {"Rb_lmn", "Zb_lmn", "L_lmn", "Ra_n", "Za_n"} <= set(
-        unique_list(flatten_list(thing.optimizable_params))[0]
+    if {"R_lmn", "Rb_lmn"} <= params and not is_any_instance(
+        constraints, BoundaryRSelfConsistency
+    ):
+        constraints += (BoundaryRSelfConsistency(eq=thing),)
+    if {"Z_lmn", "Zb_lmn"} <= params and not is_any_instance(
+        constraints, BoundaryZSelfConsistency
+    ):
+        constraints += (BoundaryZSelfConsistency(eq=thing),)
+    if {"L_lmn"} <= params and not is_any_instance(constraints, FixLambdaGauge):
+        constraints += (FixLambdaGauge(eq=thing),)
+    if {"R_lmn", "Ra_lmn"} <= params and not is_any_instance(
+        constraints, AxisRSelfConsistency
     ):
-        if not is_any_instance(constraints, BoundaryRSelfConsistency):
-            constraints += (BoundaryRSelfConsistency(eq=thing),)
-        if not is_any_instance(constraints, BoundaryZSelfConsistency):
-            constraints += (BoundaryZSelfConsistency(eq=thing),)
-        if not is_any_instance(constraints, FixLambdaGauge):
-            constraints += (FixLambdaGauge(eq=thing),)
-        if not is_any_instance(constraints, AxisRSelfConsistency):
-            constraints += (AxisRSelfConsistency(eq=thing),)
-        if not is_any_instance(constraints, AxisZSelfConsistency):
-            constraints += (AxisZSelfConsistency(eq=thing),)
+        constraints += (AxisRSelfConsistency(eq=thing),)
+    if {"Z_lmn", "Za_lmn"} <= params and not is_any_instance(
+        constraints, AxisZSelfConsistency
+    ):
+        constraints += (AxisZSelfConsistency(eq=thing),)
 
     # Curve
-    elif {"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):
-            constraints += (FixCurveRotation(curve=thing),)
+    if {"shift"} <= params and not is_any_instance(constraints, FixCurveShift):
+        constraints += (FixCurveShift(curve=thing),)
+    if {"rotmat"} <= params and not is_any_instance(constraints, FixCurveRotation):
+        constraints += (FixCurveRotation(curve=thing),)
 
     return constraints
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index b6a9a5f0d0..698f397d89 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -19,7 +19,7 @@
 from .objective_funs import _Objective
 
 
-# TODO: get rid of this class and inherit from FixParameter instead?
+# TODO: get rid of this class and inherit from FixParameters instead?
 class _FixedObjective(_Objective):
     _fixed = True
     _linear = True
@@ -64,7 +64,7 @@ def _parse_target_from_user(
 
 
 # Rename to `FixParameters` (plural) instead?
-class FixParameter(_Objective):
+class FixParameters(_Objective):
     """Fix specific degrees of freedom associated with a given Optimizable thing.
 
     Parameters
@@ -566,7 +566,7 @@ def compute(self, params, constants=None):
         return f
 
 
-class FixBoundaryR(FixParameter):
+class FixBoundaryR(FixParameters):
     """Boundary condition on the R boundary parameters.
 
     Parameters
@@ -617,7 +617,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.R_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.surface.R_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"Rb_lmn": indices},
@@ -647,7 +647,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixBoundaryZ(FixParameter):
+class FixBoundaryZ(FixParameters):
     """Boundary condition on the Z boundary parameters.
 
     Parameters
@@ -698,7 +698,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.surface.Z_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.surface.Z_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"Zb_lmn": indices},
@@ -759,8 +759,6 @@ def __init__(
             target=0,
             bounds=None,
             weight=1,
-            normalize=False,
-            normalize_target=False,
             name=name,
         )
 
@@ -824,7 +822,7 @@ def compute(self, params, constants=None):
         return jnp.dot(self._A, params["L_lmn"])
 
 
-class FixThetaSFL(FixParameter):
+class FixThetaSFL(FixParameters):
     """Fixes lambda=0 so that poloidal angle is the SFL poloidal angle.
 
     Parameters
@@ -834,12 +832,6 @@ class FixThetaSFL(FixParameter):
     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.
-    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.
     name : str, optional
         Name of the objective function.
 
@@ -852,8 +844,6 @@ def __init__(
         self,
         eq,
         weight=1,
-        normalize=True,
-        normalize_target=True,
         name="theta SFL",
     ):
         super().__init__(
@@ -862,13 +852,11 @@ def __init__(
             target=0,
             bounds=None,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
 
-class FixAxisR(FixParameter):
+class FixAxisR(FixParameters):
     """Fixes magnetic axis R coefficients.
 
     Parameters
@@ -919,7 +907,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.axis.R_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.axis.R_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"Ra_n": indices},
@@ -949,7 +937,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixAxisZ(FixParameter):
+class FixAxisZ(FixParameters):
     """Fixes magnetic axis Z coefficients.
 
     Parameters
@@ -1000,7 +988,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.axis.Z_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.axis.Z_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"Za_n": indices},
@@ -1030,7 +1018,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixModeR(FixParameter):
+class FixModeR(FixParameters):
     """Fixes Fourier-Zernike R coefficients.
 
     Parameters
@@ -1081,7 +1069,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.R_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.R_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"R_lmn": indices},
@@ -1111,7 +1099,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixModeZ(FixParameter):
+class FixModeZ(FixParameters):
     """Fixes Fourier-Zernike Z coefficients.
 
     Parameters
@@ -1162,7 +1150,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.Z_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.Z_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"Z_lmn": indices},
@@ -1192,7 +1180,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixModeLambda(FixParameter):
+class FixModeLambda(FixParameters):
     """Fixes Fourier-Zernike lambda coefficients.
 
     Parameters
@@ -1210,12 +1198,6 @@ class FixModeLambda(FixParameter):
     weight : float, ndarray, optional
         Weighting to apply to the Objective, relative to other Objectives.
         Must be broadcastable to Objective.dim_f.
-    normalize : bool
-        Whether to compute the error in physical units or non-dimensionalize.
-    normalize_target : bool
-        Whether target 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.
     modes : ndarray, optional
         Basis modes numbers [l,m,n] of Fourier-Zernike modes to fix.
         len(target) = len(weight) = len(modes).
@@ -1234,8 +1216,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=True,
-        normalize_target=True,
         modes=True,
         name="fix mode lambda",
     ):
@@ -1244,7 +1224,7 @@ def __init__(
         else:
             indices = np.array([], dtype=int)
             for mode in np.atleast_2d(modes):
-                indices = np.append(indices, eq.L_basis.get_idx(*mode, False))
+                indices = np.append(indices, eq.L_basis.get_idx(*mode))
         super().__init__(
             thing=eq,
             params={"L_lmn": indices},
@@ -1252,8 +1232,6 @@ def __init__(
             bounds=bounds,
             weight=weight,
             name=name,
-            normalize=normalize,
-            normalize_target=normalize_target,
         )
 
 
@@ -1606,12 +1584,6 @@ class FixSumModesLambda(_FixedObjective):
     weight : {float, ndarray}, optional
         Weighting to apply to the Objective, relative to other Objectives.
         Must be broadcastable to to Objective.dim_f.
-    normalize : bool
-        Whether to compute the error in physical units or non-dimensionalize.
-    normalize_target : bool
-        Whether target 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.
     sum_weight : float, ndarray, optional
         Weights on the coefficients in the sum, should be same length as modes.
         Defaults to 1 i.e. target = 1*L_111 + 1*L_222...
@@ -1637,8 +1609,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=True,
-        normalize_target=True,
         sum_weights=None,
         modes=True,
         name="Fix Sum Modes lambda",
@@ -1671,8 +1641,6 @@ def __init__(
             bounds=bounds,
             weight=weight,
             name=name,
-            normalize=normalize,
-            normalize_target=normalize_target,
         )
 
     def build(self, use_jit=False, verbose=1):
@@ -1759,7 +1727,7 @@ def compute(self, params, constants=None):
         return f
 
 
-class FixPressure(FixParameter):
+class FixPressure(FixParameters):
     """Fixes pressure coefficients.
 
     Parameters
@@ -1841,7 +1809,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixAnisotropy(FixParameter):
+class FixAnisotropy(FixParameters):
     """Fixes anisotropic pressure coefficients.
 
     Parameters
@@ -1858,12 +1826,6 @@ class FixAnisotropy(FixParameter):
     weight : {float, ndarray}, optional
         Weighting to apply to the Objective, relative to other Objectives.
         Must be broadcastable to to Objective.dim_f
-    normalize : bool
-        Whether to compute the error in physical units or non-dimensionalize.
-    normalize_target : bool
-        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.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -1884,8 +1846,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=True,
-        normalize_target=True,
         indices=True,
         name="fixed anisotropy",
     ):
@@ -1895,8 +1855,6 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
@@ -1920,7 +1878,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixIota(FixParameter):
+class FixIota(FixParameters):
     """Fixes rotational transform coefficients.
 
     Parameters
@@ -1937,13 +1895,6 @@ class FixIota(FixParameter):
     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. Has no effect for this objective.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices.
@@ -1964,8 +1915,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=False,
-        normalize_target=False,
         indices=True,
         name="fixed iota",
     ):
@@ -1975,8 +1924,6 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2000,7 +1947,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixCurrent(FixParameter):
+class FixCurrent(FixParameters):
     """Fixes toroidal current profile coefficients.
 
     Parameters
@@ -2082,7 +2029,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixElectronTemperature(FixParameter):
+class FixElectronTemperature(FixParameters):
     """Fixes electron temperature profile coefficients.
 
     Parameters
@@ -2164,7 +2111,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixElectronDensity(FixParameter):
+class FixElectronDensity(FixParameters):
     """Fixes electron density profile coefficients.
 
     Parameters
@@ -2248,7 +2195,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixIonTemperature(FixParameter):
+class FixIonTemperature(FixParameters):
     """Fixes ion temperature profile coefficients.
 
     Parameters
@@ -2330,7 +2277,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixAtomicNumber(FixParameter):
+class FixAtomicNumber(FixParameters):
     """Fixes effective atomic number profile coefficients.
 
     Parameters
@@ -2347,13 +2294,6 @@ class FixAtomicNumber(FixParameter):
     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. Has no effect for this objective.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -2374,8 +2314,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=False,
-        normalize_target=False,
         indices=True,
         name="fixed atomic number",
     ):
@@ -2385,8 +2323,6 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2410,7 +2346,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixPsi(FixParameter):
+class FixPsi(FixParameters):
     """Fixes total toroidal magnetic flux within the last closed flux surface.
 
     Parameters
@@ -2480,7 +2416,7 @@ def build(self, use_jit=False, verbose=1):
         super().build(use_jit=use_jit, verbose=verbose)
 
 
-class FixCurveShift(FixParameter):
+class FixCurveShift(FixParameters):
     """Fixes Curve.shift attribute, which is redundant with other Curve params.
 
     Parameters
@@ -2530,9 +2466,10 @@ def __init__(
             normalize_target=normalize_target,
             name=name,
         )
+        # TODO: add normalization?
 
 
-class FixCurveRotation(FixParameter):
+class FixCurveRotation(FixParameters):
     """Fixes Curve.rotmat attribute, which is redundant with other Curve params.
 
     Parameters
@@ -2548,12 +2485,6 @@ class FixCurveRotation(FixParameter):
     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.
-    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.
     name : str, optional
         Name of the objective function.
 
@@ -2568,8 +2499,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=True,
-        normalize_target=True,
         name="fixed rotation",
     ):
         super().__init__(
@@ -2578,13 +2507,11 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
 
-class FixOmniWell(FixParameter):
+class FixOmniWell(FixParameters):
     """Fixes OmnigenousField.B_lm coefficients.
 
     Parameters
@@ -2636,9 +2563,10 @@ def __init__(
             normalize_target=normalize_target,
             name=name,
         )
+        # TODO: add normalization?
 
 
-class FixOmniMap(FixParameter):
+class FixOmniMap(FixParameters):
     """Fixes OmnigenousField.x_lmn coefficients.
 
     Parameters
@@ -2651,12 +2579,6 @@ class FixOmniMap(FixParameter):
         Lower and upper bounds on the objective. Overrides target.
     weight : float, optional
         Weighting to apply to the Objective, relative to other Objectives.
-    normalize : bool
-        Whether to compute the error in physical units or non-dimensionalize.
-    normalize_target : bool
-        Whether target 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.
     indices : ndarray or bool, optional
         indices of the field.x_lmn array to fix.
         Must have len(target) = len(weight) = len(indices).
@@ -2675,8 +2597,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=False,
-        normalize_target=False,
         indices=True,
         name="fixed omnigenity map",
     ):
@@ -2686,8 +2606,6 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2705,12 +2623,6 @@ class FixOmniBmax(_FixedObjective):
         Lower and upper bounds on the objective. Overrides target.
     weight : float, optional
         Weighting to apply to the Objective, relative to other Objectives.
-    normalize : bool
-        Whether to compute the error in physical units or non-dimensionalize.
-    normalize_target : bool
-        Whether target 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.
     name : str
         Name of the objective function.
 
@@ -2726,8 +2638,6 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=False,
-        normalize_target=False,
         name="fixed omnigenity B_max",
     ):
         self._target_from_user = setdefault(bounds, target)
@@ -2736,8 +2646,6 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
-            normalize=normalize,
-            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2798,7 +2706,7 @@ def compute(self, params, constants=None):
         return f
 
 
-class FixSheetCurrent(FixParameter):
+class FixSheetCurrent(FixParameters):
     """Fixes the sheet current parameters of a free-boundary equilibrium.
 
     Note: this constraint is automatically applied when needed, and does not need to be
@@ -2852,3 +2760,4 @@ def __init__(
             normalize_target=normalize_target,
             name=name,
         )
+        # TODO: add normalization?
diff --git a/docs/api.rst b/docs/api.rst
index 7f5cb4c9ae..ae0c1efe3f 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -177,7 +177,7 @@ Objective Functions
     desc.objectives.FixOmniBmax
     desc.objectives.FixOmniMap
     desc.objectives.FixOmniWell
-    desc.objectives.FixParameter
+    desc.objectives.FixParameters
     desc.objectives.FixPressure
     desc.objectives.FixPsi
     desc.objectives.FixSumModesR
diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst
index 48c650194b..c0885210d8 100644
--- a/docs/api_objectives.rst
+++ b/docs/api_objectives.rst
@@ -126,7 +126,7 @@ Fixing degrees of freedom
     desc.objectives.FixSumModesR
     desc.objectives.FixSumModesZ
     desc.objectives.FixThetaSFL
-    desc.objectives.FixParameter
+    desc.objectives.FixParameters
 
 
 User defined objectives
diff --git a/tests/test_examples.py b/tests/test_examples.py
index de25985002..c19752c903 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -36,7 +36,7 @@
     FixIota,
     FixOmniBmax,
     FixOmniMap,
-    FixParameter,
+    FixParameters,
     FixPressure,
     FixPsi,
     FixSumModesLambda,
@@ -640,8 +640,8 @@ def test_multiobject_optimization_al():
     constraints = (
         ForceBalance(eq=eq, bounds=(-1e-4, 1e-4), normalize_target=False),
         FixPressure(eq=eq),
-        FixParameter(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}),
-        FixParameter(eq, {"Psi": True, "i_l": True}),
+        FixParameters(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}),
+        FixParameters(eq, {"Psi": True, "i_l": True}),
         FixBoundaryR(eq, modes=[[0, 0, 0]]),
         PlasmaVesselDistance(surface=surf, eq=eq, target=1),
     )
@@ -679,8 +679,8 @@ def test_multiobject_optimization_prox():
     constraints = (
         ForceBalance(eq=eq),
         FixPressure(eq=eq),
-        FixParameter(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}),
-        FixParameter(eq, {"Psi": True, "i_l": True}),
+        FixParameters(surf, {"R_lmn": np.array([0]), "Z_lmn": np.array([3])}),
+        FixParameters(eq, {"Psi": True, "i_l": True}),
         FixBoundaryR(eq, modes=[[0, 0, 0]]),
     )
 
@@ -970,7 +970,7 @@ def test_non_eq_optimization():
 
     surf.change_resolution(M=eq.M, N=eq.N)
     constraints = (
-        FixParameter(eq),
+        FixParameters(eq),
         MeanCurvature(surf, bounds=(-8, 8)),
         PrincipalCurvature(surf, bounds=(0, 15)),
     )
@@ -1001,7 +1001,7 @@ def test_only_non_eq_optimization():
     surf = eq.surface
     surf.change_resolution(M=eq.M, N=eq.N)
     constraints = (
-        FixParameter(surf, {"R_lmn": np.array(surf.R_basis.get_idx(0, 0, 0))}),
+        FixParameters(surf, {"R_lmn": np.array(surf.R_basis.get_idx(0, 0, 0))}),
     )
     obj = PrincipalCurvature(surf, target=1)
     objective = ObjectiveFunction((obj,))
@@ -1252,7 +1252,7 @@ def test_quadratic_flux_optimization_with_analytic_field():
 
     optimizer = Optimizer("lsq-exact")
 
-    constraints = (FixParameter(field, {"R0": True}),)
+    constraints = (FixParameters(field, {"R0": True}),)
     quadflux_obj = QuadraticFlux(
         eq=eq,
         field=field,
@@ -1279,14 +1279,14 @@ def test_quadratic_flux_optimization_with_analytic_field():
 def test_second_stage_optimization():
     """Test optimizing magnetic field for a fixed axisymmetric equilibrium."""
     eq = get("DSHAPE")
+    grid = LinearGrid(rho=np.array([1.0]), M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
     field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
-    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
-    constraints = FixParameter(field, [{"R0": True}, {}])
-    optimizer = Optimizer("lsq-exact")
+    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field, eval_grid=grid))
+    constraints = FixParameters(field, [{"R0": True}, {}])
+    optimizer = Optimizer("scipy-trf")
     (field,), _ = optimizer.optimize(
         things=field, objective=objective, constraints=constraints, verbose=2
     )
-    # TODO: could make this more robust instead of assuming only vertical field changes
-    np.testing.assert_allclose(field[0].R0, 3.5)
+    np.testing.assert_allclose(field[0].R0, 3.5)  # this value was fixed
     np.testing.assert_allclose(field[0].B0, 1)  # toroidal field (does not change)
     np.testing.assert_allclose(field[1].B0, -0.022, rtol=1e-2)  # vertical field
diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index 1620982b83..e3838bc0ee 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -40,7 +40,7 @@
     FixModeZ,
     FixOmniMap,
     FixOmniWell,
-    FixParameter,
+    FixParameters,
     FixPressure,
     FixPsi,
     FixSumModesLambda,
@@ -353,7 +353,7 @@ def test_factorize_linear_constraints_asserts():
         _ = factorize_linear_constraints(objective, constraint)
 
     # constraining a foreign thing
-    constraint = ObjectiveFunction(FixParameter(surf))
+    constraint = ObjectiveFunction(FixParameters(surf))
     constraint.build(verbose=0)
     with pytest.raises(UserWarning):
         _ = factorize_linear_constraints(objective, constraint)
@@ -911,7 +911,7 @@ def test_fix_omni_indices():
 
 @pytest.mark.unit
 def test_fix_subset_of_params_in_collection():
-    """Tests FixParameter fixing a subset of things in the collection."""
+    """Tests FixParameters fixing a subset of things in the collection."""
     tf_coil = FourierPlanarCoil(center=[2, 0, 0], normal=[0, 1, 0], r_n=[1])
     tf_coilset = CoilSet.linspaced_angular(tf_coil, n=4)
     vf_coil = FourierRZCoil(R_n=3, Z_n=-1)
@@ -944,6 +944,6 @@ def test_fix_subset_of_params_in_collection():
         )
     )
 
-    obj = FixParameter(full_coilset, params)
+    obj = FixParameters(full_coilset, params)
     obj.build()
     np.testing.assert_allclose(obj.target, target)
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 8955008b21..18bf4bf7a7 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -27,7 +27,7 @@
     FixBoundaryZ,
     FixCurrent,
     FixIota,
-    FixParameter,
+    FixParameters,
     FixPressure,
     FixPsi,
     ForceBalance,
@@ -1005,8 +1005,8 @@ def test_optimize_multiple_things_different_order():
         NFP=eq.NFP,
     )
     constraints = (
-        FixParameter(eq),  # don't let eq vary
-        FixParameter(  # only let the minor radius of the surface vary
+        FixParameters(eq),  # don't let eq vary
+        FixParameters(  # only let the minor radius of the surface vary
             surf, params={"R_lmn": np.array(surf.R_basis.get_idx(M=0, N=0))}
         ),
     )
@@ -1045,8 +1045,8 @@ def test_optimize_multiple_things_different_order():
 
     # fresh start
     constraints = (
-        FixParameter(eq),  # don't let eq vary
-        FixParameter(  # only let the minor radius of the surface vary
+        FixParameters(eq),  # don't let eq vary
+        FixParameters(  # only let the minor radius of the surface vary
             surf, params={"R_lmn": np.array(surf.R_basis.get_idx(M=0, N=0))}
         ),
     )
@@ -1085,7 +1085,7 @@ def test_optimize_with_single_constraint():
     eq = Equilibrium()
     optimizer = Optimizer("lsq-exact")
     objectective = ObjectiveFunction(GenericObjective("|B|", eq), use_jit=False)
-    constraints = FixParameter(
+    constraints = FixParameters(
         eq,
         {
             "R_lmn": True,

From f524cc04d67dad1553af64de6c5df9500ef499cb Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 11:23:40 -0500
Subject: [PATCH 41/50] repair getter funs

---
 desc/objectives/getters.py | 35 +++++++++++++++++------------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index 8a33b03800..ed3cdb9836 100644
--- a/desc/objectives/getters.py
+++ b/desc/objectives/getters.py
@@ -59,6 +59,7 @@ def get_equilibrium_objective(eq, mode="force", normalize=True):
     -------
     objective, ObjectiveFunction
         An objective function with default force balance objectives.
+
     """
     if mode == "energy":
         objectives = Energy(eq=eq, normalize=normalize, normalize_target=normalize)
@@ -94,17 +95,16 @@ def get_fixed_axis_constraints(eq, profiles=True, normalize=True):
         A list of the linear constraints used in fixed-axis problems.
 
     """
+    normalize_kwargs = {"normalize": normalize, "normalize_target": normalize}
     constraints = (
-        FixAxisR(eq=eq, normalize=normalize, normalize_target=normalize),
-        FixAxisZ(eq=eq, normalize=normalize, normalize_target=normalize),
-        FixPsi(eq=eq, normalize=normalize, normalize_target=normalize),
+        FixAxisR(eq=eq, **normalize_kwargs),
+        FixAxisZ(eq=eq, **normalize_kwargs),
+        FixPsi(eq=eq, **normalize_kwargs),
     )
     if profiles:
         for name, con in _PROFILE_CONSTRAINTS.items():
             if getattr(eq, name) is not None:
-                constraints += (
-                    con(eq=eq, normalize=normalize, normalize_target=normalize),
-                )
+                constraints += (con(eq=eq, **normalize_kwargs),)
     constraints += (FixSheetCurrent(eq),)
 
     return constraints
@@ -128,17 +128,16 @@ def get_fixed_boundary_constraints(eq, profiles=True, normalize=True):
         A list of the linear constraints used in fixed-boundary problems.
 
     """
+    normalize_kwargs = {"normalize": normalize, "normalize_target": normalize}
     constraints = (
-        FixBoundaryR(eq=eq, normalize=normalize, normalize_target=normalize),
-        FixBoundaryZ(eq=eq, normalize=normalize, normalize_target=normalize),
-        FixPsi(eq=eq, normalize=normalize, normalize_target=normalize),
+        FixBoundaryR(eq=eq, **normalize_kwargs),
+        FixBoundaryZ(eq=eq, **normalize_kwargs),
+        FixPsi(eq=eq, **normalize_kwargs),
     )
     if profiles:
         for name, con in _PROFILE_CONSTRAINTS.items():
             if getattr(eq, name) is not None:
-                constraints += (
-                    con(eq=eq, normalize=normalize, normalize_target=normalize),
-                )
+                constraints += (con(eq=eq, **normalize_kwargs),)
     constraints += (FixSheetCurrent(eq),)
 
     return constraints
@@ -180,21 +179,21 @@ def get_NAE_constraints(
     -------
     constraints, tuple of _Objectives
         A list of the linear constraints used in fixed-axis problems.
+
     """
+    normalize_kwargs = {"normalize": normalize, "normalize_target": normalize}
     if not isinstance(fix_lambda, bool):
         fix_lambda = int(fix_lambda)
     constraints = (
-        FixAxisR(eq=desc_eq, normalize=normalize, normalize_target=normalize),
-        FixAxisZ(eq=desc_eq, normalize=normalize, normalize_target=normalize),
-        FixPsi(eq=desc_eq, normalize=normalize, normalize_target=normalize),
+        FixAxisR(eq=desc_eq, **normalize_kwargs),
+        FixAxisZ(eq=desc_eq, **normalize_kwargs),
+        FixPsi(eq=desc_eq, **normalize_kwargs),
     )
 
     if profiles:
         for name, con in _PROFILE_CONSTRAINTS.items():
             if getattr(desc_eq, name) is not None:
-                constraints += (
-                    con(eq=desc_eq, normalize=normalize, normalize_target=normalize),
-                )
+                constraints += (con(eq=desc_eq, **normalize_kwargs),)
     constraints += (FixSheetCurrent(desc_eq),)
 
     if fix_lambda or (fix_lambda >= 0 and type(fix_lambda) is int):

From 255f242023a7dd6bd466cf88a158edc783910257 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 11:54:43 -0500
Subject: [PATCH 42/50] add test to check input order is preserved

---
 tests/test_linear_objectives.py | 43 +++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py
index e3838bc0ee..0aa7daf203 100644
--- a/tests/test_linear_objectives.py
+++ b/tests/test_linear_objectives.py
@@ -909,6 +909,49 @@ def test_fix_omni_indices():
     assert constraint.dim_f == indices.size
 
 
+@pytest.mark.unit
+def test_fix_parameters_input_order(DummyStellarator):
+    """Test that FixParameters preserves the input indices and target ordering."""
+    eq = load(load_from=str(DummyStellarator["output_path"]), file_format="hdf5")
+    default_target = eq.Rb_lmn
+
+    # default objective
+    obj = FixBoundaryR(eq)
+    obj.build()
+    np.testing.assert_allclose(obj.target, default_target)
+
+    # manually specify default
+    obj = FixBoundaryR(eq, modes=eq.surface.R_basis.modes)
+    obj.build()
+    np.testing.assert_allclose(obj.target, default_target)
+
+    # reverse order
+    obj = FixBoundaryR(eq, modes=np.flipud(eq.surface.R_basis.modes))
+    obj.build()
+    np.testing.assert_allclose(obj.target, np.flipud(default_target))
+
+    # custom order
+    obj = ObjectiveFunction(
+        FixBoundaryR(eq, modes=np.array([[0, 0, 0], [0, 1, 0], [0, 1, 1]]))
+    )
+    obj.build()
+    np.testing.assert_allclose(obj.target_scaled, np.array([3, 1, 0.3]))
+    np.testing.assert_allclose(obj.compute_scaled_error(obj.x(eq)), np.zeros(obj.dim_f))
+
+    # custom target
+    obj = ObjectiveFunction(
+        FixBoundaryR(
+            eq,
+            modes=np.array([[0, 0, 0], [0, 1, 0], [0, 1, 1]]),
+            target=np.array([0, -1, 0.5]),
+        )
+    )
+    obj.build()
+    np.testing.assert_allclose(
+        obj.compute_scaled_error(obj.x(eq)), np.array([3, 2, -0.2])
+    )
+
+
 @pytest.mark.unit
 def test_fix_subset_of_params_in_collection():
     """Tests FixParameters fixing a subset of things in the collection."""

From 1752190819369ff7fee0b7a110e38f5eb473bd1f Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 14:21:56 -0500
Subject: [PATCH 43/50] re-add normalize kwargs

---
 desc/objectives/getters.py           | 48 ++++++----------
 desc/objectives/linear_objectives.py | 85 +++++++++++++++++++++++++++-
 2 files changed, 98 insertions(+), 35 deletions(-)

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index ed3cdb9836..756a1b6acf 100644
--- a/desc/objectives/getters.py
+++ b/desc/objectives/getters.py
@@ -61,17 +61,13 @@ def get_equilibrium_objective(eq, mode="force", normalize=True):
         An objective function with default force balance objectives.
 
     """
+    kwargs = {"eq": eq, "normalize": normalize, "normalize_target": normalize}
     if mode == "energy":
-        objectives = Energy(eq=eq, normalize=normalize, normalize_target=normalize)
+        objectives = Energy(**kwargs)
     elif mode == "force":
-        objectives = ForceBalance(
-            eq=eq, normalize=normalize, normalize_target=normalize
-        )
+        objectives = ForceBalance(**kwargs)
     elif mode == "forces":
-        objectives = (
-            RadialForceBalance(eq=eq, normalize=normalize, normalize_target=normalize),
-            HelicalForceBalance(eq=eq, normalize=normalize, normalize_target=normalize),
-        )
+        objectives = (RadialForceBalance(**kwargs), HelicalForceBalance(**kwargs))
     else:
         raise ValueError("got an unknown equilibrium objective type '{}'".format(mode))
     return ObjectiveFunction(objectives)
@@ -95,17 +91,13 @@ def get_fixed_axis_constraints(eq, profiles=True, normalize=True):
         A list of the linear constraints used in fixed-axis problems.
 
     """
-    normalize_kwargs = {"normalize": normalize, "normalize_target": normalize}
-    constraints = (
-        FixAxisR(eq=eq, **normalize_kwargs),
-        FixAxisZ(eq=eq, **normalize_kwargs),
-        FixPsi(eq=eq, **normalize_kwargs),
-    )
+    kwargs = {"eq": eq, "normalize": normalize, "normalize_target": normalize}
+    constraints = (FixAxisR(**kwargs), FixAxisZ(**kwargs), FixPsi(**kwargs))
     if profiles:
         for name, con in _PROFILE_CONSTRAINTS.items():
             if getattr(eq, name) is not None:
-                constraints += (con(eq=eq, **normalize_kwargs),)
-    constraints += (FixSheetCurrent(eq),)
+                constraints += (con(**kwargs),)
+    constraints += (FixSheetCurrent(**kwargs),)
 
     return constraints
 
@@ -128,17 +120,13 @@ def get_fixed_boundary_constraints(eq, profiles=True, normalize=True):
         A list of the linear constraints used in fixed-boundary problems.
 
     """
-    normalize_kwargs = {"normalize": normalize, "normalize_target": normalize}
-    constraints = (
-        FixBoundaryR(eq=eq, **normalize_kwargs),
-        FixBoundaryZ(eq=eq, **normalize_kwargs),
-        FixPsi(eq=eq, **normalize_kwargs),
-    )
+    kwargs = {"eq": eq, "normalize": normalize, "normalize_target": normalize}
+    constraints = (FixBoundaryR(**kwargs), FixBoundaryZ(**kwargs), FixPsi(**kwargs))
     if profiles:
         for name, con in _PROFILE_CONSTRAINTS.items():
             if getattr(eq, name) is not None:
-                constraints += (con(eq=eq, **normalize_kwargs),)
-    constraints += (FixSheetCurrent(eq),)
+                constraints += (con(**kwargs),)
+    constraints += (FixSheetCurrent(**kwargs),)
 
     return constraints
 
@@ -181,20 +169,16 @@ def get_NAE_constraints(
         A list of the linear constraints used in fixed-axis problems.
 
     """
-    normalize_kwargs = {"normalize": normalize, "normalize_target": normalize}
+    kwargs = {"eq": desc_eq, "normalize": normalize, "normalize_target": normalize}
     if not isinstance(fix_lambda, bool):
         fix_lambda = int(fix_lambda)
-    constraints = (
-        FixAxisR(eq=desc_eq, **normalize_kwargs),
-        FixAxisZ(eq=desc_eq, **normalize_kwargs),
-        FixPsi(eq=desc_eq, **normalize_kwargs),
-    )
+    constraints = (FixAxisR(**kwargs), FixAxisZ(**kwargs), FixPsi(**kwargs))
 
     if profiles:
         for name, con in _PROFILE_CONSTRAINTS.items():
             if getattr(desc_eq, name) is not None:
-                constraints += (con(eq=desc_eq, **normalize_kwargs),)
-    constraints += (FixSheetCurrent(desc_eq),)
+                constraints += (con(**kwargs),)
+    constraints += (FixSheetCurrent(**kwargs),)
 
     if fix_lambda or (fix_lambda >= 0 and type(fix_lambda) is int):
         L_axis_constraints, _, _ = calc_zeroth_order_lambda(
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 698f397d89..763e6a4b9e 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -87,7 +87,6 @@ class FixParameters(_Objective):
         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,
@@ -95,6 +94,14 @@ class FixParameters(_Objective):
     name : str, optional
         Name of the objective function.
 
+    Examples
+    --------
+    .. code-block:: python
+
+        from desc.coils import FourierXYZCoil
+        from desc.grid import LinearGrid
+        import numpy as np
+
     """
 
     _scalar = False
@@ -110,8 +117,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
-        normalize=False,
-        normalize_target=False,
+        normalize=True,
+        normalize_target=True,
         name="Fixed parameters",
     ):
         self._params = params
@@ -738,6 +745,10 @@ class FixLambdaGauge(_Objective):
     ----------
     eq : Equilibrium
         Equilibrium that will be optimized to satisfy the Objective.
+    normalize : bool, optional
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     name : str, optional
         Name of the objective function.
 
@@ -752,6 +763,8 @@ class FixLambdaGauge(_Objective):
     def __init__(
         self,
         eq,
+        normalize=True,
+        normalize_target=True,
         name="lambda gauge",
     ):
         super().__init__(
@@ -759,6 +772,8 @@ def __init__(
             target=0,
             bounds=None,
             weight=1,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -832,6 +847,10 @@ class FixThetaSFL(FixParameters):
     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
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     name : str, optional
         Name of the objective function.
 
@@ -844,6 +863,8 @@ def __init__(
         self,
         eq,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         name="theta SFL",
     ):
         super().__init__(
@@ -852,6 +873,8 @@ def __init__(
             target=0,
             bounds=None,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -1198,6 +1221,10 @@ class FixModeLambda(FixParameters):
     weight : float, ndarray, optional
         Weighting to apply to the Objective, relative to other Objectives.
         Must be broadcastable to Objective.dim_f.
+    normalize : bool, optional
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     modes : ndarray, optional
         Basis modes numbers [l,m,n] of Fourier-Zernike modes to fix.
         len(target) = len(weight) = len(modes).
@@ -1216,6 +1243,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         modes=True,
         name="fix mode lambda",
     ):
@@ -1231,6 +1260,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -1584,6 +1615,10 @@ class FixSumModesLambda(_FixedObjective):
     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
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     sum_weight : float, ndarray, optional
         Weights on the coefficients in the sum, should be same length as modes.
         Defaults to 1 i.e. target = 1*L_111 + 1*L_222...
@@ -1609,6 +1644,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         sum_weights=None,
         modes=True,
         name="Fix Sum Modes lambda",
@@ -1640,6 +1677,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -1895,6 +1934,10 @@ class FixIota(FixParameters):
     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
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices.
@@ -1915,6 +1958,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         indices=True,
         name="fixed iota",
     ):
@@ -1924,6 +1969,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2294,6 +2341,10 @@ class FixAtomicNumber(FixParameters):
     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
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -2314,6 +2365,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         indices=True,
         name="fixed atomic number",
     ):
@@ -2323,6 +2376,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2485,6 +2540,10 @@ class FixCurveRotation(FixParameters):
     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
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     name : str, optional
         Name of the objective function.
 
@@ -2499,6 +2558,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         name="fixed rotation",
     ):
         super().__init__(
@@ -2507,6 +2568,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2579,6 +2642,10 @@ class FixOmniMap(FixParameters):
         Lower and upper bounds on the objective. Overrides target.
     weight : float, optional
         Weighting to apply to the Objective, relative to other Objectives.
+    normalize : bool, optional
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     indices : ndarray or bool, optional
         indices of the field.x_lmn array to fix.
         Must have len(target) = len(weight) = len(indices).
@@ -2597,6 +2664,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         indices=True,
         name="fixed omnigenity map",
     ):
@@ -2606,6 +2675,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 
@@ -2623,6 +2694,10 @@ class FixOmniBmax(_FixedObjective):
         Lower and upper bounds on the objective. Overrides target.
     weight : float, optional
         Weighting to apply to the Objective, relative to other Objectives.
+    normalize : bool, optional
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     name : str
         Name of the objective function.
 
@@ -2638,6 +2713,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         name="fixed omnigenity B_max",
     ):
         self._target_from_user = setdefault(bounds, target)
@@ -2646,6 +2723,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 

From 0615d563224e17baa4ac3aa43730d5638d5ddad2 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 14:58:46 -0500
Subject: [PATCH 44/50] update 2nd stage opt test with note

---
 tests/test_examples.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index c19752c903..7a334d2aa1 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1279,9 +1279,8 @@ def test_quadratic_flux_optimization_with_analytic_field():
 def test_second_stage_optimization():
     """Test optimizing magnetic field for a fixed axisymmetric equilibrium."""
     eq = get("DSHAPE")
-    grid = LinearGrid(rho=np.array([1.0]), M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
     field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
-    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field, eval_grid=grid))
+    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
     constraints = FixParameters(field, [{"R0": True}, {}])
     optimizer = Optimizer("scipy-trf")
     (field,), _ = optimizer.optimize(
@@ -1289,4 +1288,5 @@ def test_second_stage_optimization():
     )
     np.testing.assert_allclose(field[0].R0, 3.5)  # this value was fixed
     np.testing.assert_allclose(field[0].B0, 1)  # toroidal field (does not change)
-    np.testing.assert_allclose(field[1].B0, -0.022, rtol=1e-2)  # vertical field
+    # FIXME: fix QuadraticFlux objective so this works
+    # np.testing.assert_allclose(field[1].B0, 0)  vertical field (should vanish)

From c690222c9452c87e11b1d86a997421f32001743a Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 15:07:40 -0500
Subject: [PATCH 45/50] update FixParameters example

---
 desc/objectives/linear_objectives.py | 32 ++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 763e6a4b9e..cdf9b2e15c 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -98,9 +98,37 @@ class FixParameters(_Objective):
     --------
     .. code-block:: python
 
-        from desc.coils import FourierXYZCoil
-        from desc.grid import LinearGrid
         import numpy as np
+        from desc.coils import (
+            CoilSet, FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, MixedCoilSet
+        )
+        from desc.objectives import FixParameters
+
+        # toroidal field coil set with 3 coils
+        tf_coil = FourierPlanarCoil(center=[2, 0, 0], normal=[0, 1, 0], r_n=[1])
+        tf_coilset = CoilSet.linspaced_angular(tf_coil, n=3)
+        # vertical field coil set with 2 coils
+        vf_coil = FourierRZCoil(R_n=3, Z_n=-1)
+        vf_coilset = CoilSet.linspaced_linear(
+            vf_coil, displacement=[0, 0, 2], n=2, endpoint=True
+        )
+        # another single coil
+        coil = FourierXYZCoil()
+        # full coil set with TF coils, VF coils, and other single coil
+        full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xy_coil))
+
+        params = [
+            [
+                {"current": True},  # fix "current" in 1st TF coil
+                # fix "center" and one component of "normal" for 2nd TF coil
+                {"center": True, "normal": np.array([1])},
+                {},  # fix nothing in 3rd TF coil
+            ],
+            {"shift": True, "rotmat": True},  # fix "shift" & "rotmat" for all VF coils
+            # fix specified indices of "X_n" and "Z_n", but not "Y_n", for other coil
+            {"X_n": np.array([1, 2]), "Y_n": False, "Z_n": np.array([0])},
+        ]
+        obj = FixParameters(full_coilset, params)
 
     """
 

From 022861d49c5ef48fe16ca50180c057090a7e6a3a Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 16:04:54 -0500
Subject: [PATCH 46/50] I hate debugging code

---
 desc/objectives/getters.py           | 4 ++--
 desc/objectives/linear_objectives.py | 8 ++++++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py
index 756a1b6acf..9a7c980e32 100644
--- a/desc/objectives/getters.py
+++ b/desc/objectives/getters.py
@@ -210,11 +210,11 @@ def maybe_add_self_consistency(thing, constraints):
         constraints += (BoundaryZSelfConsistency(eq=thing),)
     if {"L_lmn"} <= params and not is_any_instance(constraints, FixLambdaGauge):
         constraints += (FixLambdaGauge(eq=thing),)
-    if {"R_lmn", "Ra_lmn"} <= params and not is_any_instance(
+    if {"R_lmn", "Ra_n"} <= params and not is_any_instance(
         constraints, AxisRSelfConsistency
     ):
         constraints += (AxisRSelfConsistency(eq=thing),)
-    if {"Z_lmn", "Za_lmn"} <= params and not is_any_instance(
+    if {"Z_lmn", "Za_n"} <= params and not is_any_instance(
         constraints, AxisZSelfConsistency
     ):
         constraints += (AxisZSelfConsistency(eq=thing),)
diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index cdf9b2e15c..678883b9d1 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -1893,6 +1893,10 @@ class FixAnisotropy(FixParameters):
     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
+        Has no effect for this objective.
+    normalize_target : bool, optional
+        Has no effect for this objective.
     indices : ndarray or bool, optional
         indices of the Profile.params array to fix.
         (e.g. indices corresponding to modes for a PowerSeriesProfile or indices
@@ -1913,6 +1917,8 @@ def __init__(
         target=None,
         bounds=None,
         weight=1,
+        normalize=True,
+        normalize_target=True,
         indices=True,
         name="fixed anisotropy",
     ):
@@ -1922,6 +1928,8 @@ def __init__(
             target=target,
             bounds=bounds,
             weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
             name=name,
         )
 

From 8f9acfca58ad4e27f27a65d01fd7218dfb2b3e09 Mon Sep 17 00:00:00 2001
From: daniel-dudt <daniel.dudt@princetonstellarators.energy>
Date: Thu, 2 May 2024 21:11:10 -0500
Subject: [PATCH 47/50] set vacuum=True for second_stage test

---
 tests/test_examples.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index 7a334d2aa1..adad7d7104 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1280,13 +1280,17 @@ def test_second_stage_optimization():
     """Test optimizing magnetic field for a fixed axisymmetric equilibrium."""
     eq = get("DSHAPE")
     field = ToroidalMagneticField(B0=1, R0=3.5) + VerticalMagneticField(B0=1)
-    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field))
+    objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=field, vacuum=True))
     constraints = FixParameters(field, [{"R0": True}, {}])
     optimizer = Optimizer("scipy-trf")
     (field,), _ = optimizer.optimize(
-        things=field, objective=objective, constraints=constraints, verbose=2
+        things=field,
+        objective=objective,
+        constraints=constraints,
+        ftol=0,
+        xtol=0,
+        verbose=2,
     )
     np.testing.assert_allclose(field[0].R0, 3.5)  # this value was fixed
-    np.testing.assert_allclose(field[0].B0, 1)  # toroidal field (does not change)
-    # FIXME: fix QuadraticFlux objective so this works
-    # np.testing.assert_allclose(field[1].B0, 0)  vertical field (should vanish)
+    np.testing.assert_allclose(field[0].B0, 1)  # toroidal field (no change)
+    np.testing.assert_allclose(field[1].B0, 0, atol=1e-12)  # vertical field (vanishes)

From fee6990831b6c694a876ee17fbd4e04e34ff14ce Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 8 May 2024 10:44:49 -0400
Subject: [PATCH 48/50] remove old comment

---
 desc/objectives/linear_objectives.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 678883b9d1..3b2604afc3 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -63,7 +63,6 @@ def _parse_target_from_user(
         return target, bounds
 
 
-# Rename to `FixParameters` (plural) instead?
 class FixParameters(_Objective):
     """Fix specific degrees of freedom associated with a given Optimizable thing.
 

From d02532452a78bf5c5db1894045d761953abb8b26 Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 8 May 2024 10:52:56 -0400
Subject: [PATCH 49/50] add note about True to FixParameters

---
 desc/objectives/linear_objectives.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 3b2604afc3..4fac503790 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -72,7 +72,8 @@ class FixParameters(_Objective):
         Object whose degrees of freedom are being fixed.
     params : nested list of dicts
         Dict keys are the names of parameters to fix (str), and dict values are the
-        indices to fix for each corresonding parameter (int array).
+        indices to fix for each corresponding parameter (int array).
+        Use True instead of an int array to fix all indices for that parameter.
         Must have the same pytree structure as thing.params_dict.
         The default is to fix all indices of all parameters.
     target : dict of {float, ndarray}, optional

From 4015affcb35648cabdadd9c8de5f7a41f52cee06 Mon Sep 17 00:00:00 2001
From: Dario Panici <dpanici@princeton.edu>
Date: Wed, 8 May 2024 10:57:34 -0400
Subject: [PATCH 50/50] and note of False

---
 desc/objectives/linear_objectives.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py
index 4fac503790..ef761dfc36 100644
--- a/desc/objectives/linear_objectives.py
+++ b/desc/objectives/linear_objectives.py
@@ -73,7 +73,8 @@ class FixParameters(_Objective):
     params : nested list of dicts
         Dict keys are the names of parameters to fix (str), and dict values are the
         indices to fix for each corresponding parameter (int array).
-        Use True instead of an int array to fix all indices for that parameter.
+        Use True (False) instead of an int array to fix all (none) of the indices
+        for that parameter.
         Must have the same pytree structure as thing.params_dict.
         The default is to fix all indices of all parameters.
     target : dict of {float, ndarray}, optional