diff --git a/desc/objectives/_neoclassical.py b/desc/objectives/_neoclassical.py index b32507d9d..9c668edd2 100644 --- a/desc/objectives/_neoclassical.py +++ b/desc/objectives/_neoclassical.py @@ -29,6 +29,20 @@ from .objective_funs import _Objective, collect_docs from .utils import _parse_callable_target_bounds +_bounce_overwrite = { + "deriv_mode": """ + deriv_mode : {"auto", "fwd", "rev"} + Specify how to compute Jacobian matrix, either forward mode or reverse mode AD. + "auto" selects forward or reverse mode based on the size of the input and output + of the objective. Has no effect on ``self.grad`` or ``self.hess`` which always + use reverse mode and forward over reverse mode respectively. + + Default is ``fwd``. If ``rev`` is chosen, then ``jac_chunk_size=1`` is chosen + by default. In ``rev`` mode, reducing the pitch angle parameter ``batch_size`` + does not reduce memory, so it is recommended to retain the default for that. + """ +} + class EffectiveRipple(_Objective): """The effective ripple is a proxy for neoclassical transport. @@ -96,6 +110,7 @@ class EffectiveRipple(_Objective): bounds_default="``target=0``.", normalize_detail=" Note: Has no effect for this objective.", normalize_target_detail=" Note: Has no effect for this objective.", + overwrite=_bounce_overwrite, ) _coordinates = "r" @@ -111,7 +126,7 @@ def __init__( normalize=True, normalize_target=True, loss_function=None, - deriv_mode="auto", + deriv_mode="fwd", grid=None, name="Effective ripple", jac_chunk_size=None, @@ -142,6 +157,10 @@ def __init__( "num_pitch": num_pitch, "batch_size": batch_size, } + if deriv_mode == "rev" and jac_chunk_size is None: + # Reverse mode is bottlenecked by coordinate mapping. + # Compute Jacobian one flux surface at a time. + jac_chunk_size = 1 super().__init__( things=eq, @@ -171,14 +190,19 @@ def build(self, use_jit=True, verbose=1): if self._grid is None: self._grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) assert self._grid.can_fft + # Constants can only hold 1D arrays for some reason. self._constants["clebsch"] = FourierChebyshevSeries.nodes( self._X, self._Y, self._grid.compress(self._grid.nodes[:, 0]), domain=(0, 2 * np.pi), + ).ravel() + self._constants["fieldline quad x"], self._constants["fieldline quad w"] = ( + leggauss(self._hyperparam["Y_B"] // 2) + ) + self._constants["quad x"], self._constants["quad w"] = chebgauss2( + self._hyperparam.pop("num_quad") ) - self._constants["fieldline_quad"] = leggauss(self._hyperparam["Y_B"] // 2) - self._constants["quad"] = chebgauss2(self._hyperparam.pop("num_quad")) self._dim_f = self._grid.num_rho self._target, self._bounds = _parse_callable_target_bounds( @@ -239,13 +263,16 @@ def compute(self, params, constants=None): self._X, self._Y, iota=constants["transforms"]["grid"].compress(data["iota"]), - clebsch=constants["clebsch"], + clebsch=constants["clebsch"].reshape(-1, 3), # Pass in params so that root finding is done with the new # perturbed λ coefficients and not the original equilibrium's. params=params, ), - fieldline_quad=constants["fieldline_quad"], - quad=constants["quad"], + fieldline_quad=( + constants["fieldline quad x"], + constants["fieldline quad w"], + ), + quad=(constants["quad x"], constants["quad w"]), **self._hyperparam, ) return constants["transforms"]["grid"].compress(data["effective ripple"]) @@ -313,10 +340,10 @@ class GammaC(_Objective): Note that Nemov's Γ_c converges to a finite nonzero value in the infinity limit of the number of toroidal transits. - Velasco's expression has a secular term that will drive the result - to zero as the number of toroidal transits increases unless the - secular term is averaged out from all the singular integrals. - Therefore, an optimization using Velasco's metric should be evaluated by + Velasco's expression has a secular term that may drive the result + to zero as the number of toroidal transits increases if the quadratures + are unable to average out the secular term from all the singular integrals. + So an optimization using Velasco's metric may need to be evaluated by measuring decrease in Γ_c at a fixed number of toroidal transits until unless an adaptive quadrature is used. @@ -327,6 +354,7 @@ class GammaC(_Objective): bounds_default="``target=0``.", normalize_detail=" Note: Has no effect for this objective.", normalize_target_detail=" Note: Has no effect for this objective.", + overwrite=_bounce_overwrite, ) _coordinates = "r" @@ -342,7 +370,7 @@ def __init__( normalize=True, normalize_target=True, loss_function=None, - deriv_mode="auto", + deriv_mode="fwd", grid=None, name="Gamma_c", jac_chunk_size=None, @@ -378,6 +406,10 @@ def __init__( self._key = "Gamma_c" else: self._key = "Gamma_c Velasco" + if deriv_mode == "rev" and jac_chunk_size is None: + # Reverse mode is bottlenecked by coordinate mapping. + # Compute Jacobian one flux surface at a time. + jac_chunk_size = 1 super().__init__( things=eq, @@ -407,19 +439,22 @@ def build(self, use_jit=True, verbose=1): if self._grid is None: self._grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False) assert self._grid.can_fft + # Constants can only hold 1D arrays for some reason. self._constants["clebsch"] = FourierChebyshevSeries.nodes( self._X, self._Y, self._grid.compress(self._grid.nodes[:, 0]), domain=(0, 2 * np.pi), + ).ravel() + self._constants["fieldline quad x"], self._constants["fieldline quad w"] = ( + leggauss(self._hyperparam["Y_B"] // 2) ) - self._constants["fieldline_quad"] = leggauss(self._hyperparam["Y_B"] // 2) num_quad = self._hyperparam.pop("num_quad") - self._constants["quad"] = get_quadrature( + self._constants["quad x"], self._constants["quad w"] = get_quadrature( leggauss(num_quad), (automorphism_sin, grad_automorphism_sin), ) - self._constants["quad2"] = chebgauss2(num_quad) + self._constants["quad2 x"], self._constants["quad2 w"] = chebgauss2(num_quad) self._dim_f = self._grid.num_rho self._target, self._bounds = _parse_callable_target_bounds( @@ -459,7 +494,7 @@ def compute(self, params, constants=None): if constants is None: constants = self.constants if "quad2" in constants: - self._hyperparam["quad2"] = constants["quad2"] + self._hyperparam["quad2"] = (constants["quad2 x"], constants["quad2 w"]) eq = self.things[0] data = compute_fun( @@ -478,13 +513,16 @@ def compute(self, params, constants=None): self._X, self._Y, iota=constants["transforms"]["grid"].compress(data["iota"]), - clebsch=constants["clebsch"], + clebsch=constants["clebsch"].reshape(-1, 3), # Pass in params so that root finding is done with the new # perturbed λ coefficients and not the original equilibrium's. params=params, ), - fieldline_quad=constants["fieldline_quad"], - quad=constants["quad"], + fieldline_quad=( + constants["fieldline quad x"], + constants["fieldline quad w"], + ), + quad=(constants["quad x"], constants["quad w"]), **self._hyperparam, ) return constants["transforms"]["grid"].compress(data[self._key]) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 31a550166..0074e090e 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -66,9 +66,9 @@ doc_deriv_mode = """ deriv_mode : {"auto", "fwd", "rev"} Specify how to compute Jacobian matrix, either forward mode or reverse mode AD. - "auto" selects forward or reverse mode based on the size of the input and output - of the objective. Has no effect on ``self.grad`` or ``self.hess`` which always - use reverse mode and forward over reverse mode respectively. + ``auto`` selects forward or reverse mode based on the size of the input and + output of the objective. Has no effect on ``self.grad`` or ``self.hess`` which + always use reverse mode and forward over reverse mode respectively. """ doc_name = """ name : str, optional diff --git a/docs/notebooks/tutorials/EffectiveRipple.ipynb b/docs/notebooks/tutorials/EffectiveRipple.ipynb index 00bccd724..389f0920f 100644 --- a/docs/notebooks/tutorials/EffectiveRipple.ipynb +++ b/docs/notebooks/tutorials/EffectiveRipple.ipynb @@ -385,9 +385,9 @@ " num_well=30 * 10,\n", " num_quad=32,\n", " num_pitch=45,\n", - " # TODO: It seems batch_size has no effect on memory consumed by Jacobian in reverse mode.\n", - " batch_size=1,\n", - " deriv_mode=\"fwd\",\n", + " # Uncomment to compute at only batch_size pitch angles at a time.\n", + " # This will reduce peak memory by 2.5 GB.\n", + " # batch_size=1,\n", " ),\n", " AspectRatio(eq1, bounds=(8, 11), weight=1e3),\n", " GenericObjective(\n",