Fix broadcasting for coil objectives

-Handles case when objective is computed per-grid node and at least one weight is zero.

Author: singh-jaydeep

.gitattributes

@@ -1,3 +1,4 @@
 *.ipynb -linguist-detectable
 documentation/* linguist-documentation
 desc/_version.py export-subst
+*.sh text eol=lf

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["coil_mask"] contains the indices in [0,self._num_coils]
+                for which the corresponding weight is positive. Similarly,
+                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),
+                "coil_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
+                coil_mask = self._coilset_broadcast(self._weight)
+                objective_mask = self._coilset_broadcast(
+                    self._weight, self._broadcast_input
+                )
+                self._coilset_tree["coil_mask"] = np.nonzero(coil_mask)[0]
+                self._coilset_tree["objective_mask"] = np.nonzero(objective_mask)[0]
 
         coil = self.things[0]
         grid = self._grid
@@ -211,33 +217,31 @@ 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["coil_mask"]
             ]
             self._dim_f = np.sum(grid_nodes_unmasked)
         else:
-            coils_unmasked = np.ones(self._num_coils)[
-                self._coilset_tree["coilset_mask"]
-            ]
+            coils_unmasked = np.ones(self._num_coils)[self._coilset_tree["coil_mask"]]
             self._dim_f = len(coils_unmasked)
 
         # map grid to the same structure as coil and then remove unnecessary members
         grid = tree_unflatten(structure, grid)
         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 +298,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 +319,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
         -------
@@ -331,10 +341,10 @@ def _coilset_broadcast(self, x):
             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 +443,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 +545,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 +645,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 +751,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 +858,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 +1576,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):

tests/test_objective_funs.py

@@ -4614,31 +4614,64 @@ def test_coil_objective_indices(DummyMixedCoilSet):
     compB = objB.compute_scaled_error(None)
     np.testing.assert_allclose(compA, compB, atol=1e-13)
 
+    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)
+
+    objE = CoilTorsion(coilsetA, weight=weight)
+    objE.build()
+
 
 @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)
+    coilsetA = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5")
+    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