Update rules for chunking jacobian (#1283)

Resolves #1276

desc/objectives/getters.py

@@ -1,6 +1,6 @@
 """Utilities for getting standard groups of objectives and constraints."""
 
-from desc.utils import flatten_list, is_any_instance, unique_list
+from desc.utils import flatten_list, is_any_instance, isposint, unique_list
 
 from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance
 from .linear_objectives import (
@@ -86,7 +86,10 @@ def get_equilibrium_objective(eq, mode="force", normalize=True, jac_chunk_size="
         objectives = (RadialForceBalance(**kwargs), HelicalForceBalance(**kwargs))
     else:
         raise ValueError("got an unknown equilibrium objective type '{}'".format(mode))
-    return ObjectiveFunction(objectives, jac_chunk_size=jac_chunk_size)
+    deriv_mode = "batched" if isposint(jac_chunk_size) else "auto"
+    return ObjectiveFunction(
+        objectives, jac_chunk_size=jac_chunk_size, deriv_mode=deriv_mode
+    )
 
 
 def get_fixed_axis_constraints(eq, profiles=True, normalize=True):

desc/objectives/objective_funs.py

@@ -260,14 +260,6 @@ def __init__(
         self._compiled = False
         self._name = name
 
-    def _set_derivatives(self):
-        """Choose derivative mode based on mode of sub-objectives."""
-        if self._deriv_mode == "auto":
-            if all((obj._deriv_mode == "fwd") for obj in self.objectives):
-                self._deriv_mode = "batched"
-            else:
-                self._deriv_mode = "blocked"
-
     def _unjit(self):
         """Remove jit compiled methods."""
         methods = [
@@ -325,35 +317,34 @@ def build(self, use_jit=None, verbose=1):
         else:
             self._scalar = False
 
-        self._set_derivatives()
+        self._set_things()
+
+        # setting derivative mode and chunking.
+        errorif(
+            isposint(self._jac_chunk_size) and self._deriv_mode in ["auto", "blocked"],
+            ValueError,
+            "'jac_chunk_size' was passed into ObjectiveFunction, but the "
+            "ObjectiveFunction is not using 'batched' deriv_mode",
+        )
         sub_obj_jac_chunk_sizes_are_ints = [
             isposint(obj._jac_chunk_size) for obj in self.objectives
         ]
         errorif(
-            any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode != "blocked",
+            any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode == "batched",
             ValueError,
             "'jac_chunk_size' was passed into one or more sub-objectives, but the"
-            " ObjectiveFunction is  using 'batched' deriv_mode, so sub-objective "
+            " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective "
             "'jac_chunk_size' will be ignored in favor of the ObjectiveFunction's "
             f"'jac_chunk_size' of {self._jac_chunk_size}."
             " Specify 'blocked' deriv_mode if each sub-objective is desired to have a "
             "different 'jac_chunk_size' for its Jacobian computation.",
         )
-        errorif(
-            self._jac_chunk_size not in ["auto", None]
-            and self._deriv_mode == "blocked",
-            ValueError,
-            "'jac_chunk_size' was passed into ObjectiveFunction, but the "
-            "ObjectiveFunction is using 'blocked' deriv_mode, so sub-objective "
-            "'jac_chunk_size' are used to compute each sub-objective's Jacobian, "
-            "`ignoring the ObjectiveFunction's 'jac_chunk_size'.",
-        )
 
-        if not self.use_jit:
-            self._unjit()
-
-        self._set_things()
-        self._built = True
+        if self._deriv_mode == "auto":
+            if all((obj._deriv_mode == "fwd") for obj in self.objectives):
+                self._deriv_mode = "batched"
+            else:
+                self._deriv_mode = "blocked"
 
         if self._jac_chunk_size == "auto":
             # Heuristic estimates of fwd mode Jacobian memory usage,
@@ -365,6 +356,15 @@ def build(self, use_jit=None, verbose=1):
                 * self.dim_x
             )
             self._jac_chunk_size = max([1, max_chunk_size])
+            if self._deriv_mode == "blocked":
+                for obj in self.objectives:
+                    if obj._jac_chunk_size is None:
+                        obj._jac_chunk_size = self._jac_chunk_size
+
+        if not self.use_jit:
+            self._unjit()
+
+        self._built = True
 
         timer.stop("Objective build")
         if verbose > 1:
@@ -946,7 +946,7 @@ def dim_x(self):
     @property
     def dim_f(self):
         """int: Number of objective equations."""
-        if not self.built:
+        if not hasattr(self, "_dim_f"):
             raise RuntimeError("ObjectiveFunction must be built first.")
         return self._dim_f
 

tests/test_examples.py

@@ -1043,6 +1043,7 @@ def test_freeb_vacuum():
     objective = ObjectiveFunction(
         VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True),
         jac_chunk_size=1000,
+        deriv_mode="batched",
     )
     eq, _ = eq.optimize(
         objective,

tests/test_objective_funs.py

@@ -1323,6 +1323,7 @@ def test_derivative_modes():
             AspectRatio(eq),
         ],
         deriv_mode="batched",
+        jac_chunk_size="auto",
         use_jit=False,
     )
     obj2 = ObjectiveFunction(
@@ -1332,13 +1333,15 @@ def test_derivative_modes():
             AspectRatio(eq, jac_chunk_size=None),
         ],
         deriv_mode="blocked",
+        jac_chunk_size="auto",
         use_jit=False,
     )
     obj1.build()
     obj2.build()
     # check that default size works for blocked
-    assert obj2.objectives[1]._jac_chunk_size is None
-    assert obj2.objectives[2]._jac_chunk_size is None
+    assert obj2.objectives[0]._jac_chunk_size == 2
+    assert obj2.objectives[1]._jac_chunk_size > 0
+    assert obj2.objectives[2]._jac_chunk_size > 0
     # hard to say what size auto will give, just check it is >0
     assert obj1._jac_chunk_size > 0
     obj3.build()