Allow Optimization with OptimizableCollections to Work

desc/magnetic_fields.py

@@ -478,6 +478,12 @@
             [type(field) for field in fields]
         )
         self._fields = fields
+        offset = jnp.concatenate(
+            [jnp.array([0]), jnp.cumsum(jnp.array([s.dim_x for s in self]))[:-1]]
+        )
+        # store split indices for unpacking
+        # as a numpy array to avoid jax issues later
+        self._split_idx = np.asarray(offset)
 
     def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None):
         """Compute magnetic field at a set of points.

desc/objectives/__init__.py

@@ -13,6 +13,7 @@
 from ._geometry import (
     AspectRatio,
     BScaleLength,
+    DummyFields,
     Elongation,
     GoodCoordinates,
     MeanCurvature,

desc/objectives/_geometry.py

@@ -1339,3 +1339,144 @@
         f = data["g_rr"]
 
         return jnp.concatenate([g, constants["sigma"] * f])
+
+
+class DummyFields(_Objective):
+    """Target a quantity at points from magnetic fields.
+
+    Parameters
+    ----------
+    fields : MagneticField
+        MagneticField object.
+    eq : Equilibrium
+        eq to calc coords to target Bphi over BZ from fields at.
+    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
+        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.
+    loss_function : {None, 'mean', 'min', 'max'}, optional
+        Loss function to apply to the objective values once computed. This loss function
+        is called on the raw compute value, before any shifting, scaling, or
+        normalization.
+    deriv_mode : {"auto", "fwd", "rev"}
+        Specify how to compute jacobian matrix, either forward mode or reverse mode AD.
+        "auto" selects forward or reverse mode based on the size of the input and output
+        of the objective. Has no effect on self.grad or self.hess which always use
+        reverse mode and forward over reverse mode respectively.
+    eval_grid : Grid, optional
+        Collocation grid containing the nodes to evaluate at.
+    source_grid : Grid, optional
+        Grid to discretize field objects with.
+    name : str, optional
+        Name of the objective function.
+
+    """
+
+    _coordinates = "rtz"
+    _units = "unitless"
+    _print_value_fmt = "Magnetic field quantity: {:10.3e} "
+
+    def __init__(
+        self,
+        field,
+        eq,
+        target=None,
+        bounds=None,
+        weight=1,
+        normalize=True,
+        normalize_target=True,
+        loss_function=None,
+        deriv_mode="auto",
+        eval_grid=None,
+        source_grid=None,
+        name="Bphi over BZ",
+    ):
+        if target is None and bounds is None:
+            bounds = (1, np.inf)
+        self._eval_grid = eval_grid
+        self._source_grid = source_grid
+        self._eq = eq
+
+        super().__init__(
+            things=field,
+            target=target,
+            bounds=bounds,
+            weight=weight,
+            normalize=normalize,
+            normalize_target=normalize_target,
+            loss_function=loss_function,
+            deriv_mode=deriv_mode,
+            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._eq
+        if self._eval_grid is None:
+            eval_grid = LinearGrid(M=eq.M * 2, N=eq.N * 2, NFP=eq.NFP)
+        else:
+            eval_grid = self._eval_grid
+
+        self._dim_f = eval_grid.num_nodes
+        self._data_keys = ["R", "phi", "Z", "e^theta", "e^zeta"]
+
+        timer = Timer()
+        if verbose > 0:
+            print("Precomputing transforms")
+        timer.start("Precomputing transforms")
+        data = eq.compute(self._data_keys, grid=eval_grid)
+        self._constants = {
+            "coords_rpz": np.vstack([data["R"], data["phi"], data["Z"]]).T,
+            "quad_weights": eval_grid.weights,
+        }
+
+        timer.stop("Precomputing transforms")
+        if verbose > 1:
+            timer.disp("Precomputing transforms")
+
+        super().build(use_jit=use_jit, verbose=verbose)
+
+    def compute(self, params, constants=None):
+        """Compute magnetic field scale length.
+
+        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
+            quantity.
+
+        """
+        if constants is None:
+            constants = self.constants
+        B = self.things[0].compute_magnetic_field(
+            constants["coords_rpz"], basis="rpz", params=params
+        )
+        # return Bphi over BZ (some random quantity we can directly control)
+        return B[:, 1] / B[:, 2]