Change defaults for reverse mode jacobian chunk size for eps_eff

desc/objectives/_neoclassical.py

@@ -29,6 +29,20 @@
 from .objective_funs import _Objective, collect_docs
 from .utils import _parse_callable_target_bounds
 
+_bounce_overwrite = {
+    "deriv_mode": """
+    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.
+
+        Default is ``fwd``. If ``rev`` is chosen, then ``jac_chunk_size=1`` is chosen
+        by default. In ``rev`` mode, reducing the pitch angle parameter ``batch_size``
+        does not reduce memory, so it is recommended to retain the default for that.
+        """
+}
+
 
 class EffectiveRipple(_Objective):
     """The effective ripple is a proxy for neoclassical transport.
@@ -96,6 +110,7 @@ class EffectiveRipple(_Objective):
         bounds_default="``target=0``.",
         normalize_detail=" Note: Has no effect for this objective.",
         normalize_target_detail=" Note: Has no effect for this objective.",
+        overwrite=_bounce_overwrite,
     )
 
     _coordinates = "r"
@@ -111,7 +126,7 @@ def __init__(
         normalize=True,
         normalize_target=True,
         loss_function=None,
-        deriv_mode="auto",
+        deriv_mode="fwd",
         grid=None,
         name="Effective ripple",
         jac_chunk_size=None,
@@ -142,6 +157,10 @@ def __init__(
             "num_pitch": num_pitch,
             "batch_size": batch_size,
         }
+        if deriv_mode == "rev" and jac_chunk_size is None:
+            # Reverse mode is bottlenecked by coordinate mapping.
+            # Compute Jacobian one flux surface at a time.
+            jac_chunk_size = 1
 
         super().__init__(
             things=eq,
@@ -171,14 +190,19 @@ def build(self, use_jit=True, verbose=1):
         if self._grid is None:
             self._grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False)
         assert self._grid.can_fft
+        # Constants can only hold 1D arrays for some reason.
         self._constants["clebsch"] = FourierChebyshevSeries.nodes(
             self._X,
             self._Y,
             self._grid.compress(self._grid.nodes[:, 0]),
             domain=(0, 2 * np.pi),
+        ).ravel()
+        self._constants["fieldline quad x"], self._constants["fieldline quad w"] = (
+            leggauss(self._hyperparam["Y_B"] // 2)
+        )
+        self._constants["quad x"], self._constants["quad w"] = chebgauss2(
+            self._hyperparam.pop("num_quad")
         )
-        self._constants["fieldline_quad"] = leggauss(self._hyperparam["Y_B"] // 2)
-        self._constants["quad"] = chebgauss2(self._hyperparam.pop("num_quad"))
 
         self._dim_f = self._grid.num_rho
         self._target, self._bounds = _parse_callable_target_bounds(
@@ -239,13 +263,16 @@ def compute(self, params, constants=None):
                 self._X,
                 self._Y,
                 iota=constants["transforms"]["grid"].compress(data["iota"]),
-                clebsch=constants["clebsch"],
+                clebsch=constants["clebsch"].reshape(-1, 3),
                 # Pass in params so that root finding is done with the new
                 # perturbed λ coefficients and not the original equilibrium's.
                 params=params,
             ),
-            fieldline_quad=constants["fieldline_quad"],
-            quad=constants["quad"],
+            fieldline_quad=(
+                constants["fieldline quad x"],
+                constants["fieldline quad w"],
+            ),
+            quad=(constants["quad x"], constants["quad w"]),
             **self._hyperparam,
         )
         return constants["transforms"]["grid"].compress(data["effective ripple"])
@@ -313,10 +340,10 @@ class GammaC(_Objective):
 
         Note that Nemov's Γ_c converges to a finite nonzero value in the
         infinity limit of the number of toroidal transits.
-        Velasco's expression has a secular term that will drive the result
-        to zero as the number of toroidal transits increases unless the
-        secular term is averaged out from all the singular integrals.
-        Therefore, an optimization using Velasco's metric should be evaluated by
+        Velasco's expression has a secular term that may drive the result
+        to zero as the number of toroidal transits increases if the quadratures
+        are unable to average out the secular term from all the singular integrals.
+        So an optimization using Velasco's metric may need to be evaluated by
         measuring decrease in Γ_c at a fixed number of toroidal transits until
         unless an adaptive quadrature is used.
 
@@ -327,6 +354,7 @@ class GammaC(_Objective):
         bounds_default="``target=0``.",
         normalize_detail=" Note: Has no effect for this objective.",
         normalize_target_detail=" Note: Has no effect for this objective.",
+        overwrite=_bounce_overwrite,
     )
 
     _coordinates = "r"
@@ -342,7 +370,7 @@ def __init__(
         normalize=True,
         normalize_target=True,
         loss_function=None,
-        deriv_mode="auto",
+        deriv_mode="fwd",
         grid=None,
         name="Gamma_c",
         jac_chunk_size=None,
@@ -378,6 +406,10 @@ def __init__(
             self._key = "Gamma_c"
         else:
             self._key = "Gamma_c Velasco"
+        if deriv_mode == "rev" and jac_chunk_size is None:
+            # Reverse mode is bottlenecked by coordinate mapping.
+            # Compute Jacobian one flux surface at a time.
+            jac_chunk_size = 1
 
         super().__init__(
             things=eq,
@@ -407,19 +439,22 @@ def build(self, use_jit=True, verbose=1):
         if self._grid is None:
             self._grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False)
         assert self._grid.can_fft
+        # Constants can only hold 1D arrays for some reason.
         self._constants["clebsch"] = FourierChebyshevSeries.nodes(
             self._X,
             self._Y,
             self._grid.compress(self._grid.nodes[:, 0]),
             domain=(0, 2 * np.pi),
+        ).ravel()
+        self._constants["fieldline quad x"], self._constants["fieldline quad w"] = (
+            leggauss(self._hyperparam["Y_B"] // 2)
         )
-        self._constants["fieldline_quad"] = leggauss(self._hyperparam["Y_B"] // 2)
         num_quad = self._hyperparam.pop("num_quad")
-        self._constants["quad"] = get_quadrature(
+        self._constants["quad x"], self._constants["quad w"] = get_quadrature(
             leggauss(num_quad),
             (automorphism_sin, grad_automorphism_sin),
         )
-        self._constants["quad2"] = chebgauss2(num_quad)
+        self._constants["quad2 x"], self._constants["quad2 w"] = chebgauss2(num_quad)
 
         self._dim_f = self._grid.num_rho
         self._target, self._bounds = _parse_callable_target_bounds(
@@ -459,7 +494,7 @@ def compute(self, params, constants=None):
         if constants is None:
             constants = self.constants
         if "quad2" in constants:
-            self._hyperparam["quad2"] = constants["quad2"]
+            self._hyperparam["quad2"] = (constants["quad2 x"], constants["quad2 w"])
 
         eq = self.things[0]
         data = compute_fun(
@@ -478,13 +513,16 @@ def compute(self, params, constants=None):
                 self._X,
                 self._Y,
                 iota=constants["transforms"]["grid"].compress(data["iota"]),
-                clebsch=constants["clebsch"],
+                clebsch=constants["clebsch"].reshape(-1, 3),
                 # Pass in params so that root finding is done with the new
                 # perturbed λ coefficients and not the original equilibrium's.
                 params=params,
             ),
-            fieldline_quad=constants["fieldline_quad"],
-            quad=constants["quad"],
+            fieldline_quad=(
+                constants["fieldline quad x"],
+                constants["fieldline quad w"],
+            ),
+            quad=(constants["quad x"], constants["quad w"]),
             **self._hyperparam,
         )
         return constants["transforms"]["grid"].compress(data[self._key])

desc/objectives/objective_funs.py

@@ -66,9 +66,9 @@
 doc_deriv_mode = """
     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.
+        ``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.
 """
 doc_name = """
     name : str, optional

docs/notebooks/tutorials/EffectiveRipple.ipynb

@@ -385,9 +385,9 @@
     "            num_well=30 * 10,\n",
     "            num_quad=32,\n",
     "            num_pitch=45,\n",
-    "            # TODO: It seems batch_size has no effect on memory consumed by Jacobian in reverse mode.\n",
-    "            batch_size=1,\n",
-    "            deriv_mode=\"fwd\",\n",
+    "            # Uncomment to compute at only batch_size pitch angles at a time.\n",
+    "            # This will reduce peak memory by 2.5 GB.\n",
+    "            # batch_size=1,\n",
     "        ),\n",
     "        AspectRatio(eq1, bounds=(8, 11), weight=1e3),\n",
     "        GenericObjective(\n",