diff --git a/CHANGELOG.md b/CHANGELOG.md index e48e91e741..0945c872c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ Bug Fixes - Fixes a bug in `OmnigenousField.change_resolution` when changing `L_B`. - Scaling a `ScaledProfile` or taking power of a `PowerProfile` now only updates the `scale`/`power` attributes instead of nesting the `ScaledProfile`/`PowerProfile`s. - `jax.Array`s in `_static_attrs` will be automatically converted to `np.ndarray` to prevent stalling code. In general, jax arrays should be omitted in `_static_attrs`. +- Fixes a bug in `_CoilObjective` for objectives which are computed per-grid node when at least one entry of `weight` is zero. Performance Improvements @@ -35,6 +36,11 @@ Deprecations - `constants` argument of `compute`, `jvp`, `jac`, `grad` and `hess` methods (including all of their variants) to all objective classes (including `ObjectiveFunction` and wrappers) is deprecated and will be removed in a future release. This argument was not necessary, and the code will still work if user doesn't pass it. Users should update their custom objectives for this change. In addition, `constants` property of the `ObjectiveFunction` and all sub-classes of `_Objective` is deprecated. +Breaking Changes + +- Name change in `_CoilObjective` replacing `coilset_mask` with `objective_mask`. Custom + subclasses with `_broadcast_input="Node"` that previously used `coilset_mask` should + switch to `objective_mask`. v0.17.1 ------- @@ -73,6 +79,10 @@ or if multiple things are being optimized, `x_scale` can be a list of dict, one - Changes the import paths for ``desc.external`` to require reference to the sub-modules. - Adds a differentiable utility for finding constant offset toroidal surfaces inside of optimizations. See [PR](https://github.com/PlasmaControl/DESC/pull/2016) for more details. - Add support for Python 3.14 +- Adds support for optimization targeting individual coils in a coilset. + - Coil objectives accept pytree inputs for `target`, `bounds`, and `weight`. + - Able to set weights to zero, excluding certain coils from the objective. + Bug Fixes @@ -108,10 +118,6 @@ New Features - `field_line_integrate` function doesn't accept additional keyword-arguments related to `diffrax`, if it is necessary, they must be given through `options` dictionary. - ``poincare_plot`` and ``plot_field_lines`` functions can now plot partial results if the integration failed. Previously, user had to pass ``throw=False`` or change the integration parameters. Users can ignore the warnings that are caused by hitting the bounds (i.e. `Terminating differential equation solve because an event occurred.`). - `chunk_size` argument is now used for chunking the number of field lines. For the chunking of Biot-Savart integration for the magnetic field, users can use `bs_chunk_size` instead. -- Adds support for optimization targeting individual coils in a coilset. - - Coil objectives accept pytree inputs for `target`, `bounds`, and `weight`. - - Able to set weights to zero, excluding certain coils from the objective. - Bug Fixes diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index edfbdf8f2c..51e8666204 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -59,7 +59,7 @@ class _CoilObjective(_Objective): "Coil" if the objective returns a single scalar per coil, and "Node" if it returns a scalar at every grid point. To be compatible with masking, compute function should apply the mask - self._coilset_tree["coilset_mask"] before returning data. + self._coilset_tree["objective_mask"] before returning data. """ __doc__ = __doc__.rstrip() + collect_docs(coil=True) @@ -133,7 +133,9 @@ def _build_coilset_tree(): params_tree["coils"] contains a nested list of 0s representing individual coils and the coilsets to which they belong. Similarly, params_tree["nodes"] lists the grid nodes associated with each coil. - params_tree["coilset_mask"] contains the indices in [0,self._dim_f-1] + params_tree["coilset_mask"] contains the indices in + [0,self._num_coils-1] for which the corresponding weight is positive. + params_tree["objective_mask"] contains the indices in [0,self._dim_f-1] for which the corresponding weight is positive. If all weights are positive (i.e. no masking needed), contains default slice(None). """ @@ -166,12 +168,16 @@ def expand(t, idx=0): self._coilset_tree = { "coils": tree[0], "nodes": tree[1], - "coilset_mask": slice(None), + "coilset_mask": np.arange(self._num_coils), + "objective_mask": slice(None), } if np.any([w == 0 for w in tree_leaves(self._weight)]): - mask = self._coilset_broadcast(self._weight) - mask = np.nonzero(mask)[0] - self._coilset_tree["coilset_mask"] = mask + coilset_mask = self._coilset_broadcast(self._weight) + objective_mask = self._coilset_broadcast( + self._weight, self._broadcast_input + ) + self._coilset_tree["coilset_mask"] = np.nonzero(coilset_mask)[0] + self._coilset_tree["objective_mask"] = np.nonzero(objective_mask)[0] coil = self.things[0] grid = self._grid @@ -211,12 +217,12 @@ def expand(t, idx=0): _build_coilset_tree() quad_weights = np.concatenate([g.spacing[:, 2] for g in grid])[ - self._coilset_tree["coilset_mask"] + self._coilset_tree["objective_mask"] ] if self._broadcast_input == "Node": grid_nodes_unmasked = [ - g.num_nodes for g in grid[self._coilset_tree["coilset_mask"]] + grid[i].num_nodes for i in self._coilset_tree["coilset_mask"] ] self._dim_f = np.sum(grid_nodes_unmasked) else: @@ -230,14 +236,14 @@ def expand(t, idx=0): grid = _prune_coilset_tree(grid) coil = _prune_coilset_tree(coil) - self._weight = self._coilset_broadcast(self._weight) + self._weight = self._coilset_broadcast(self._weight, self._broadcast_input) if self._bounds: self._bounds = ( - self._coilset_broadcast(self._bounds[0]), - self._coilset_broadcast(self._bounds[1]), + self._coilset_broadcast(self._bounds[0], self._broadcast_input), + self._coilset_broadcast(self._bounds[1], self._broadcast_input), ) elif self._target: - self._target = self._coilset_broadcast(self._target) + self._target = self._coilset_broadcast(self._target, self._broadcast_input) timer = Timer() if verbose > 0: @@ -294,14 +300,18 @@ def bounds(self, bounds): assert (bounds is None) or (isinstance(bounds, tuple) and len(bounds) == 2) if bounds: self._bounds = ( - self._coilset_broadcast(bounds[0]), - self._coilset_broadcast(bounds[1]), + self._coilset_broadcast(bounds[0], self._broadcast_input), + self._coilset_broadcast(bounds[1], self._broadcast_input), ) self._check_dimensions() @_Objective.target.setter def target(self, target): - self._target = self._coilset_broadcast(target) if target is not None else target + self._target = ( + self._coilset_broadcast(target, self._broadcast_input) + if target is not None + else target + ) self._check_dimensions() @_Objective.weight.setter @@ -311,13 +321,15 @@ def weight(self, weight): # objective should be rebuilt to account for masking self._built = False - def _coilset_broadcast(self, x): - """Expand an array in accordance with the attribute _broadcast_input. + def _coilset_broadcast(self, x, target="Coil"): + """Broadcast an array to dimensions consistent with "target". Parameters ---------- x : float or list[float] Must be broadcastable to the structure of self._things[0]. + target: str, optional + Optional string taking values "Coil" or "Node". Defaults to "Coil". Returns ------- @@ -325,16 +337,18 @@ def _coilset_broadcast(self, x): Float inputs are returned unchanged, and list inputs are expanded to size self._dim_f. """ + assert target in ["Node", "Coil"] + # No need to broadcast if input is a scalar arr_flat = tree_leaves(x) if len(arr_flat) == 1: return np.atleast_1d(arr_flat[0]) arr = jax_tree_broadcast(x, self._coilset_tree["coils"]) - if self._broadcast_input == "Node": + if target == "Node": arr = tree_map(lambda a, b: [a] * b, arr, self._coilset_tree["nodes"]) arr, _ = tree_flatten(arr) - return np.asarray(arr)[self._coilset_tree["coilset_mask"]] + return np.asarray(arr)[self._coilset_tree["objective_mask"]] class CoilLength(_CoilObjective): @@ -433,7 +447,7 @@ def compute(self, params, constants=None): data = super().compute(params, constants=constants) data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict)) out = jnp.array([dat["length"] for dat in data]) - return out[self._coilset_tree["coilset_mask"]] + return out[self._coilset_tree["objective_mask"]] class CoilCurvature(_CoilObjective): @@ -535,7 +549,7 @@ def compute(self, params, constants=None): data = super().compute(params, constants=constants) data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict)) out = jnp.concatenate([dat["curvature"] for dat in data]) - return out[self._coilset_tree["coilset_mask"]] + return out[self._coilset_tree["objective_mask"]] class CoilTorsion(_CoilObjective): @@ -635,7 +649,7 @@ def compute(self, params, constants=None): data = super().compute(params, constants=constants) data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict)) out = jnp.concatenate([dat["torsion"] for dat in data]) - return out[self._coilset_tree["coilset_mask"]] + return out[self._coilset_tree["objective_mask"]] class CoilCurrentLength(CoilLength): @@ -741,7 +755,7 @@ def compute(self, params, constants=None): lengths = super().compute(params, constants=constants) params = tree_leaves(params, is_leaf=lambda x: isinstance(x, dict)) currents = jnp.concatenate([param["current"] for param in params]) - out = jnp.atleast_1d(lengths * currents[self._coilset_tree["coilset_mask"]]) + out = jnp.atleast_1d(lengths * currents[self._coilset_tree["objective_mask"]]) return out @@ -848,7 +862,7 @@ def compute(self, params, constants=None): for dat in data ] ) - return out[self._coilset_tree["coilset_mask"]] + return out[self._coilset_tree["objective_mask"]] class CoilSetMinDistance(_Objective): @@ -1566,7 +1580,7 @@ def compute(self, params, constants=None): constants = self._get_deprecated_constants(constants) data = tree_leaves(data, is_leaf=lambda x: isinstance(x, dict)) out = jnp.array([jnp.var(jnp.linalg.norm(dat["x_s"], axis=1)) for dat in data]) - return (out * constants["mask"])[self._coilset_tree["coilset_mask"]] + return (out * constants["mask"])[self._coilset_tree["objective_mask"]] class QuadraticFlux(_Objective): diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 5403aef6ba..717dc7e508 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -4571,6 +4571,7 @@ def test_deflation_operator_all_Nones(): @pytest.mark.unit def test_coil_objective_input(DummyMixedCoilSet): """Tests broadcasting for inputs to _CoilObjectives.""" + # Consists of [coilset with 1 coil, coilset with 3 coils, single coil, single coil] coilset = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5") weight = [1.0, 2.0, 3.0, 4.0] @@ -4602,7 +4603,9 @@ def test_coil_objective_input(DummyMixedCoilSet): @pytest.mark.unit def test_coil_objective_indices(DummyMixedCoilSet): """Tests that setting "weights" to zero correctly masks _CoilObjectives errors.""" + # Consists of [coilset with 1 coil, coilset with 3 coils, single coil, single coil] coilsetA = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5") + # Consists of [single coil, single coil, single coil] coilsetB = MixedCoilSet((coilsetA[1][1], coilsetA[2], coilsetA[3])) weight = [0, [1, 0, 0], 1, 1] @@ -4613,32 +4616,75 @@ def test_coil_objective_indices(DummyMixedCoilSet): compA = objA.compute_scaled_error(None) compB = objB.compute_scaled_error(None) np.testing.assert_allclose(compA, compB, atol=1e-13) + assert objA.dim_f == 3 + + objC = CoilCurvature(coilsetA, weight=weight, normalize=False) + objD = CoilCurvature(coilsetB, normalize=False) + objC.build() + objD.build() + compC = objC.compute_scaled_error(None) + compD = objD.compute_scaled_error(None) + np.testing.assert_allclose(compC, compD, atol=1e-13) + assert objC.dim_f == objD.dim_f + + objE = CoilTorsion(coilsetA, weight=weight, normalize=False) + objF = CoilTorsion(coilsetB, normalize=False) + objE.build() + objF.build() + compE = objE.compute_scaled_error(None) + compF = objF.compute_scaled_error(None) + np.testing.assert_allclose(compE, compF, atol=1e-13) + assert objE.dim_f == objF.dim_f @pytest.mark.unit def test_coil_objective_setter(DummyMixedCoilSet): """Tests setters for _CoilObjectives.""" - coilset = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5") - obj = CoilLength(coilset) + # Consists of [coilset with 1 coil, coilset with 3 coils, single coil, single coil] + coilsetA = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5") + # Consists of [coilset with 1 coil, coilset with 3 coils] + coilsetB = MixedCoilSet((coilsetA[0], coilsetA[1]), check_intersection=False) - weight = [1.0, 1.0, 0.0, 0.0] - weight_expanded = [1.0, 1.0, 1.0, 1.0] - obj.weight = weight - obj.build() - np.testing.assert_allclose(obj.weight, weight_expanded, atol=1e-13) + objA = CoilLength(coilsetA) + objB = CoilLength(coilsetB) + + weightA = [1.0, 2.0, 0.0, 0.0] + weightB = [1.0, 2.0] + weight_expanded = [1.0, 2.0, 2.0, 2.0] + objA.weight = weightA + objB.weight = weightB + objA.build() + objB.build() + np.testing.assert_allclose(objA.weight, objB.weight, atol=1e-13) + np.testing.assert_allclose(objA.weight, weight_expanded, atol=1e-13) bounds = ([0.0, 1.0, 2.0, 3.0], [[4.0], [5.0, 6.0, 7.0], 8.0, 9.0]) bounds_expanded = ([0.0, 1.0, 1.0, 1.0], [4.0, 5.0, 6.0, 7.0]) target = [[4.0], 5.0, 8.0, 9.0] target_expanded = [4.0, 5.0, 5.0, 5.0] - obj.bounds = bounds - np.testing.assert_allclose(obj.bounds[0], bounds_expanded[0], atol=1e-13) - np.testing.assert_allclose(obj.bounds[1], bounds_expanded[1], atol=1e-13) - - obj.bounds = None - obj.target = target - np.testing.assert_allclose(obj.target, target_expanded, atol=1e-13) + objA.bounds = bounds + np.testing.assert_allclose(objA.bounds[0], bounds_expanded[0], atol=1e-13) + np.testing.assert_allclose(objA.bounds[1], bounds_expanded[1], atol=1e-13) + + objA.bounds = None + objA.target = target + np.testing.assert_allclose(objA.target, target_expanded, atol=1e-13) + + objC = CoilCurvature(coilsetA) + objD = CoilCurvature(coilsetB) + + objC.weight = weightA + objD.weight = weightB + objC.build() + objD.build() + objC.bounds = bounds + objC.target = target + num_nodes = len(objD.weight) + np.testing.assert_allclose(objC.weight, objD.weight, atol=1e-13) + assert len(objC.weight) == num_nodes + assert len(objC.bounds[0]) == num_nodes + assert len(objC.target) == num_nodes @pytest.mark.unit