Merge pull request #361 from hiddenSymmetries/ag/surface_dofs_fix

ag/surface_dofs_fix

src/simsopt/_core/derivative.py

@@ -173,22 +173,38 @@ def __call__(self, optim, as_derivative=False):
 
         Args:
             optim: An Optimizable object
+            as_derivative: if True return as a Derivative object. The keys
+                           of the returned Derivative dictionary can include `optim`, the ancestors of `optim`
+                           or Optimizables that share DOFs with ancestors of `optim`.  The values are numpy arrays
+                           corresponding to the derivatives with respect to all degrees of freedom, both free and fixed.
+                           If False, return as a numpy array.  The entries of the array correspond only to free
+                           DOFs, and fixed ones are removed out.
         """
         from .optimizable import Optimizable  # Import here to avoid circular import
         assert isinstance(optim, Optimizable)
         derivs = []
-        keys = []
-        for k in optim.unique_dof_lineage:
-            if np.any(k.dofs_free_status):
-                local_derivs = np.zeros(k.local_dof_size)
-                for opt in k.dofs.dep_opts():
-                    local_derivs += self.data[opt][opt.local_dofs_free_status]
-                    keys.append(opt)
-                derivs.append(local_derivs)
 
         if as_derivative:
+            keys = []
+
+            for k in optim.unique_dof_lineage:
+                for opt in k.dofs.dep_opts():
+                    # the next if-statament is there to avoid the dictionary from accumulating 
+                    # empty values e.g. if there are no local DOFs to opt, then self.data[opt] 
+                    # returns np.array([]).
+                    if opt.local_full_dof_size > 0:
+                        derivs.append(self.data[opt])
+                        keys.append(opt)
+
             return Derivative({k: d for k, d in zip(keys, derivs)})
+
         else:
+            for k in optim.unique_dof_lineage:
+                if np.any(k.dofs_free_status):
+                    local_derivs = np.zeros(k.local_dof_size)
+                    for opt in k.dofs.dep_opts():
+                        local_derivs += self.data[opt][opt.local_dofs_free_status]
+                    derivs.append(local_derivs)
             return np.concatenate(derivs)
 
     # https://stackoverflow.com/questions/11624955/avoiding-python-sum-default-start-arg-behavior

tests/core/test_optimizable.py

@@ -1422,6 +1422,35 @@ def test_derivative(self):
         self.assertTrue((sum_obj.dJ(partials=True)(adder_orig) == sum_obj.dJ()).all())
         self.assertTrue((sum_obj.dJ(partials=True)(adder_shared_dofs) == sum_obj.dJ()).all())
 
+    def test_as_derivative1(self):
+        # this test checks that you can restrict the Derivative dictionary
+        # to the proper subset of Optimizables when as_derivative=True
+
+        optA = OptClassSharedDOFs(x0=[1, 2, 3], names=["x", "y", "z"],
+                                        fixed=[False, False, True])
+        optA_shared_dofs = OptClassSharedDOFs(dofs=optA.dofs)
+
+        optB = OptClassSharedDOFs(x0=[np.pi, 1, 1.21], names=["xx", "yy", "zz"],
+                                        fixed=[False, False, True])
+        sum_opt = optA + optA_shared_dofs + optB
+        deriv = sum_opt.dJ(partials=True)(sum_opt, as_derivative=True)
+
+        # restrict to optA 
+        np.testing.assert_allclose(deriv(optA), optA.dJ()*2, atol=1e-14)
+        # restrict to optA_shared_dofs
+        np.testing.assert_allclose(deriv(optA_shared_dofs), optA.dJ()*2, atol=1e-14)
+        # restrict to sum_opt
+        np.testing.assert_allclose(deriv(sum_opt), np.concatenate((optA.dJ()*2, optB.dJ())), atol=1e-14)
+
+    def test_as_derivative2(self):
+        # this test checks that when you sum a Derivative dictionary generated using as_derivative=True,
+        # to another that things work as expected when some DOFs are fixed.
+
+        opt = OptClassSharedDOFs(x0=[1, 2, 3], names=["x", "y", "z"],
+                                        fixed=[False, False, True])
+        deriv = opt.dJ(partials=True)(opt, as_derivative=True) + opt.dJ(partials=True)
+        np.testing.assert_allclose(deriv(opt), opt.dJ()*2, atol=1e-14)
+
     def test_load_save(self):
         import tempfile
         from pathlib import Path

tests/geo/test_surface_objectives.py

@@ -337,18 +337,23 @@ def test_nonQSratio_derivative(self):
             for label in ["Volume", "ToroidalFlux"]:
                 for weight_inv_modB in [True, False]:
                     for optimize_G in [True, False]:
-                        if boozer_type == 'ls' and label == 'ToroidalFlux':
-                           continue
-                        if boozer_type == 'exact' and optimize_G is False:
-                            continue
-                        if boozer_type == 'exact' and weight_inv_modB:
-                            continue
-                        for axis in [False, True]:
-                            with self.subTest(label=label, axis=axis, boozer_type=boozer_type, optimize_G=optimize_G, weight_inv_modB=weight_inv_modB):
-                                self.subtest_nonQSratio_derivative(label, axis, boozer_type, optimize_G, weight_inv_modB)
+                        for fix_coil_dof in [True, False]:
+                            if boozer_type == 'ls' and label == 'ToroidalFlux':
+                               continue
+                            if boozer_type == 'exact' and optimize_G is False:
+                                continue
+                            if boozer_type == 'exact' and weight_inv_modB:
+                                continue
+                            for axis in [False, True]:
+                                with self.subTest(label=label, axis=axis, boozer_type=boozer_type, optimize_G=optimize_G, weight_inv_modB=weight_inv_modB, fix_coil_dof=fix_coil_dof):
+                                    self.subtest_nonQSratio_derivative(label, axis, boozer_type, optimize_G, weight_inv_modB, fix_coil_dof)
 
-    def subtest_nonQSratio_derivative(self, label, axis, boozer_type, optimize_G, weight_inv_modB):
+    def subtest_nonQSratio_derivative(self, label, axis, boozer_type, optimize_G, weight_inv_modB, fix_coil_dof):
         bs, boozer_surface = get_boozer_surface(label=label, boozer_type=boozer_type, optimize_G=optimize_G, weight_inv_modB=weight_inv_modB)
+
+        if fix_coil_dof:
+            bs.coils[0].curve.fix('xc(0)')
+
         coeffs = bs.x
         io = NonQuasiSymmetricRatio(boozer_surface, bs, quasi_poloidal=axis)