Add surf_batch_size to Boozer transform to control memory usage

CHANGELOG.md

@@ -4,6 +4,8 @@ Changelog
 Performance Improvements
 
 - Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep. On ``jax >= 0.10.0`` this uses ``qr_multiply`` to additionally avoid forming ``Q`` explicitly; on older versions a fallback preserves the same results.
+- Adds ``surf_batch_size`` kwarg to Boozer and omnigenous field compute variables, to allow for tuning of the memory usage when computing these quantities by choosing how many surfaces to simultaneously compute.
+  - Also adds ``surf_batch_size`` as an additional kwarg to ``make_boozmn_output``, as well as the objectives ``Omnigenity`` and ``QuasisymmetryBoozer``
 
 Bug Fixes
 

desc/compute/_omnigenity.py

@@ -14,6 +14,7 @@
 from interpax import interp1d
 
 from desc.backend import jnp, sign, vmap
+from desc.batching import vmap_chunked
 
 from ..utils import cross, dot, safediv
 from .data_index import register_compute_fun
@@ -36,14 +37,17 @@
     grid_requirement={"is_meshgrid": True, "sym": False},
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to ``grid.num_rho`` e.g. compute all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _B_theta_mn(params, transforms, profiles, data, **kwargs):
     B_theta = transforms["grid"].meshgrid_reshape(data["B_theta"], "rtz")
 
     def fitfun(x):
         return transforms["B"].fit(x.flatten(order="F"))
 
-    B_theta_mn = vmap(fitfun)(B_theta)
+    B_theta_mn = vmap_chunked(fitfun, chunk_size=kwargs.get("surf_batch_size"))(B_theta)
     # modes stored as shape(rho, mn) flattened
     data["B_theta_mn"] = B_theta_mn.flatten()
     return data
@@ -68,14 +72,17 @@
     aliases="B_zeta_mn",  # TODO(#568): remove when phi != zeta
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to ``grid.num_rho`` e.g. compute all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _B_phi_mn(params, transforms, profiles, data, **kwargs):
     B_phi = transforms["grid"].meshgrid_reshape(data["B_phi|r,t"], "rtz")
 
     def fitfun(x):
         return transforms["B"].fit(x.flatten(order="F"))
 
-    B_zeta_mn = vmap(fitfun)(B_phi)
+    B_zeta_mn = vmap_chunked(fitfun, chunk_size=kwargs.get("surf_batch_size"))(B_phi)
     # modes stored as shape(rho, mn) flattened
     data["B_phi_mn"] = B_zeta_mn.flatten()
     return data
@@ -141,11 +148,18 @@
     grid_requirement={"is_meshgrid": True, "sym": False},
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to ``grid.num_rho`` e.g. compute all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _w(params, transforms, profiles, data, **kwargs):
     grid = transforms["grid"]
     w_mn = data["w_Boozer_mn"].reshape((grid.num_rho, -1))
-    w = vmap(transforms["w"].transform)(w_mn)  # shape(rho, theta*zeta)
+    w = vmap_chunked(
+        transforms["w"].transform, chunk_size=kwargs.get("surf_batch_size")
+    )(
+        w_mn
+    )  # shape(rho, theta*zeta)
     w = w.reshape((grid.num_rho, grid.num_theta, grid.num_zeta), order="F")
     w = jnp.moveaxis(w, 0, 1)
     data["w_Boozer"] = w.flatten(order="F")
@@ -169,13 +183,18 @@
     grid_requirement={"is_meshgrid": True, "sym": False},
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to ``grid.num_rho`` e.g. compute all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _w_t(params, transforms, profiles, data, **kwargs):
     grid = transforms["grid"]
     w_mn = data["w_Boozer_mn"].reshape((grid.num_rho, -1))
     # need to close over dt which can't be vmapped
     fun = lambda x: transforms["w"].transform(x, dt=1)
-    w_t = vmap(fun)(w_mn)  # shape(rho, theta*zeta)
+    w_t = vmap_chunked(fun, chunk_size=kwargs.get("surf_batch_size"))(
+        w_mn
+    )  # shape(rho, theta*zeta)
     w_t = w_t.reshape((grid.num_rho, grid.num_theta, grid.num_zeta), order="F")
     w_t = jnp.moveaxis(w_t, 0, 1)
     data["w_Boozer_t"] = w_t.flatten(order="F")
@@ -199,13 +218,18 @@
     grid_requirement={"is_meshgrid": True, "sym": False},
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to ``grid.num_rho`` e.g. compute all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _w_z(params, transforms, profiles, data, **kwargs):
     grid = transforms["grid"]
     w_mn = data["w_Boozer_mn"].reshape((grid.num_rho, -1))
     # need to close over dz which can't be vmapped
     fun = lambda x: transforms["w"].transform(x, dz=1)
-    w_z = vmap(fun)(w_mn)  # shape(rho, theta*zeta)
+    w_z = vmap_chunked(fun, chunk_size=kwargs.get("surf_batch_size"))(
+        w_mn
+    )  # shape(rho, theta*zeta)
     w_z = w_z.reshape((grid.num_rho, grid.num_theta, grid.num_zeta), order="F")
     w_z = jnp.moveaxis(w_z, 0, 1)
     data["w_Boozer_z"] = w_z.flatten(order="F")
@@ -254,6 +278,9 @@
     grid_requirement={"is_meshgrid": True, "sym": False},
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to computing all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _nu_B_mn(params, transforms, profiles, data, **kwargs):
     norm = data["Boozer transform modes norm"]
@@ -282,7 +309,9 @@
             data["nu"],
         ),
     )
-    nu_B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, nu)
+    nu_B_mn = vmap_chunked(
+        fun, in_axes=(0, 0, 0, 0, 0), chunk_size=kwargs.get("surf_batch_size")
+    )(rho, theta_B, zeta_B, sqrtg_B_desc, nu)
     data["nu_B_mn"] = nu_B_mn.flatten()
     return data
 
@@ -428,6 +457,9 @@
     ],
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to computing all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _sqrtg_Boozer_mn(params, transforms, profiles, data, **kwargs):
     norm = data["Boozer transform modes norm"]
@@ -456,7 +488,9 @@
             data["sqrt(g)_Boozer"],
         ),
     )
-    sqrtg_B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, sqrtg_B)
+    sqrtg_B_mn = vmap_chunked(
+        fun, in_axes=(0, 0, 0, 0, 0), chunk_size=kwargs.get("surf_batch_size")
+    )(rho, theta_B, zeta_B, sqrtg_B_desc, sqrtg_B)
     data["sqrt(g)_Boozer_mn"] = sqrtg_B_mn.flatten()
     return data
 
@@ -485,6 +519,9 @@
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
     aliases=["|B|_mn"],
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to computing all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _B_mn(params, transforms, profiles, data, **kwargs):
     norm = data["Boozer transform modes norm"]
@@ -513,7 +550,9 @@
             data["|B|"],
         ),
     )
-    B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, B)
+    B_mn = vmap_chunked(
+        fun, in_axes=(0, 0, 0, 0, 0), chunk_size=kwargs.get("surf_batch_size")
+    )(rho, theta_B, zeta_B, sqrtg_B_desc, B)
     data["|B|_mn_B"] = B_mn.flatten()
     return data
 
@@ -541,6 +580,9 @@
     ],
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to computing all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _R_mn(params, transforms, profiles, data, **kwargs):
     norm = data["Boozer transform modes norm"]
@@ -569,7 +611,9 @@
             data["R"],
         ),
     )
-    R_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, R)
+    R_mn = vmap_chunked(
+        fun, in_axes=(0, 0, 0, 0, 0), chunk_size=kwargs.get("surf_batch_size")
+    )(rho, theta_B, zeta_B, sqrtg_B_desc, R)
     data["R_mn_B"] = R_mn.flatten()
     return data
 
@@ -597,6 +641,9 @@
     ],
     M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to computing all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _Z_mn(params, transforms, profiles, data, **kwargs):
     norm = data["Boozer transform modes norm"]
@@ -625,7 +672,9 @@
             data["Z"],
         ),
     )
-    Z_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B_desc, Z)
+    Z_mn = vmap_chunked(
+        fun, in_axes=(0, 0, 0, 0, 0), chunk_size=kwargs.get("surf_batch_size")
+    )(rho, theta_B, zeta_B, sqrtg_B_desc, Z)
     data["Z_mn_B"] = Z_mn.flatten()
     return data
 
@@ -886,6 +935,9 @@
     coordinates="rtz",
     data=["eta"],
     parameterization="desc.magnetic_fields._core.OmnigenousField",
+    surf_batch_size="int: Number of flux surfaces to compute simultaneously. Defaults"
+    " to computing all flux surfaces simultaneously. Decrease "
+    "to reduce memory required for computation.",
 )
 def _B_omni(params, transforms, profiles, data, **kwargs):
     # reshaped to size (L_B, M_B)
@@ -906,7 +958,9 @@
         return interp1d(x, eta_input, B, method="monotonic-0")
 
     # |B|_omnigeneous is an even function so B(-eta) = B(+eta) = B(|eta|)
-    B = vmap(_interp)(jnp.abs(eta), B_input.T)  # shape (nr, nt*nz)
+    B = vmap_chunked(_interp, in_axes=(0, 0), chunk_size=kwargs.get("surf_batch_size"))(
+        jnp.abs(eta), B_input.T
+    )  # shape (nr, nt*nz)
     B = B.reshape(
         (
             transforms["grid"].num_rho,

desc/objectives/_omnigenity.py

@@ -2,7 +2,8 @@
 
 import warnings
 
-from desc.backend import jnp, vmap
+from desc.backend import jnp
+from desc.batching import vmap_chunked
 from desc.compute import get_profiles, get_transforms
 from desc.compute._omnigenity import _omnigenity_mapping
 from desc.compute.utils import _compute as compute_fun
@@ -31,6 +32,10 @@ class QuasisymmetryBoozer(_Objective):
         Poloidal resolution of Boozer transformation. Default = 2 * eq.M.
     N_booz : int, optional
         Toroidal resolution of Boozer transformation. Default = 2 * eq.N.
+    surf_batch_size: int
+        Number of flux surfaces to compute simultaneously. Defaults to
+        computing all flux surfaces simultaneously. Decrease to reduce
+        memory required for computation.
 
     """
 
@@ -40,7 +45,7 @@ class QuasisymmetryBoozer(_Objective):
 
     _units = "(T)"
     _print_value_fmt = "Quasi-symmetry Boozer error: "
-    _static_attrs = _Objective._static_attrs + ["_helicity"]
+    _static_attrs = _Objective._static_attrs + ["_helicity", "_surf_batch_size"]
 
     def __init__(
         self,
@@ -58,13 +63,15 @@ def __init__(
         N_booz=None,
         name="QS Boozer",
         jac_chunk_size=None,
+        surf_batch_size=None,
     ):
         if target is None and bounds is None:
             target = 0
         self._grid = grid
         self.helicity = helicity
         self.M_booz = M_booz
         self.N_booz = N_booz
+        self._surf_batch_size = surf_batch_size
         super().__init__(
             things=eq,
             target=target,
@@ -142,6 +149,7 @@ def build(self, use_jit=True, verbose=1):
             "profiles": profiles,
             "matrix": matrix,
             "idx": idx,
+            "surf_batch_size": self._surf_batch_size,
         }
 
         timer.stop("Precomputing transforms")
@@ -180,6 +188,7 @@ def compute(self, params, constants=None):
             params=params,
             transforms=constants["transforms"],
             profiles=constants["profiles"],
+            surf_batch_size=constants["surf_batch_size"],
         )
         B_mn = data["|B|_mn_B"].reshape((constants["transforms"]["grid"].num_rho, -1))
         B_mn = constants["matrix"] @ B_mn.T
@@ -538,6 +547,10 @@ class Omnigenity(_Objective):
         computation time during optimization and only ``eq`` is allowed to change.
         If False, the field is allowed to change during the optimization and its
         associated data are re-computed at every iteration (Default).
+    surf_batch_size: int
+        Number of flux surfaces to compute simultaneously. Defaults to
+        computing all flux surfaces simultaneously. Decrease to reduce
+        memory required for computation.
 
     """
 
@@ -551,6 +564,7 @@ class Omnigenity(_Objective):
         "_field_data_keys",
         "_field_fixed",
         "_helicity",
+        "_surf_batch_size",
     ]
 
     _coordinates = "rtz"
@@ -577,6 +591,7 @@ def __init__(
         field_fixed=False,
         name="omnigenity",
         jac_chunk_size=None,
+        surf_batch_size=None,
     ):
         if target is None and bounds is None:
             target = 0
@@ -590,6 +605,7 @@ def __init__(
         self.eta_weight = eta_weight
         self._eq_fixed = eq_fixed
         self._field_fixed = field_fixed
+        self._surf_batch_size = surf_batch_size
         if not eq_fixed and not field_fixed:
             things = [eq, field]
         elif eq_fixed and not field_fixed:
@@ -709,6 +725,7 @@ def build(self, use_jit=True, verbose=1):
             "field_transforms": field_transforms,
             "quad_weights": w,
             "helicity": self.helicity,
+            "surf_batch_size": self._surf_batch_size,
         }
 
         if self._eq_fixed:
@@ -719,6 +736,7 @@ def build(self, use_jit=True, verbose=1):
                 params=self._eq.params_dict,
                 transforms=self._constants["eq_transforms"],
                 profiles=self._constants["eq_profiles"],
+                surf_batch_size=self._surf_batch_size,
             )
             self._constants["eq_data"] = eq_data
         if self._field_fixed:
@@ -730,6 +748,7 @@ def build(self, use_jit=True, verbose=1):
                 transforms=self._constants["field_transforms"],
                 profiles={},
                 helicity=self._constants["helicity"],
+                surf_batch_size=self._surf_batch_size,
             )
             self._constants["field_data"] = field_data
 
@@ -789,6 +808,7 @@ def compute(self, params_1=None, params_2=None, constants=None):
                 params=eq_params,
                 transforms=constants["eq_transforms"],
                 profiles=constants["eq_profiles"],
+                surf_batch_size=constants["surf_batch_size"],
             )
 
         # compute field data
@@ -814,6 +834,7 @@ def compute(self, params_1=None, params_2=None, constants=None):
                 profiles={},
                 helicity=constants["helicity"],
                 iota=eq_data["iota"][eq_grid.unique_rho_idx],
+                surf_batch_size=constants["surf_batch_size"],
             )
             theta_B = field_data["theta_B"]
             zeta_B = field_data["zeta_B"]
@@ -840,7 +861,11 @@ def _compute_B_eta_alpha(theta_B, zeta_B, B_mn):
             (field_grid.num_rho, -1)
         )
         B_mn = eq_data["|B|_mn_B"].reshape((eq_grid.num_rho, -1))
-        B_eta_alpha = vmap(_compute_B_eta_alpha)(theta_B, zeta_B, B_mn)
+        B_eta_alpha = vmap_chunked(
+            _compute_B_eta_alpha,
+            in_axes=(0, 0, 0),
+            chunk_size=constants["surf_batch_size"],
+        )(theta_B, zeta_B, B_mn)
         B_eta_alpha = B_eta_alpha.reshape(
             (field_grid.num_rho, field_grid.num_theta, field_grid.num_zeta)
         )