diff --git a/desc/objectives/_bootstrap.py b/desc/objectives/_bootstrap.py index c76a003b2..9f6850cb2 100644 --- a/desc/objectives/_bootstrap.py +++ b/desc/objectives/_bootstrap.py @@ -9,7 +9,7 @@ from desc.utils import Timer, errorif, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class BootstrapRedlConsistency(_Objective): @@ -31,31 +31,6 @@ class BootstrapRedlConsistency(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` @@ -64,22 +39,13 @@ class BootstrapRedlConsistency(_Objective): First entry must be M=1. Second entry is the toroidal mode number N, used for evaluating the Redl bootstrap current formula. Set to 0 for axisymmetry or quasi-axisymmetry; set to +/-NFP for quasi-helical symmetry. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "r" _units = "(T A m^-2)" _print_value_fmt = "Bootstrap current self-consistency error: " diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 40c53bd85..9d6267e53 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -17,7 +17,7 @@ from desc.utils import Timer, errorif, safenorm, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class _CoilObjective(_Objective): @@ -29,50 +29,16 @@ class _CoilObjective(_Objective): Coil for which the data keys will be optimized. data_keys : list of str data keys that will be optimized when this class is inherited. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, list, optional Collocation grid containing the nodes to evaluate at. If a list, must have the same structure as coil. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + coil=True, + ) + def __init__( self, coil, @@ -235,52 +201,18 @@ class CoilLength(_CoilObjective): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``target=2*np.pi``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=2*np.pi``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=2*np.pi``.", + bounds_default="``target=2*np.pi``.", + coil=True, + ) + _scalar = False # Not always a scalar, if a coilset is passed in _units = "(m)" _print_value_fmt = "Coil length: " @@ -373,52 +305,18 @@ class CoilCurvature(_CoilObjective): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``bounds=(0,1)``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(0,1)``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(0,1).``", + bounds_default="``bounds=(0,1).``", + coil=True, + ) + _scalar = False _units = "(m^-1)" _print_value_fmt = "Coil curvature: " @@ -506,52 +404,18 @@ class CoilTorsion(_CoilObjective): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``target=0``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + coil=True, + ) + _scalar = False _units = "(m^-1)" _print_value_fmt = "Coil torsion: " @@ -639,52 +503,18 @@ class CoilCurrentLength(CoilLength): ---------- coil : CoilSet or Coil Coil(s) that are to be optimized - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. Defaults to ``target=0``. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individual coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(N=2 * coil.N + 5)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + coil=True, + ) + _scalar = False _units = "(A*m)" _print_value_fmt = "Coil current length: " @@ -780,52 +610,19 @@ class CoilSetMinDistance(_Objective): ---------- coil : CoilSet Coil(s) that are to be optimized. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individial coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, list, optional Collocation grid used to discretize each coil. Defaults to the default grid for the given coil-type, see ``coils.py`` and ``curve.py`` for more details. If a list, must have the same structure as coils. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + coil=True, + ) + _scalar = False _units = "(m)" _print_value_fmt = "Minimum coil-coil distance: " @@ -951,32 +748,6 @@ class PlasmaCoilSetMinDistance(_Objective): to satisfy the Objective. coil : CoilSet Coil(s) that are to be optimized. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If array, it has to - be flattened according to the number of inputs. - bounds : tuple of float, ndarray, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Operates over all coils, not each individial coil. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. plasma_grid : Grid, optional Collocation grid containing the nodes to evaluate plasma geometry at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. @@ -997,22 +768,15 @@ class PlasmaCoilSetMinDistance(_Objective): during optimization, and self.things = [eq] only. If False, the coil coordinates are computed at every iteration. False by default, so that self.things = [coil, eq]. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + coil=True, + ) + _scalar = False _units = "(m)" _print_value_fmt = "Minimum plasma-coil distance: " @@ -1209,22 +973,6 @@ class QuadraticFlux(_Objective): field : MagneticField External field produced by coils or other source, which will be optimized to minimize the normal field error on the provided equilibrium's surface. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - len(target) must be equal to Objective.dim_f. - Default target is zero. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - len(bounds[0]) and len(bounds[1]) must be equal to Objective.dim_f - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - len(weight) must be equal to Objective.dim_f - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. source_grid : Grid, optional Collocation grid containing the nodes for plasma source terms. Default grid is detailed in the docs for ``compute_B_plasma`` @@ -1240,22 +988,14 @@ class QuadraticFlux(_Objective): vacuum : bool If true, B_plasma (the contribution to the normal field on the boundary from the plasma currents) is set to zero. - name : str - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + ) + _scalar = False _linear = False _print_value_fmt = "Boundary normal field error: " @@ -1430,30 +1170,6 @@ class ToroidalFlux(_Objective): field : MagneticField MagneticField object, the parameters of this will be optimized to minimize the objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Defaults to eq.Psi. Must be broadcastable to Objective.dim_f. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: has no effect for this objective - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. field_grid : Grid, optional Grid containing the nodes to evaluate field source at on the winding surface. (used if e.g. field is a CoilSet or @@ -1463,23 +1179,15 @@ class ToroidalFlux(_Objective): Collocation grid containing the nodes to evaluate the normal magnetic field at plasma geometry at. Defaults to a LinearGrid(L=eq.L_grid, M=eq.M_grid, zeta=jnp.array(0.0), NFP=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 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=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 use the largest size i.e ``obj.dim_x``. - """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=eq.Psi``.", + bounds_default="``target=eq.Psi``.", + loss_detail=" Note: has no effect for this objective.", + ) + _coordinates = "rtz" _units = "(Wb)" _print_value_fmt = "Toroidal Flux: " diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index 0124a84b1..7be04509e 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -7,7 +7,7 @@ from desc.utils import Timer from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class ForceBalance(_Objective): @@ -33,50 +33,16 @@ class ForceBalance(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(N)" @@ -222,46 +188,16 @@ class ForceBalanceAnisotropic(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : float, ndarray, optional - Target value(s) of the objective. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. grid : Grid, ndarray, optional - Collocation grid containing the nodes to evaluate at. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. grid : Grid, ndarray, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - name : str - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _units = "(N)" _coordinates = "rtz" _equilibrium = True @@ -389,50 +325,16 @@ class RadialForceBalance(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(N)" @@ -560,50 +462,16 @@ class HelicalForceBalance(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(N)" @@ -727,52 +595,18 @@ class Energy(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` gamma : float, optional Adiabatic (compressional) index. Default = 0. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = True _coordinates = "" _equilibrium = True @@ -909,50 +743,16 @@ class CurrentDensity(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``ConcentricGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _equilibrium = True _coordinates = "rtz" _units = "(A*m)" diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index 5752f730f..9429b328e 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -11,7 +11,7 @@ from desc.grid import LinearGrid from desc.integrals import DFTInterpolator, FFTInterpolator, virtual_casing_biot_savart from desc.nestor import Nestor -from desc.objectives.objective_funs import _Objective +from desc.objectives.objective_funs import _Objective, collect_docs from desc.utils import PRINT_WIDTH, Timer, errorif, parse_argname_change, warnif from .normalization import compute_scaling_factors @@ -38,31 +38,6 @@ class VacuumBoundaryError(_Objective): Equilibrium that will be optimized to satisfy the Objective. field : MagneticField External field produced by coils or other sources outside the plasma. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. eval_grid : Grid, optional Collocation grid containing the nodes to evaluate error at. Should be at rho=1. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` @@ -71,22 +46,13 @@ class VacuumBoundaryError(_Objective): field_fixed : bool Whether to assume the field is fixed. For free boundary solve, should be fixed. For single stage optimization, should be False (default). - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = False _linear = False _print_value_fmt = "Boundary Error: " @@ -381,31 +347,6 @@ class BoundaryError(_Objective): Equilibrium that will be optimized to satisfy the Objective. field : MagneticField External field produced by coils. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. s, q : integer Hyperparameters for singular integration scheme, s is roughly equal to the size of the local singular grid with respect to the global grid, q is the order of @@ -423,21 +364,13 @@ class BoundaryError(_Objective): loop : bool If True, evaluate integral using loops, as opposed to vmap. Slower, but uses less memory. - 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 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=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 use the largest size i.e ``obj.dim_x``. + """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + __doc__ += """ Examples -------- Assigning a surface current to the equilibrium: @@ -862,53 +795,19 @@ class BoundaryErrorNESTOR(_Objective): Equilibrium that will be optimized to satisfy the Objective. field : MagneticField External field produced by coils. - target : float, ndarray, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : float, ndarray, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to Objective.dim_f. mf, nf : integer maximum poloidal and toroidal mode numbers to use for NESTOR scalar potential. ntheta, nzeta : int number of grid points in poloidal, toroidal directions to use in NESTOR. field_grid : Grid, optional Grid used to discretize field. - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = False _linear = False _print_value_fmt = "Boundary magnetic pressure error: " diff --git a/desc/objectives/_generic.py b/desc/objectives/_generic.py index b2e568fd2..8e16e61cf 100644 --- a/desc/objectives/_generic.py +++ b/desc/objectives/_generic.py @@ -14,7 +14,7 @@ from desc.utils import errorif, parse_argname_change from .linear_objectives import _FixedObjective -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class GenericObjective(_Objective): @@ -26,51 +26,16 @@ class GenericObjective(_Objective): Name of the quantity to compute. thing : Optimizable Object that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f. - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect on this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _print_value_fmt = "Generic objective value: " def __init__( @@ -200,28 +165,13 @@ class LinearObjectiveFromUser(_FixedObjective): Custom objective function. thing : Optimizable Object whose degrees of freedom are being constrained. - target : dict of {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of dict {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : dict of {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Should be a scalar or have the same tree structure as thing.params. - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Has no effect for this objective. - name : str, optional - Name of the objective function. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _scalar = False _linear = True _fixed = True @@ -326,49 +276,16 @@ class ObjectiveFromUser(_Objective): Custom objective function. thing : Optimizable Object that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` if thing is an Equilibrium. - 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 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=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 use the largest size i.e ``obj.dim_x``. + """ + + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + __doc__ += """ Examples -------- .. code-block:: python diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 6c8c1376d..65fa1cf82 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -9,7 +9,7 @@ from desc.utils import Timer, errorif, parse_argname_change, safenorm, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs from .utils import check_if_points_are_inside_perimeter, softmin @@ -21,52 +21,21 @@ class AspectRatio(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=2``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=2``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=2``.", + bounds_default="``target=2``.", + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(dimensionless)" _print_value_fmt = "Aspect ratio: " @@ -196,52 +165,20 @@ class Elongation(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - Has no effect for this objective. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1``.", + bounds_default="``target=1``.", + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(dimensionless)" _print_value_fmt = "Elongation: " @@ -370,52 +307,19 @@ class Volume(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1``.", + bounds_default="``target=1``.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(m^3)" _print_value_fmt = "Plasma volume: " @@ -570,31 +474,6 @@ class PlasmaVesselDistance(_Objective): Equilibrium that will be optimized to satisfy the Objective. surface : Surface Bounding surface to penalize distance to. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``bounds=(1,np.inf)``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(1,np.inf)``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. surface_grid : Grid, optional Collocation grid containing the nodes to evaluate surface geometry at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. @@ -623,22 +502,14 @@ class PlasmaVesselDistance(_Objective): the array by 2/min_val to ensure that softmin_alpha*array>1. Making softmin_alpha larger than this minimum value will make the softmin a more accurate approximation of the true min. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + ) + _coordinates = "rtz" _units = "(m)" _print_value_fmt = "Plasma-vessel distance: " @@ -946,51 +817,18 @@ class MeanCurvature(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``bounds=(-np.inf, 0)``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(-np.inf, 0)``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(-np.inf,0)``.", + bounds_default="``bounds=(-np.inf,0)``.", + ) + _coordinates = "rtz" _units = "(m^-1)" _print_value_fmt = "Mean curvature: " @@ -1119,51 +957,18 @@ class PrincipalCurvature(_Objective): eq : Equilibrium or FourierRZToroidalSurface Equilibrium or FourierRZToroidalSurface that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)`` for ``Equilibrium`` or ``LinearGrid(M=2*eq.M, N=2*eq.N)`` for ``FourierRZToroidalSurface``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1``.", + bounds_default="``target=1``.", + ) + _coordinates = "rtz" _units = "(m^-1)" _print_value_fmt = "Principal curvature: " @@ -1288,50 +1093,17 @@ class BScaleLength(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``bounds=(1,np.inf)``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``bounds=(1,np.inf)``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target should be normalized before comparing to computed values. - if `normalize` is `True` and the target is in physical units, this should also - be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``bounds=(1,np.inf)``.", + bounds_default="``bounds=(1,np.inf)``.", + ) + _coordinates = "rtz" _units = "(m)" _print_value_fmt = "Magnetic field scale length: " @@ -1453,49 +1225,16 @@ class GoodCoordinates(_Objective): Equilibrium that will be optimized to satisfy the Objective. sigma : float Relative weight between the Jacobian and radial terms. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + ) + _scalar = False _units = "(dimensionless)" _print_value_fmt = "Coordinate goodness : " diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index b720608b4..8ac2faa9e 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -10,7 +10,7 @@ from desc.vmec_utils import ptolemy_linear_transform from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class QuasisymmetryBoozer(_Objective): @@ -20,31 +20,6 @@ class QuasisymmetryBoozer(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Must be a LinearGrid with sym=False. @@ -55,22 +30,13 @@ class QuasisymmetryBoozer(_Objective): Poloidal resolution of Boozer transformation. Default = 2 * eq.M. N_booz : int, optional Toroidal resolution of Boozer transformation. Default = 2 * eq.N. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _units = "(T)" _print_value_fmt = "Quasi-symmetry Boozer error: " @@ -250,52 +216,18 @@ class QuasisymmetryTwoTerm(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. helicity : tuple, optional Type of quasi-symmetry (M, N). - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(T^3)" _print_value_fmt = "Quasi-symmetry two-term error: " @@ -450,50 +382,16 @@ class QuasisymmetryTripleProduct(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(T^4/m^2)" _print_value_fmt = "Quasi-symmetry error: " @@ -614,31 +512,6 @@ class Omnigenity(_Objective): Equilibrium to be optimized to satisfy the Objective. field : OmnigenousField Omnigenous magnetic field to be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. eq_grid : Grid, optional Collocation grid containing the nodes to evaluate at for equilibrium data. Defaults to a linearly space grid on the rho=1 surface. @@ -669,22 +542,13 @@ class Omnigenity(_Objective): computation time during optimization and self.things = [eq] only. If False, the field is allowed to change during the optimization and its associated data are re-computed at every iteration (Default). - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(T)" _print_value_fmt = "Omnigenity error: " @@ -988,51 +852,16 @@ class Isodynamicity(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) + _coordinates = "rtz" _units = "(dimensionless)" _print_value_fmt = "Isodynamicity error: " diff --git a/desc/objectives/_power_balance.py b/desc/objectives/_power_balance.py index 313aa8734..bc291cb8a 100644 --- a/desc/objectives/_power_balance.py +++ b/desc/objectives/_power_balance.py @@ -6,7 +6,7 @@ from desc.utils import Timer, errorif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs class FusionPower(_Objective): @@ -24,52 +24,20 @@ class FusionPower(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=1e9``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=1e9``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. fuel : str, optional Fusion fuel, assuming a 50/50 mix. One of {'DT'}. Default = 'DT'. grid : Grid, optional Collocation grid used to compute the intermediate quantities. 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=1e9``.", + bounds_default="``target=1e9``.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(W)" _print_value_fmt = "Fusion power: " @@ -220,31 +188,6 @@ class HeatingPowerISS04(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Defaults to ``target=0``. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. Note: Has no effect for this objective. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. H_ISS04 : float, optional ISS04 confinement enhancement factor. Default = 1. gamma : float, optional @@ -252,22 +195,15 @@ class HeatingPowerISS04(_Objective): grid : Grid, optional Collocation grid used to compute the intermediate quantities. 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", + bounds_default="``target=0``.", + loss_detail=" Note: Has no effect for this objective.", + ) + _scalar = True _units = "(W)" _print_value_fmt = "Heating power: " diff --git a/desc/objectives/_profiles.py b/desc/objectives/_profiles.py index 1103421d4..0d0a522e4 100644 --- a/desc/objectives/_profiles.py +++ b/desc/objectives/_profiles.py @@ -8,65 +8,43 @@ from desc.utils import Timer, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs from .utils import _parse_callable_target_bounds - -class Pressure(_Objective): - """Target pressure profile. - - Parameters - ---------- - eq : Equilibrium - Equilibrium that will be optimized to satisfy the Objective. +profile_overwrite = { + "target": """ target : {float, ndarray, callable}, optional Target value(s) of the objective. Only used if bounds is None. Must be broadcastable to Objective.dim_f. If a callable, should take a single argument `rho` and return the desired value of the profile at those locations. Defaults to ``target=0``. + """, + "bounds": """ bounds : tuple of {float, ndarray, callable}, optional Lower and upper bounds on the objective. Overrides target. Both bounds must be broadcastable to to Objective.dim_f If a callable, each should take a single argument `rho` and return the desired bound (lower or upper) of the profile at those locations. Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. + """, +} + + +class Pressure(_Objective): + """Target pressure profile. + + Parameters + ---------- + eq : Equilibrium + Equilibrium that will be optimized to satisfy the Objective. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid)``. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs(overwrite=profile_overwrite) + _coordinates = "r" _units = "(Pa)" _print_value_fmt = "Pressure: " @@ -188,56 +166,20 @@ class RotationalTransform(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``target=0``. - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + overwrite=profile_overwrite, + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _coordinates = "r" _units = "(dimensionless)" _print_value_fmt = "Rotational transform: " @@ -372,56 +314,20 @@ class Shear(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``target=0``. - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + overwrite=profile_overwrite, + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _coordinates = "r" _units = "(dimensionless)" _print_value_fmt = "Shear: " @@ -552,56 +458,16 @@ class ToroidalCurrent(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``target=0``. - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``target=0``. - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs(overwrite=profile_overwrite) + _coordinates = "r" _units = "(A)" _print_value_fmt = "Toroidal current: " diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index 600359b40..02da27459 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -9,9 +9,27 @@ from desc.utils import Timer, errorif, setdefault, warnif from .normalization import compute_scaling_factors -from .objective_funs import _Objective +from .objective_funs import _Objective, collect_docs from .utils import _parse_callable_target_bounds +overwrite_stability = { + "target": """ + target : {float, ndarray, callable}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. If a callable, should take a + single argument `rho` and return the desired value of the profile at those + locations. Defaults to ``bounds=(0, np.inf)`` + """, + "bounds": """ + bounds : tuple of {float, ndarray, callable}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to to Objective.dim_f + If a callable, each should take a single argument `rho` and return the + desired bound (lower or upper) of the profile at those locations. + Defaults to ``bounds=(0, np.inf)`` + """, +} + class MercierStability(_Objective): """The Mercier criterion is a fast proxy for MHD stability. @@ -28,56 +46,16 @@ class MercierStability(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``bounds=(0, np.inf)`` - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``bounds=(0, np.inf)`` - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs(overwrite=overwrite_stability) + _coordinates = "r" _units = "(Wb^-2)" _print_value_fmt = "Mercier Stability: " @@ -222,56 +200,20 @@ class MagneticWell(_Objective): ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray, callable}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. If a callable, should take a - single argument `rho` and return the desired value of the profile at those - locations. Defaults to ``bounds=(0, np.inf)`` - bounds : tuple of {float, ndarray, callable}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - If a callable, each should take a single argument `rho` and return the - desired bound (lower or upper) of the profile at those locations. - Defaults to ``bounds=(0, np.inf)`` - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. Note: Has no effect for this objective. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. - deriv_mode : {"auto", "fwd", "rev"} - Specify how to compute jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on self.grad or self.hess which always use - reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. Defaults to ``LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid)``. Note that it should have poloidal and toroidal resolution, as flux surface averages are required. - 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 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=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 use the largest size i.e ``obj.dim_x``. """ + __doc__ = __doc__.rstrip() + collect_docs( + overwrite=overwrite_stability, + normalize_detail=" Note: Has no effect for this objective.", + normalize_target_detail=" Note: Has no effect for this objective.", + ) + _coordinates = "r" _units = "(dimensionless)" _print_value_fmt = "Magnetic Well: " diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py index 727c6f183..7d4772a8b 100644 --- a/desc/objectives/getters.py +++ b/desc/objectives/getters.py @@ -1,6 +1,6 @@ """Utilities for getting standard groups of objectives and constraints.""" -from desc.utils import flatten_list, is_any_instance, unique_list +from desc.utils import flatten_list, is_any_instance, isposint, unique_list from ._equilibrium import Energy, ForceBalance, HelicalForceBalance, RadialForceBalance from .linear_objectives import ( @@ -86,7 +86,10 @@ def get_equilibrium_objective(eq, mode="force", normalize=True, jac_chunk_size=" objectives = (RadialForceBalance(**kwargs), HelicalForceBalance(**kwargs)) else: raise ValueError("got an unknown equilibrium objective type '{}'".format(mode)) - return ObjectiveFunction(objectives, jac_chunk_size=jac_chunk_size) + deriv_mode = "batched" if isposint(jac_chunk_size) else "auto" + return ObjectiveFunction( + objectives, jac_chunk_size=jac_chunk_size, deriv_mode=deriv_mode + ) def get_fixed_axis_constraints(eq, profiles=True, normalize=True): diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 9cdfb61c2..28d7530d1 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -33,6 +33,153 @@ unique_list, ) +doc_target = """ + target : {float, ndarray}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. +""" +doc_bounds = """ + bounds : tuple of {float, ndarray}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to Objective.dim_f +""" +doc_weight = """ + weight : {float, ndarray}, optional + Weighting to apply to the Objective, relative to other Objectives. + Must be broadcastable to Objective.dim_f +""" +doc_normalize = """ + normalize : bool, optional + Whether to compute the error in physical units or non-dimensionalize. +""" +doc_normalize_target = """ + normalize_target : bool, optional + Whether target and bounds should be normalized before comparing to computed + values. If `normalize` is `True` and the target is in physical units, + this should also be set to True. +""" +doc_loss_function = """ + loss_function : {None, 'mean', 'min', 'max'}, optional + Loss function to apply to the objective values once computed. This loss function + is called on the raw compute value, before any shifting, scaling, or + normalization. +""" +doc_deriv_mode = """ + deriv_mode : {"auto", "fwd", "rev"} + Specify how to compute Jacobian matrix, either forward mode or reverse mode AD. + "auto" selects forward or reverse mode based on the size of the input and output + of the objective. Has no effect on self.grad or self.hess which always use + reverse mode and forward over reverse mode respectively. +""" +doc_name = """ + name : str, optional + Name of the objective. +""" +doc_jac_chunk_size = """ + jac_chunk_size : int or "auto", optional + Will calculate the Jacobian + ``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 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= 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 use the largest size i.e ``obj.dim_x``. + Defaults to ``chunk_size=None``. +""" +docs = { + "target": doc_target, + "bounds": doc_bounds, + "weight": doc_weight, + "normalize": doc_normalize, + "normalize_target": doc_normalize_target, + "loss_function": doc_loss_function, + "deriv_mode": doc_deriv_mode, + "name": doc_name, + "jac_chunk_size": doc_jac_chunk_size, +} + + +def collect_docs( + overwrite=None, + target_default="", + bounds_default="", + normalize_detail=None, + normalize_target_detail=None, + loss_detail=None, + coil=False, +): + """Collect default parameters for the docstring of Objective. + + Parameters + ---------- + overwrite : dict, optional + Dict of strings to overwrite from the _Objective's docstring. If None, + all default parameters are included as they are. Use this argument if + you want to specify a special docstring for a specific parameter in + your objective definition. + target_default : str, optional + Default value for the target parameter. + bounds_default : str, optional + Default value for the bounds parameter. + normalize_detail : str, optional + Additional information about the normalize parameter. + normalize_target_detail : str, optional + Additional information about the normalize_target parameter. + loss_detail : str, optional + Additional information about the loss function. + coil : bool, optional + Whether the objective is a coil objective. If True, adds extra docs to + target and loss_function. + + Returns + ------- + doc_params : str + String of default parameters for the docstring. + + """ + doc_params = "" + for key in docs.keys(): + if overwrite is not None and key in overwrite.keys(): + doc_params += overwrite[key].rstrip() + else: + if key == "target": + target = "" + if coil: + target += ( + "If array, it has to be flattened according to the " + + "number of inputs." + ) + if target_default != "": + target = target + " Defaults to " + target_default + doc_params += docs[key].rstrip() + target + elif key == "bounds" and bounds_default != "": + doc_params = ( + doc_params + docs[key].rstrip() + " Defaults to " + bounds_default + ) + elif key == "loss_function": + loss = "" + if coil: + loss = " Operates over all coils, not each individual coil." + if loss_detail is not None: + loss += loss_detail + doc_params += docs[key].rstrip() + loss + elif key == "normalize": + norm = "" + if normalize_detail is not None: + norm += normalize_detail + doc_params += docs[key].rstrip() + norm + elif key == "normalize_target": + norm_target = "" + if normalize_target_detail is not None: + norm_target = normalize_target_detail + doc_params += docs[key].rstrip() + norm_target + else: + doc_params += docs[key].rstrip() + + return doc_params + class ObjectiveFunction(IOAble): """Objective function comprised of one or more Objectives. @@ -113,14 +260,6 @@ def __init__( self._compiled = False self._name = name - def _set_derivatives(self): - """Choose derivative mode based on mode of sub-objectives.""" - if self._deriv_mode == "auto": - if all((obj._deriv_mode == "fwd") for obj in self.objectives): - self._deriv_mode = "batched" - else: - self._deriv_mode = "blocked" - def _unjit(self): """Remove jit compiled methods.""" methods = [ @@ -178,35 +317,34 @@ def build(self, use_jit=None, verbose=1): else: self._scalar = False - self._set_derivatives() + self._set_things() + + # setting derivative mode and chunking. + errorif( + isposint(self._jac_chunk_size) and self._deriv_mode in ["auto", "blocked"], + ValueError, + "'jac_chunk_size' was passed into ObjectiveFunction, but the " + "ObjectiveFunction is not using 'batched' deriv_mode", + ) sub_obj_jac_chunk_sizes_are_ints = [ isposint(obj._jac_chunk_size) for obj in self.objectives ] errorif( - any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode != "blocked", + any(sub_obj_jac_chunk_sizes_are_ints) and self._deriv_mode == "batched", ValueError, "'jac_chunk_size' was passed into one or more sub-objectives, but the" - " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " + " ObjectiveFunction is using 'batched' deriv_mode, so sub-objective " "'jac_chunk_size' will be ignored in favor of the ObjectiveFunction's " f"'jac_chunk_size' of {self._jac_chunk_size}." " Specify 'blocked' deriv_mode if each sub-objective is desired to have a " "different 'jac_chunk_size' for its Jacobian computation.", ) - errorif( - self._jac_chunk_size not in ["auto", None] - and self._deriv_mode == "blocked", - ValueError, - "'jac_chunk_size' was passed into ObjectiveFunction, but the " - "ObjectiveFunction is using 'blocked' deriv_mode, so sub-objective " - "'jac_chunk_size' are used to compute each sub-objective's Jacobian, " - "`ignoring the ObjectiveFunction's 'jac_chunk_size'.", - ) - - if not self.use_jit: - self._unjit() - self._set_things() - self._built = True + if self._deriv_mode == "auto": + if all((obj._deriv_mode == "fwd") for obj in self.objectives): + self._deriv_mode = "batched" + else: + self._deriv_mode = "blocked" if self._jac_chunk_size == "auto": # Heuristic estimates of fwd mode Jacobian memory usage, @@ -218,6 +356,15 @@ def build(self, use_jit=None, verbose=1): * self.dim_x ) self._jac_chunk_size = max([1, max_chunk_size]) + if self._deriv_mode == "blocked": + for obj in self.objectives: + if obj._jac_chunk_size is None: + obj._jac_chunk_size = self._jac_chunk_size + + if not self.use_jit: + self._unjit() + + self._built = True timer.stop("Objective build") if verbose > 1: @@ -799,7 +946,7 @@ def dim_x(self): @property def dim_f(self): """int: Number of objective equations.""" - if not self.built: + if not hasattr(self, "_dim_f"): raise RuntimeError("ObjectiveFunction must be built first.") return self._dim_f diff --git a/docs/adding_objectives.rst b/docs/adding_objectives.rst index 31192c991..2fa3028ba 100644 --- a/docs/adding_objectives.rst +++ b/docs/adding_objectives.rst @@ -47,42 +47,16 @@ A full example objective with comments describing the key points is given below: ---------- eq : Equilibrium Equilibrium that will be optimized to satisfy the Objective. - target : {float, ndarray}, optional - Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. - bounds : tuple of {float, ndarray}, optional - Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f - weight : {float, ndarray}, optional - Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f - normalize : bool, optional - Whether to compute the error in physical units or non-dimensionalize. - normalize_target : bool, optional - Whether target and bounds should be normalized before comparing to computed - values. If `normalize` is `True` and the target is in physical units, - this should also be set to True. - loss_function : {None, 'mean', 'min', 'max'}, optional - Loss function to apply to the objective values once computed. This loss function - is called on the raw compute value, before any shifting, scaling, or - normalization. grid : Grid, optional Collocation grid containing the nodes to evaluate at. - name : str, optional - Name of the objective function. - jac_chunk_size : int or "auto", 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 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=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 use the largest possible size. """ + # Most of the documentation is shared among all objectives, so we just inherit + # the docstring from the base class and add a few details specific to this objective. + # See the documentation of `collect_docs` for more details. + __doc__ = __doc__.rstrip() + collect_docs( + target_default="``target=0``.", bounds_default="``target=0``." + ) _coordinates = "rtz" # What coordinates is this objective a function of, with r=rho, t=theta, z=zeta? # i.e. if only a profile, it is "r" , while if all 3 coordinates it is "rtz" @@ -229,12 +203,12 @@ you will have to manually convert these vectors using the geometry utility funct ``rpz2xyz`` and/or ``rpz2xyz_vec``. See the ``PlasmaVesselDistance`` objective for an example of this. -Adapting Existing Objectives with Different Loss Funtions ---------------------------------------------------------- +Adapting Existing Objectives with Different Loss Functions +---------------------------------------------------------- If your desired objective is already implemented in DESC, but not in the correct form, -a few different loss functions are available through the the ``loss_function`` kwarg -when instantiating an Objective objective to modify the objective cost in order to adapt +a few different loss functions are available through the ``loss_function`` kwarg +when instantiating an Objective, to modify the objective cost in order to adapt the objective to your desired purpose. For example, the DESC ``RotationalTransform`` objective with ``target=iota_target`` by default forms the residual by taking the target and subtracting it from the profile at the points in the grid, resulting in a residual diff --git a/tests/test_examples.py b/tests/test_examples.py index 53fb6479b..93c700b36 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1043,6 +1043,7 @@ def test_freeb_vacuum(): objective = ObjectiveFunction( VacuumBoundaryError(eq=eq, field=ext_field, field_fixed=True), jac_chunk_size=1000, + deriv_mode="batched", ) eq, _ = eq.optimize( objective, diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index cfbf09090..1c842bd11 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -78,7 +78,7 @@ ) from desc.objectives._free_boundary import BoundaryErrorNESTOR from desc.objectives.normalization import compute_scaling_factors -from desc.objectives.objective_funs import _Objective +from desc.objectives.objective_funs import _Objective, collect_docs from desc.objectives.utils import softmax, softmin from desc.profiles import FourierZernikeProfile, PowerSeriesProfile from desc.utils import PRINT_WIDTH @@ -1323,6 +1323,7 @@ def test_derivative_modes(): AspectRatio(eq), ], deriv_mode="batched", + jac_chunk_size="auto", use_jit=False, ) obj2 = ObjectiveFunction( @@ -1332,13 +1333,15 @@ def test_derivative_modes(): AspectRatio(eq, jac_chunk_size=None), ], deriv_mode="blocked", + jac_chunk_size="auto", use_jit=False, ) obj1.build() obj2.build() # check that default size works for blocked - assert obj2.objectives[1]._jac_chunk_size is None - assert obj2.objectives[2]._jac_chunk_size is None + assert obj2.objectives[0]._jac_chunk_size == 2 + assert obj2.objectives[1]._jac_chunk_size > 0 + assert obj2.objectives[2]._jac_chunk_size > 0 # hard to say what size auto will give, just check it is >0 assert obj1._jac_chunk_size > 0 obj3.build() @@ -2930,3 +2933,19 @@ def test_objective_print_widths(): + "change the name or increase the PRINT_WIDTH in the " + "desc/utils.py file. The former is preferred." ) + + +def test_objective_docstring(): + """Test that the objective docstring and collect_docs are consistent.""" + objective_docs = _Objective.__doc__.rstrip() + doc_header = ( + "Objective (or constraint) used in the optimization of an Equilibrium.\n\n" + + " Parameters\n" + + " ----------\n" + + " things : Optimizable or tuple/list of Optimizable\n" + + " Objects that will be optimized to satisfy the Objective.\n" + ) + collected_docs = collect_docs().strip() + collected_docs = doc_header + " " + collected_docs + + assert objective_docs == collected_docs