change chunk_size to jac_chunk_size, change docstring

PlasmaControl · Sep 9, 2024 · f3b5f25 · f3b5f25
1 parent 6f2bbdb
commit f3b5f25
Show file tree

Hide file tree

Showing 17 changed files with 722 additions and 684 deletions.
diff --git a/desc/continuation.py b/desc/continuation.py
@@ -29,7 +29,7 @@ def _solve_axisym(
     maxiter=100,
     verbose=1,
     checkpoint_path=None,
-    chunk_size=None,
+    jac_chunk_size=None,
 ):
     """Solve initial axisymmetric case with adaptive step sizing."""
     timer = Timer()
@@ -101,7 +101,7 @@ def _solve_axisym(
 
         constraints_i = get_fixed_boundary_constraints(eq=eqi)
         objective_i = get_equilibrium_objective(
-            eq=eqi, mode=objective, chunk_size=chunk_size
+            eq=eqi, mode=objective, jac_chunk_size=jac_chunk_size
         )
 
         if verbose:
@@ -199,7 +199,7 @@ def _add_pressure(
     maxiter=100,
     verbose=1,
     checkpoint_path=None,
-    chunk_size=None,
+    jac_chunk_size=None,
 ):
     """Add pressure with adaptive step sizing."""
     timer = Timer()
@@ -229,7 +229,7 @@ def _add_pressure(
 
         constraints_i = get_fixed_boundary_constraints(eq=eqi)
         objective_i = get_equilibrium_objective(
-            eq=eqi, mode=objective, chunk_size=chunk_size
+            eq=eqi, mode=objective, jac_chunk_size=jac_chunk_size
         )
 
         if verbose:
@@ -330,7 +330,7 @@ def _add_shaping(
     maxiter=100,
     verbose=1,
     checkpoint_path=None,
-    chunk_size=None,
+    jac_chunk_size=None,
 ):
     """Add 3D shaping with adaptive step sizing."""
     timer = Timer()
@@ -361,7 +361,7 @@ def _add_shaping(
 
         constraints_i = get_fixed_boundary_constraints(eq=eqi)
         objective_i = get_equilibrium_objective(
-            eq=eqi, mode=objective, chunk_size=chunk_size
+            eq=eqi, mode=objective, jac_chunk_size=jac_chunk_size
         )
 
         if verbose:
@@ -460,7 +460,7 @@ def solve_continuation_automatic(  # noqa: C901
     maxiter=100,
     verbose=1,
     checkpoint_path=None,
-    chunk_size=None,
+    jac_chunk_size=None,
     **kwargs,
 ):
     """Solve for an equilibrium using an automatic continuation method.
@@ -539,7 +539,7 @@ def solve_continuation_automatic(  # noqa: C901
         maxiter,
         verbose,
         checkpoint_path,
-        chunk_size=chunk_size,
+        jac_chunk_size=jac_chunk_size,
     )
 
     # for zero current we want to do shaping before pressure to avoid having a
@@ -558,7 +558,7 @@ def solve_continuation_automatic(  # noqa: C901
             maxiter,
             verbose,
             checkpoint_path,
-            chunk_size=chunk_size,
+            jac_chunk_size=jac_chunk_size,
         )
 
         eqfam = _add_pressure(
@@ -574,7 +574,7 @@ def solve_continuation_automatic(  # noqa: C901
             maxiter,
             verbose,
             checkpoint_path,
-            chunk_size=chunk_size,
+            jac_chunk_size=jac_chunk_size,
         )
 
     # for other cases such as fixed iota or nonzero current we do pressure first
@@ -593,7 +593,7 @@ def solve_continuation_automatic(  # noqa: C901
             maxiter,
             verbose,
             checkpoint_path,
-            chunk_size=chunk_size,
+            jac_chunk_size=jac_chunk_size,
         )
 
         eqfam = _add_shaping(
@@ -609,7 +609,7 @@ def solve_continuation_automatic(  # noqa: C901
             maxiter,
             verbose,
             checkpoint_path,
-            chunk_size=chunk_size,
+            jac_chunk_size=jac_chunk_size,
         )
     eq.params_dict = eqfam[-1].params_dict
     eqfam[-1] = eq
@@ -640,7 +640,7 @@ def solve_continuation(  # noqa: C901
     maxiter=100,
     verbose=1,
     checkpoint_path=None,
-    chunk_size=None,
+    jac_chunk_size=None,
 ):
     """Solve for an equilibrium by continuation method.
 
@@ -702,7 +702,7 @@ def solve_continuation(  # noqa: C901
     if not isinstance(optimizer, Optimizer):
         optimizer = Optimizer(optimizer)
     objective_i = get_equilibrium_objective(
-        eq=eqfam[0], mode=objective, chunk_size=chunk_size
+        eq=eqfam[0], mode=objective, jac_chunk_size=jac_chunk_size
     )
     constraints_i = get_fixed_boundary_constraints(eq=eqfam[0])
 
@@ -751,7 +751,7 @@ def solve_continuation(  # noqa: C901
             # TODO: pass Jx if available
             eqp = eqfam[ii - 1].copy()
             objective_i = get_equilibrium_objective(
-                eq=eqp, mode=objective, chunk_size=chunk_size
+                eq=eqp, mode=objective, jac_chunk_size=jac_chunk_size
             )
             constraints_i = get_fixed_boundary_constraints(eq=eqp)
             eqp.change_resolution(**eqi.resolution)
@@ -771,9 +771,8 @@ def solve_continuation(  # noqa: C901
             stop = True
 
         if not stop:
-            # TODO: add ability to rebind objectives
             objective_i = get_equilibrium_objective(
-                eq=eqi, mode=objective, chunk_size=chunk_size
+                eq=eqi, mode=objective, jac_chunk_size=jac_chunk_size
             )
             constraints_i = get_fixed_boundary_constraints(eq=eqi)
             eqi.solve(

diff --git a/desc/objectives/_bootstrap.py b/desc/objectives/_bootstrap.py
@@ -66,18 +66,18 @@ class BootstrapRedlConsistency(_Objective):
         or quasi-axisymmetry; set to +/-NFP for quasi-helical symmetry.
     name : str, optional
         Name of the objective function.
-    chunk_size : int, optional
-        If `"blocked"` deriv_mode is used in the ObjectiveFunction, will
-        calculate the Jacobian for this objective ``chunk_size`` columns at a time,
-        instead of all at once.  The memory usage of the Jacobian calculation is
-        linearly proportional to ``chunk_size``: the smaller the ``chunk_size``, the
-        less memory the Jacobian calculation will require (with some baseline memory
-        usage). The time to compute the Jacobian is roughly ``t ~1/chunk_size``
-        with some baseline time, so the larger the ``chunk_size``, the faster the
-        calculation takes, at the cost of requiring more memory. A ``chunk_size``
-        of 1 corresponds to the least memory intensive,  but slowest method of
-        calculating the Jacobian.
-        If None, it will default to the largest possible `chunk_size` i.e. ``dim_x``
+    jac_chunk_size : int, optional
+        Will calculate the Jacobian for this objective ``jac_chunk_size``
+        columns at a time, instead of all at once.  The memory usage of the
+        Jacobian calculation is roughly ``memory usage = m0 + m1*jac_chunk_size``:
+        the higher the chunk size, the less memory the Jacobian calculation
+        will require (with some baseline memory usage). The time to compute the
+        Jacobian is roughly ``t=t0 +t1/jac_chunk_size``, so the larger the
+        ``jac_chunk_size``, the faster the calculation takes, at the cost of
+        requiring more memory. A ``jac_chunk_size`` of 1 corresponds to the least
+        memory intensive,  but slowest method of calculating the Jacobian.
+        If None, it will default to the largest possible
+        `jac_chunk_size` i.e. ``np.ceil(dim_x/4)``
 
     """
 
@@ -98,7 +98,7 @@ def __init__(
         grid=None,
         helicity=(1, 0),
         name="Bootstrap current self-consistency (Redl)",
-        chunk_size=None,
+        jac_chunk_size=None,
     ):
         if target is None and bounds is None:
             target = 0
@@ -116,7 +116,7 @@ def __init__(
             loss_function=loss_function,
             deriv_mode=deriv_mode,
             name=name,
-            chunk_size=chunk_size,
+            jac_chunk_size=jac_chunk_size,
         )
 
     def build(self, use_jit=True, verbose=1):