diff --git a/desc/continuation.py b/desc/continuation.py index 635fcb876..e4d38ab72 100644 --- 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 index e76b56cca..0933fc14b 100644 --- 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): diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index c88427bb8..565bf3bdd 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -59,18 +59,18 @@ class _CoilObjective(_Objective): If a list, must have the same structure as coil. 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)`` """ @@ -88,7 +88,7 @@ def __init__( deriv_mode="auto", grid=None, name=None, - chunk_size=None, + jac_chunk_size=None, ): self._grid = grid self._data_keys = data_keys @@ -103,7 +103,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): # noqa:C901 @@ -269,18 +269,18 @@ class CoilLength(_CoilObjective): Defaults to ``LinearGrid(N=2 * coil.N + 5)`` 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)`` """ @@ -301,7 +301,7 @@ def __init__( deriv_mode="auto", grid=None, name="coil length", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 2 * np.pi @@ -318,7 +318,7 @@ def __init__( deriv_mode=deriv_mode, grid=grid, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): @@ -409,18 +409,18 @@ class CoilCurvature(_CoilObjective): Defaults to ``LinearGrid(N=2 * coil.N + 5)`` 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)`` """ @@ -441,7 +441,7 @@ def __init__( deriv_mode="auto", grid=None, name="coil curvature", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (0, 1) @@ -458,7 +458,7 @@ def __init__( deriv_mode=deriv_mode, grid=grid, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): @@ -544,18 +544,18 @@ class CoilTorsion(_CoilObjective): Defaults to ``LinearGrid(N=2 * coil.N + 5)`` 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)`` """ @@ -576,7 +576,7 @@ def __init__( deriv_mode="auto", grid=None, name="coil torsion", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -593,7 +593,7 @@ def __init__( deriv_mode=deriv_mode, grid=grid, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): @@ -679,18 +679,18 @@ class CoilCurrentLength(CoilLength): Defaults to ``LinearGrid(N=2 * coil.N + 5)`` 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)`` """ @@ -711,7 +711,7 @@ def __init__( deriv_mode="auto", grid=None, name="coil current length", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -727,7 +727,7 @@ def __init__( deriv_mode=deriv_mode, grid=grid, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): @@ -822,18 +822,18 @@ class CoilSetMinDistance(_Objective): If a list, must have the same structure as coils. 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)`` """ @@ -854,7 +854,7 @@ def __init__( deriv_mode="auto", grid=None, name="coil-coil minimum distance", - chunk_size=None, + jac_chunk_size=None, ): from desc.coils import CoilSet @@ -876,7 +876,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): @@ -1011,18 +1011,18 @@ class PlasmaCoilSetMinDistance(_Objective): False by default, so that self.things = [coil, eq]. 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)`` """ @@ -1047,7 +1047,7 @@ def __init__( eq_fixed=False, coils_fixed=False, name="plasma-coil minimum distance", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (1, np.inf) @@ -1073,7 +1073,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): @@ -1256,18 +1256,18 @@ class QuadraticFlux(_Objective): plasma currents) is set to zero. name : str 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)`` """ @@ -1291,7 +1291,7 @@ def __init__( field_grid=None, vacuum=False, name="Quadratic flux", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -1310,7 +1310,7 @@ def __init__( normalize=normalize, normalize_target=normalize_target, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): @@ -1480,18 +1480,18 @@ class ToroidalFlux(_Objective): zeta=jnp.array(0.0), NFP=eq.NFP). 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)`` @@ -1515,7 +1515,7 @@ def __init__( field_grid=None, eval_grid=None, name="toroidal-flux", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = eq.Psi @@ -1535,7 +1535,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): diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index 352674d5e..bb2c4dc7a 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -63,18 +63,18 @@ class ForceBalance(_Objective): Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` 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)`` """ @@ -95,7 +95,7 @@ def __init__( deriv_mode="auto", grid=None, name="force", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -110,7 +110,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): @@ -249,18 +249,18 @@ class ForceBalanceAnisotropic(_Objective): Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` name : str 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)`` """ @@ -282,7 +282,7 @@ def __init__( deriv_mode="auto", grid=None, name="force-anisotropic", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -297,7 +297,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): @@ -422,18 +422,18 @@ class RadialForceBalance(_Objective): Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` 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)`` """ @@ -455,7 +455,7 @@ def __init__( deriv_mode="auto", grid=None, name="radial force", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -470,7 +470,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): @@ -595,18 +595,18 @@ class HelicalForceBalance(_Objective): Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` 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)`` """ @@ -628,7 +628,7 @@ def __init__( deriv_mode="auto", grid=None, name="helical force", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -643,7 +643,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): @@ -766,18 +766,18 @@ class Energy(_Objective): Adiabatic (compressional) index. Default = 0. 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)`` """ @@ -802,7 +802,7 @@ def __init__( grid=None, gamma=0, name="energy", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -818,7 +818,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): @@ -948,18 +948,18 @@ class CurrentDensity(_Objective): Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` 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)`` """ @@ -981,7 +981,7 @@ def __init__( deriv_mode="auto", grid=None, name="current density", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -996,7 +996,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): diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index fb8858812..2a923f03e 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -73,18 +73,18 @@ class VacuumBoundaryError(_Objective): be fixed. For single stage optimization, should be False (default). 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)`` """ @@ -110,7 +110,7 @@ def __init__( field_grid=None, field_fixed=False, name="Vacuum boundary error", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -133,7 +133,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): @@ -425,18 +425,18 @@ class BoundaryError(_Objective): less memory. 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)`` @@ -483,7 +483,7 @@ def __init__( field_fixed=False, loop=True, name="Boundary error", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -510,7 +510,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): @@ -897,18 +897,18 @@ class BoundaryErrorNESTOR(_Objective): reverse mode and forward over reverse mode respectively. 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)`` """ @@ -936,7 +936,7 @@ def __init__( loss_function=None, deriv_mode="auto", name="NESTOR Boundary", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -956,7 +956,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): diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index 6147ccf04..9dc7b05f6 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -57,18 +57,18 @@ class GenericObjective(_Objective): ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium. 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)`` """ @@ -88,7 +88,7 @@ def __init__( deriv_mode="auto", grid=None, name="generic", - chunk_size=None, + jac_chunk_size=None, **kwargs, ): errorif( @@ -111,7 +111,7 @@ def __init__( loss_function=loss_function, deriv_mode=deriv_mode, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) self._p = _parse_parameterization(thing) self._scalar = not bool(data_index[self._p][self.f]["dim"]) @@ -240,7 +240,7 @@ def __init__( normalize=False, normalize_target=False, name="custom linear", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -253,7 +253,7 @@ def __init__( normalize=normalize, normalize_target=normalize_target, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=False, verbose=1): @@ -359,18 +359,18 @@ class ObjectiveFromUser(_Objective): ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium. 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)`` Examples @@ -408,7 +408,7 @@ def __init__( deriv_mode="auto", grid=None, name="custom", - chunk_size=None, + jac_chunk_size=None, **kwargs, ): errorif( @@ -431,7 +431,7 @@ def __init__( loss_function=loss_function, deriv_mode=deriv_mode, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) self._p = _parse_parameterization(thing) diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 7e63c67b8..8eb7d9b8e 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -53,18 +53,18 @@ class AspectRatio(_Objective): or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. 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)`` """ @@ -85,7 +85,7 @@ def __init__( deriv_mode="auto", grid=None, name="aspect ratio", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 2 @@ -100,7 +100,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): @@ -230,18 +230,18 @@ class Elongation(_Objective): or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. 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)`` """ @@ -262,7 +262,7 @@ def __init__( deriv_mode="auto", grid=None, name="elongation", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 1 @@ -277,7 +277,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): @@ -406,18 +406,18 @@ class Volume(_Objective): or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. 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)`` """ @@ -438,7 +438,7 @@ def __init__( deriv_mode="auto", grid=None, name="volume", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 1 @@ -453,7 +453,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): @@ -631,18 +631,18 @@ class PlasmaVesselDistance(_Objective): more accurate approximation of the true min. 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)`` """ @@ -669,7 +669,7 @@ def __init__( softmin_alpha=1.0, name="plasma-vessel distance", use_signed_distance=False, - chunk_size=None, + jac_chunk_size=None, **kwargs, ): if target is None and bounds is None: @@ -709,7 +709,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): @@ -984,18 +984,18 @@ class MeanCurvature(_Objective): or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. 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)`` """ @@ -1016,7 +1016,7 @@ def __init__( deriv_mode="auto", grid=None, name="mean curvature", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (-np.inf, 0) @@ -1031,7 +1031,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): @@ -1159,18 +1159,18 @@ class PrincipalCurvature(_Objective): or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. 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)`` """ @@ -1191,7 +1191,7 @@ def __init__( deriv_mode="auto", grid=None, name="principal-curvature", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 1 @@ -1206,7 +1206,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): @@ -1329,18 +1329,18 @@ class BScaleLength(_Objective): ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. 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)`` """ @@ -1361,7 +1361,7 @@ def __init__( deriv_mode="auto", grid=None, name="B-scale-length", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (1, np.inf) @@ -1376,7 +1376,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): @@ -1495,18 +1495,18 @@ class GoodCoordinates(_Objective): Collocation grid containing the nodes to evaluate at. 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)`` """ @@ -1528,7 +1528,7 @@ def __init__( deriv_mode="auto", grid=None, name="coordinate goodness", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -1544,7 +1544,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): diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index fb35c37f6..7c72d7f3d 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -57,18 +57,18 @@ class QuasisymmetryBoozer(_Objective): Toroidal resolution of Boozer transformation. Default = 2 * eq.N. 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)`` """ @@ -91,7 +91,7 @@ def __init__( M_booz=None, N_booz=None, name="QS Boozer", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -109,7 +109,7 @@ def __init__( loss_function=loss_function, deriv_mode=deriv_mode, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) self._print_value_fmt = "Quasi-symmetry ({},{}) Boozer error: ".format( @@ -284,18 +284,18 @@ class QuasisymmetryTwoTerm(_Objective): Type of quasi-symmetry (M, N). 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)`` """ @@ -317,7 +317,7 @@ def __init__( grid=None, helicity=(1, 0), name="QS two-term", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -333,7 +333,7 @@ def __init__( loss_function=loss_function, deriv_mode=deriv_mode, name=name, - chunk_size=chunk_size, + jac_chunk_size=jac_chunk_size, ) self._print_value_fmt = "Quasi-symmetry ({},{}) two-term error: ".format( @@ -484,18 +484,18 @@ class QuasisymmetryTripleProduct(_Objective): Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. 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)`` """ @@ -516,7 +516,7 @@ def __init__( deriv_mode="auto", grid=None, name="QS triple product", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -531,7 +531,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): @@ -677,18 +677,18 @@ class Omnigenity(_Objective): associated data are re-computed at every iteration (Default). 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)`` """ @@ -716,7 +716,7 @@ def __init__( eq_fixed=False, field_fixed=False, name="omnigenity", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -748,7 +748,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): @@ -1027,18 +1027,18 @@ class Isodynamicity(_Objective): Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. 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)`` """ @@ -1059,7 +1059,7 @@ def __init__( deriv_mode="auto", grid=None, name="Isodynamicity", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -1074,7 +1074,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): diff --git a/desc/objectives/_power_balance.py b/desc/objectives/_power_balance.py index 299b24835..63faf689c 100644 --- a/desc/objectives/_power_balance.py +++ b/desc/objectives/_power_balance.py @@ -56,6 +56,18 @@ class FusionPower(_Objective): Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP)``. name : str, optional Name of the objective function. + 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)`` """ @@ -76,6 +88,7 @@ def __init__( fuel="DT", grid=None, name="fusion power", + jac_chunk_size=None, ): errorif( fuel not in ["DT"], ValueError, f"fuel must be one of ['DT'], got {fuel}." @@ -94,6 +107,7 @@ def __init__( loss_function=loss_function, deriv_mode=deriv_mode, name=name, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): @@ -241,6 +255,18 @@ class HeatingPowerISS04(_Objective): Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP)``. name : str, optional Name of the objective function. + 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)`` """ @@ -262,6 +288,7 @@ def __init__( gamma=0, grid=None, name="heating power", + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -278,6 +305,7 @@ def __init__( loss_function=loss_function, deriv_mode=deriv_mode, name=name, + jac_chunk_size=jac_chunk_size, ) def build(self, use_jit=True, verbose=1): diff --git a/desc/objectives/_profiles.py b/desc/objectives/_profiles.py index bf72f6210..de9926d3a 100644 --- a/desc/objectives/_profiles.py +++ b/desc/objectives/_profiles.py @@ -53,18 +53,18 @@ class Pressure(_Objective): Defaults to ``LinearGrid(L=eq.L_grid)``. 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)`` """ @@ -85,7 +85,7 @@ def __init__( deriv_mode="auto", grid=None, name="pressure", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -100,7 +100,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): @@ -226,18 +226,18 @@ class RotationalTransform(_Objective): are required. 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)`` """ @@ -258,7 +258,7 @@ def __init__( deriv_mode="auto", grid=None, name="rotational transform", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -273,7 +273,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): @@ -412,18 +412,18 @@ class Shear(_Objective): are required. 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)`` """ @@ -444,7 +444,7 @@ def __init__( deriv_mode="auto", grid=None, name="shear", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (-np.inf, 0) @@ -459,7 +459,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): @@ -594,18 +594,18 @@ class ToroidalCurrent(_Objective): are required. 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)`` """ @@ -626,7 +626,7 @@ def __init__( deriv_mode="auto", grid=None, name="toroidal current", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: target = 0 @@ -641,7 +641,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): diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index 0f4336f01..1520b73e2 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -63,18 +63,18 @@ class MercierStability(_Objective): are required. 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)`` """ @@ -95,7 +95,7 @@ def __init__( deriv_mode="auto", grid=None, name="Mercier Stability", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (0, np.inf) @@ -110,7 +110,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): @@ -259,18 +259,18 @@ class MagneticWell(_Objective): are required. 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)`` """ @@ -291,7 +291,7 @@ def __init__( deriv_mode="auto", grid=None, name="Magnetic Well", - chunk_size=None, + jac_chunk_size=None, ): if target is None and bounds is None: bounds = (0, np.inf) @@ -306,7 +306,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): diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py index 5875176ca..a0cc95967 100644 --- a/desc/objectives/getters.py +++ b/desc/objectives/getters.py @@ -44,7 +44,9 @@ } -def get_equilibrium_objective(eq, mode="force", normalize=True, **kwargs): +def get_equilibrium_objective( + eq, mode="force", normalize=True, jac_chunk_size=None, **kwargs +): """Get the objective function for a typical force balance equilibrium problem. Parameters @@ -57,6 +59,8 @@ def get_equilibrium_objective(eq, mode="force", normalize=True, **kwargs): for minimizing MHD energy. normalize : bool Whether to normalize units of objective. + jac_chunk_size : int, optional + Returns ------- @@ -76,7 +80,7 @@ def get_equilibrium_objective(eq, mode="force", normalize=True, **kwargs): objectives = (RadialForceBalance(**kwargs), HelicalForceBalance(**kwargs)) else: raise ValueError("got an unknown equilibrium objective type '{}'".format(mode)) - return ObjectiveFunction(objectives, chunk_size=kwargs.get("chunk_size", None)) + return ObjectiveFunction(objectives, jac_chunk_size=jac_chunk_size) def get_fixed_axis_constraints(eq, profiles=True, normalize=True): diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index de408f91f..87b390db9 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -34,26 +34,26 @@ class ObjectiveFunction(IOAble): deriv_mode : {"auto", "batched", "blocked", "looped"} Method for computing Jacobian matrices. "batched" uses forward mode, applied to the entire objective at once, and is generally the fastest for vector valued - objectives, though most memory intensive. "blocked" builds the Jacobian for each - objective separately, using each objective's preferred AD mode (and - each objective's `chunk_size`). Generally the most efficient option when mixing - scalar and vector valued objectives. "looped" uses forward mode jacobian vector - products in a loop to build the Jacobian column by column. Generally the + objectives, though most memory intensive. "blocked" builds the Jacobian for + each objective separately, using each objective's preferred AD mode (and + each objective's `jac_chunk_size`). Generally the most efficient option when + mixing scalar and vector valued objectives. "looped" uses forward mode jacobian + vector products in a loop to build the Jacobian column by column. Generally the slowest, but most memory efficient. "auto" defaults to "batched" if all sub-objectives are set to "fwd", otherwise "blocked". name : str Name of the objective function. - chunk_size : int, optional + jac_chunk_size : int, optional If `"batched"` deriv_mode is used, will calculate the Jacobian - ``chunk_size`` columns at a time, instead of all at once. A - ``chunk_size`` of 1 is equivalent to using `"looped"` deriv_mode. - 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 it takes - 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. + ``jac_chunk_size`` columns at a time, instead of all at once. A + ``jac_chunk_size`` of 1 is equivalent to using `"looped"` deriv_mode. + 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 it takes 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. If None, it will default to ``np.ceil(dim_x/4)`` """ @@ -66,7 +66,7 @@ def __init__( use_jit=True, deriv_mode="auto", name="ObjectiveFunction", - chunk_size=None, + jac_chunk_size=None, ): if not isinstance(objectives, (tuple, list)): objectives = (objectives,) @@ -76,8 +76,7 @@ def __init__( assert use_jit in {True, False} assert deriv_mode in {"auto", "batched", "looped", "blocked"} - self.chunk_size = chunk_size - + self._jac_chunk_size = jac_chunk_size self._objectives = objectives self._use_jit = use_jit self._deriv_mode = deriv_mode @@ -151,43 +150,43 @@ def build(self, use_jit=None, verbose=1): self._scalar = False self._set_derivatives() - sub_obj_chunk_sizes = [obj.chunk_size for obj in self.objectives] + sub_obj_jac_chunk_sizes = [obj._jac_chunk_size for obj in self.objectives] warnif( - np.any(sub_obj_chunk_sizes) and self._deriv_mode != "blocked", + np.any(sub_obj_jac_chunk_sizes) and self._deriv_mode != "blocked", UserWarning, - "'chunk_size' was passed into one or more sub-objectives, but the" + "'jac_chunk_size' was passed into one or more sub-objectives, but the" "ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " - "'chunk_size' will be ignored in favor of the ObjectiveFunction's " - f"'chunk_size' of {self.chunk_size}." + "'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 'chunk_size' for its Jacobian computation.", + "different 'jac_chunk_size' for its Jacobian computation.", ) warnif( - self.chunk_size is not None and self._deriv_mode == "blocked", + self._jac_chunk_size is not None and self._deriv_mode == "blocked", UserWarning, - "'chunk_size' was passed into ObjectiveFunction, but the" + "'jac_chunk_size' was passed into ObjectiveFunction, but the" "ObjectiveFunction is using 'blocked' deriv_mode, so sub-objective " - "'chunk_size' are used to compute each sub-objective's Jacobian, " - "`ignoring the ObjectiveFunction's 'chunk_size'.", + "'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() - if self.chunk_size is None and self._deriv_mode == "batched": - # set chunk_size to 1/4 of number columns of Jacobian + if self._jac_chunk_size is None and self._deriv_mode == "batched": + # set jac_chunk_size to 1/4 of number columns of Jacobian # as the default for batched deriv_mode - self.chunk_size = int(np.ceil(self.dim_x / 4)) + self._jac_chunk_size = int(np.ceil(self.dim_x / 4)) if self._deriv_mode == "blocked": - # set chunk_size for each sub-objective + # set jac_chunk_size for each sub-objective # to 1/4 of number columns of Jacobian # as the default for batched deriv_mode for obj in self.objectives: - obj.chunk_size = ( + obj._jac_chunk_size = ( int(np.ceil(sum(t.dim_x for t in obj.things) / 4)) - if obj.chunk_size is None - else obj.chunk_size + if obj._jac_chunk_size is None + else obj._jac_chunk_size ) self._built = True @@ -558,19 +557,21 @@ def _jvp(self, v, x, constants=None, op="compute_scaled"): if len(v) == 1: jvpfun = lambda dx: Derivative.compute_jvp(fun, 0, dx, x) return batched_vectorize( - jvpfun, signature="(n)->(k)", chunk_size=self.chunk_size + jvpfun, signature="(n)->(k)", jac_chunk_size=self._jac_chunk_size )(v[0]) elif len(v) == 2: jvpfun = lambda dx1, dx2: Derivative.compute_jvp2(fun, 0, 0, dx1, dx2, x) return batched_vectorize( - jvpfun, signature="(n),(n)->(k)", chunk_size=self.chunk_size + jvpfun, signature="(n),(n)->(k)", jac_chunk_size=self._jac_chunk_size )(v[0], v[1]) elif len(v) == 3: jvpfun = lambda dx1, dx2, dx3: Derivative.compute_jvp3( fun, 0, 0, 0, dx1, dx2, dx3, x ) return batched_vectorize( - jvpfun, signature="(n),(n),(n)->(k)", chunk_size=self.chunk_size + jvpfun, + signature="(n),(n),(n)->(k)", + jac_chunk_size=self._jac_chunk_size, )(v[0], v[1], v[2]) else: raise NotImplementedError("Cannot compute JVP higher than 3rd order.") @@ -908,7 +909,7 @@ def __init__( loss_function=None, deriv_mode="auto", name=None, - chunk_size=None, + jac_chunk_size=None, ): if self._scalar: assert self._coordinates == "" @@ -919,7 +920,7 @@ def __init__( assert (bounds is None) or (target is None), "Cannot use both bounds and target" assert loss_function in [None, "mean", "min", "max"] assert deriv_mode in {"auto", "fwd", "rev"} - self.chunk_size = chunk_size + self._jac_chunk_size = jac_chunk_size self._target = target self._bounds = bounds @@ -1160,7 +1161,9 @@ def _jvp(self, v, x, constants=None, op="compute_scaled"): fun = lambda *x: getattr(self, op)(*x, constants=constants) jvpfun = lambda *dx: Derivative.compute_jvp(fun, tuple(range(len(x))), dx, *x) sig = ",".join(f"(n{i})" for i in range(len(x))) + "->(k)" - return batched_vectorize(jvpfun, signature=sig, chunk_size=self.chunk_size)(*v) + return batched_vectorize( + jvpfun, signature=sig, jac_chunk_size=self._jac_chunk_size + )(*v) @jit def jvp_scaled(self, v, x, constants=None): diff --git a/desc/utils.py b/desc/utils.py index bf609b700..c3d0d014c 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -727,28 +727,29 @@ def _unchunk(x): @_treeify -def _chunk(x, chunk_size=None): - # chunk_size=None -> add just a dummy chunk dimension, same as np.expand_dims(x, 0) +def _chunk(x, jac_chunk_size=None): + # jac_chunk_size=None -> add just a dummy chunk dimension, + # same as np.expand_dims(x, 0) if x.ndim == 0: raise ValueError("x cannot be chunked as it has 0 dimensions.") n = x.shape[0] - if chunk_size is None: - chunk_size = n + if jac_chunk_size is None: + jac_chunk_size = n - n_chunks, residual = divmod(n, chunk_size) + n_chunks, residual = divmod(n, jac_chunk_size) if residual != 0: raise ValueError( - "The first dimension of x must be divisible by chunk_size." - + f"\n Got x.shape={x.shape} but chunk_size={chunk_size}." + "The first dimension of x must be divisible by jac_chunk_size." + + f"\n Got x.shape={x.shape} but jac_chunk_size={jac_chunk_size}." ) - return x.reshape((n_chunks, chunk_size) + x.shape[1:]) + return x.reshape((n_chunks, jac_chunk_size) + x.shape[1:]) -def _chunk_size(x): +def _jac_chunk_size(x): b = set(map(lambda x: x.shape[:2], jax.tree_util.tree_leaves(x))) if len(b) != 1: raise ValueError( - "The arrays in x have inconsistent chunk_size or number of chunks" + "The arrays in x have inconsistent jac_chunk_size or number of chunks" ) return b.pop()[1] @@ -768,27 +769,27 @@ def unchunk(x_chunked): """ return _unchunk(x_chunked), functools.partial( - _chunk, chunk_size=_chunk_size(x_chunked) + _chunk, jac_chunk_size=_jac_chunk_size(x_chunked) ) -def chunk(x, chunk_size=None): +def chunk(x, jac_chunk_size=None): """Split an array (or a pytree of arrays) into chunks along the first axis. Parameters ---------- x: an array (or pytree of arrays) - chunk_size: an integer or None (default) - The first axis in x must be a multiple of chunk_size + jac_chunk_size: an integer or None (default) + The first axis in x must be a multiple of jac_chunk_size Returns ------- (x_chunked, unchunk_fn): tuple - - x_chunked is x reshaped to (-1, chunk_size)+x.shape[1:] - if chunk_size is None then it defaults to x.shape[0], i.e. just one chunk + - x_chunked is x reshaped to (-1, jac_chunk_size)+x.shape[1:] + if jac_chunk_size is None then it defaults to x.shape[0], i.e. just one chunk - unchunk_fn is a function which restores x given x_chunked """ - return _chunk(x, chunk_size), _unchunk + return _chunk(x, jac_chunk_size), _unchunk #### @@ -947,11 +948,11 @@ def f_(*args, **kwargs): # Licensed under the Apache License, Version 2.0 (the "License"); -def _eval_fun_in_chunks(vmapped_fun, chunk_size, argnums, *args, **kwargs): +def _eval_fun_in_chunks(vmapped_fun, jac_chunk_size, argnums, *args, **kwargs): n_elements = jax.tree_util.tree_leaves(args[argnums[0]])[0].shape[0] - n_chunks, n_rest = divmod(n_elements, chunk_size) + n_chunks, n_rest = divmod(n_elements, jac_chunk_size) - if n_chunks == 0 or chunk_size >= n_elements: + if n_chunks == 0 or jac_chunk_size >= n_elements: y = vmapped_fun(*args, **kwargs) else: # split inputs @@ -959,7 +960,7 @@ def _get_chunks(x): x_chunks = jax.tree_util.tree_map( lambda x_: x_[: n_elements - n_rest, ...], x ) - x_chunks = _chunk(x_chunks, chunk_size) + x_chunks = _chunk(x_chunks, jac_chunk_size) return x_chunks def _get_rest(x): @@ -989,16 +990,16 @@ def _get_rest(x): def _chunk_vmapped_function( vmapped_fun: Callable, - chunk_size: Optional[int], + jac_chunk_size: Optional[int], argnums=0, ) -> Callable: """Takes a vmapped function and computes it in chunks.""" - if chunk_size is None: + if jac_chunk_size is None: return vmapped_fun if isinstance(argnums, int): argnums = (argnums,) - return functools.partial(_eval_fun_in_chunks, vmapped_fun, chunk_size, argnums) + return functools.partial(_eval_fun_in_chunks, vmapped_fun, jac_chunk_size, argnums) def _parse_in_axes(in_axes): @@ -1018,7 +1019,7 @@ def apply_chunked( f: Callable, in_axes=0, *, - chunk_size: Optional[int], + jac_chunk_size: Optional[int], ) -> Callable: """Compute f in smaller chunks over axis 0. @@ -1045,14 +1046,14 @@ def apply_chunked( ---------- f: A function that satisfies the condition above in_axes: The axes that should be scanned along. Only supports `0` or `None` - chunk_size: The maximum size of the chunks to be used. If it is `None`, + jac_chunk_size: The maximum size of the chunks to be used. If it is `None`, chunking is disabled """ _, argnums = _parse_in_axes(in_axes) return _chunk_vmapped_function( f, - chunk_size, + jac_chunk_size, argnums, ) @@ -1061,7 +1062,7 @@ def vmap_chunked( f: Callable, in_axes=0, *, - chunk_size: Optional[int], + jac_chunk_size: Optional[int], ) -> Callable: """Behaves like jax.vmap but uses scan to chunk the computations in smaller chunks. @@ -1069,7 +1070,7 @@ def vmap_chunked( .. code-block:: python - nk.jax.apply_chunked(jax.vmap(f, in_axes), in_axes, chunk_size) + nk.jax.apply_chunked(jax.vmap(f, in_axes), in_axes, jac_chunk_size) Some limitations to `in_axes` apply. @@ -1077,7 +1078,7 @@ def vmap_chunked( ---------- f: The function to be vectorised. in_axes: The axes that should be scanned along. Only supports `0` or `None` - chunk_size: The maximum size of the chunks to be used. If it is `None`, + jac_chunk_size: The maximum size of the chunks to be used. If it is `None`, chunking is disabled @@ -1087,10 +1088,12 @@ def vmap_chunked( """ in_axes, argnums = _parse_in_axes(in_axes) vmapped_fun = jax.vmap(f, in_axes=in_axes) - return _chunk_vmapped_function(vmapped_fun, chunk_size, argnums) + return _chunk_vmapped_function(vmapped_fun, jac_chunk_size, argnums) -def batched_vectorize(pyfunc, *, excluded=frozenset(), signature=None, chunk_size=None): +def batched_vectorize( + pyfunc, *, excluded=frozenset(), signature=None, jac_chunk_size=None +): """Define a vectorized function with broadcasting and batching. below is taken from JAX @@ -1120,7 +1123,7 @@ def batched_vectorize(pyfunc, *, excluded=frozenset(), signature=None, chunk_siz provided, ``pyfunc`` will be called with (and expected to return) arrays with shapes given by the size of corresponding core dimensions. By default, pyfunc is assumed to take scalars arrays as input and output. - chunk_size: the size of the batches to pass to vmap. if 1, will only + jac_chunk_size: the size of the batches to pass to vmap. if 1, will only Returns ------- @@ -1202,7 +1205,7 @@ def wrapped(*args, **kwargs): else: # change the vmap here to chunked_vmap vectorized_func = vmap_chunked( - vectorized_func, in_axes, chunk_size=chunk_size + vectorized_func, in_axes, jac_chunk_size=jac_chunk_size ) result = vectorized_func(*squeezed_args) diff --git a/docs/adding_objectives.rst b/docs/adding_objectives.rst index 04bebcf79..d97c31fbf 100644 --- a/docs/adding_objectives.rst +++ b/docs/adding_objectives.rst @@ -70,18 +70,18 @@ A full example objective with comments describing the key points is given below: Collocation grid containing the nodes to evaluate at. name : str, optional Name of the objective function. - chunk_size : int, optional + jac_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, + 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 - linearly proportional to ``chunk_size``: the smaller the ``chunk_size``, the + linearly proportional to ``jac_chunk_size``: the smaller the ``jac_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`` + usage). The time to compute the Jacobian is roughly ``t ~1/jac_chunk_size`` + with some baseline time, 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 `chunk_size` i.e. ``dim_x`` + If None, it will default to the largest possible `jac_chunk_size` i.e. ``dim_x`` """ @@ -100,7 +100,7 @@ A full example objective with comments describing the key points is given below: normalize_target=True, grid=None, name="QS triple product", - chunk_size=None, + jac_chunk_size=None, ): # we don't have to do much here, mostly just call ``super().__init__()`` if target is None and bounds is None: @@ -114,7 +114,7 @@ A full example objective with comments describing the key points is given below: normalize=normalize, normalize_target=normalize_target, name=name, - chunk_size=chunk_size + jac_chunk_size=jac_chunk_size ) def build(self, use_jit=True, verbose=1): @@ -255,23 +255,23 @@ Reducing Memory Size of Objective Jacobian Calculation During optimization, one of the most memory-intensive steps is the calculation of the Jacobian of the cost function. This memory cost comes from attempting to calculate the entire Jacobian matrix in one vectorized operation. However, this can be tuned between high memory usage but quick (default) -and low memory usage but slower with the ``chunk_size`` keyword argument. By default, where this matters +and low memory usage but slower with the ``jac_chunk_size`` keyword argument. By default, where this matters is when creating the overall ``ObjectiveFunction`` to be used in the optimization (where by default ``deriv_mode="batched"``). The Jacobian is a matrix of shape [``obj.dim_f`` x ``obj.dim_x``], and the calculation of the Jacobian is vectorized over -the columns (the ``obj.dim_x`` dimension), where ``obj`` is the ``ObjectiveFunction`` object. Passing in the ``chunk_size`` attribute allows one to split up -the vectorized computation into chunks of ``chunk_size`` columns at a time, allowing one to compute the Jacobian +the columns (the ``obj.dim_x`` dimension), where ``obj`` is the ``ObjectiveFunction`` object. Passing in the ``jac_chunk_size`` attribute allows one to split up +the vectorized computation into chunks of ``jac_chunk_size`` columns at a time, allowing one to compute the Jacobian in a slightly slower, but more memory-efficient manner. 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 ``chunk_size=None``, it will default to the largest -possible `chunk_size` i.e. ``obj.dim_x``. +linearly proportional to ``jac_chunk_size``: the smaller the ``jac_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/jac_chunk_size`` +with some baseline time, 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 ``jac_chunk_size=None``, it will default to the largest +possible `jac_chunk_size` i.e. ``obj.dim_x``. If ``deriv_mode="blocked"`` is specified when the ``ObjectiveFunction`` is created, then the Jacobian will be calculated individually for each of the sub-objectives inside of the ``ObjectiveFunction``, and in that case -the ``chunk_size`` of the individual ``_Objective`` objects inside of the ``ObjectiveFunction`` will be used. -For example, if ``obj1 = QuasisymmetryTripleProduct(eq, chunk_size=100)``, ``obj2 = MeanCurvature(eq, chunk_size=2000)`` +the ``jac_chunk_size`` of the individual ``_Objective`` objects inside of the ``ObjectiveFunction`` will be used. +For example, if ``obj1 = QuasisymmetryTripleProduct(eq, jac_chunk_size=100)``, ``obj2 = MeanCurvature(eq, jac_chunk_size=2000)`` and ``obj = ObjectiveFunction((obj1, obj2), deriv_mode="blocked")``, then the Jacobian will be calculated with a -``chunk_size=100`` for the quasisymmetry part and a ``chunk_size=2000`` for the curvature part, then the full Jacobian +``jac_chunk_size=100`` for the quasisymmetry part and a ``jac_chunk_size=2000`` for the curvature part, then the full Jacobian will be formed as a block diagonal matrix with the individual Jacobians of these two objectives. diff --git a/tests/test_examples.py b/tests/test_examples.py index 4d3308d2a..0ac1ed8f6 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -301,7 +301,7 @@ def test_ATF_results(tmpdir_factory): spectral_indexing=eq0.spectral_indexing, ) eqf = EquilibriaFamily.solve_continuation_automatic( - eq, verbose=2, checkpoint_path=output_dir.join("ATF.h5"), chunk_size=500 + eq, verbose=2, checkpoint_path=output_dir.join("ATF.h5"), jac_chunk_size=500 ) eqf = load(output_dir.join("ATF.h5")) rho_err, theta_err = area_difference_desc(eq0, eqf[-1]) @@ -1040,7 +1040,8 @@ def test_freeb_vacuum(): FixPsi(eq=eq), ) objective = ObjectiveFunction( - VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True), chunk_size=1000 + VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True), + jac_chunk_size=1000, ) eq, _ = eq.optimize( objective, diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 9140ed4f9..1d9d548d1 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -1284,7 +1284,7 @@ def test_derivative_modes(): surf = FourierRZToroidalSurface() obj1 = ObjectiveFunction( [ - PlasmaVesselDistance(eq, surf, chunk_size=1), + PlasmaVesselDistance(eq, surf, jac_chunk_size=1), MagneticWell(eq), ], deriv_mode="batched", @@ -1292,11 +1292,11 @@ def test_derivative_modes(): ) obj2 = ObjectiveFunction( [ - PlasmaVesselDistance(eq, surf, chunk_size=2), + PlasmaVesselDistance(eq, surf, jac_chunk_size=2), MagneticWell(eq), ], deriv_mode="blocked", - chunk_size=10, + jac_chunk_size=10, use_jit=False, ) obj3 = ObjectiveFunction( @@ -1307,12 +1307,12 @@ def test_derivative_modes(): deriv_mode="looped", use_jit=False, ) - with pytest.warns(UserWarning, match="chunk_size"): + with pytest.warns(UserWarning, match="jac_chunk_size"): obj1.build() - with pytest.warns(UserWarning, match="chunk_size"): + with pytest.warns(UserWarning, match="jac_chunk_size"): obj2.build() # check that default size works for blocked - assert obj2.objectives[1].chunk_size == np.ceil( + assert obj2.objectives[1].jac_chunk_size == np.ceil( sum(t.dim_x for t in obj2.objectives[1].things) / 4 ) obj3.build()