diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 1b345a6492..c926e891b5 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -35,6 +35,7 @@ _field, _geometry, _metric, + _neoclassical, _omnigenity, _profiles, _stability, diff --git a/desc/compute/_equil.py b/desc/compute/_equil.py index de0fad7797..b1e9142253 100644 --- a/desc/compute/_equil.py +++ b/desc/compute/_equil.py @@ -658,12 +658,12 @@ def _F_anisotropic(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="", - data=["|B|", "sqrt(g)"], + data=["|B|^2", "sqrt(g)"], resolution_requirement="rtz", ) def _W_B(params, transforms, profiles, data, **kwargs): data["W_B"] = jnp.sum( - data["|B|"] ** 2 * data["sqrt(g)"] * transforms["grid"].weights + data["|B|^2"] * data["sqrt(g)"] * transforms["grid"].weights ) / (2 * mu_0) return data diff --git a/desc/compute/_neoclassical.py b/desc/compute/_neoclassical.py new file mode 100644 index 0000000000..4461efeeaa --- /dev/null +++ b/desc/compute/_neoclassical.py @@ -0,0 +1,274 @@ +"""Compute functions for neoclassical transport. + +Notes +----- +Some quantities require additional work to compute at the magnetic axis. +A Python lambda function is used to lazily compute the magnetic axis limits +of these quantities. These lambda functions are evaluated only when the +computational grid has a node on the magnetic axis to avoid potentially +expensive computations. +""" + +from functools import partial + +from quadax import simpson + +from desc.backend import imap, jit, jnp + +from ..integrals.bounce_integral import Bounce1D +from ..integrals.quad_utils import chebgauss2 +from ..utils import safediv +from .data_index import register_compute_fun + +_bounce_doc = { + "quad": ( + "tuple[jnp.ndarray] : Quadrature points and weights for bounce integrals. " + "Default option is well tested." + ), + "num_quad": ( + "int : Resolution for quadrature of bounce integrals. " + "Default is 32. This option is ignored if given ``quad``." + ), + "num_pitch": "int : Resolution for quadrature over velocity coordinate.", + "num_well": ( + "int : Maximum number of wells to detect for each pitch and field line. " + "Default is to detect all wells, but due to limitations in JAX this option " + "may consume more memory. Specifying a number that tightly upper bounds " + "the number of wells will increase performance." + ), + "batch": "bool : Whether to vectorize part of the computation. Default is true.", +} + + +def _alpha_mean(f): + """Simple mean over field lines. + + Simple mean rather than integrating over α and dividing by 2π + (i.e. f.T.dot(dα) / dα.sum()), because when the toroidal angle extends + beyond one transit we need to weight all field lines uniformly, regardless + of their spacing wrt α. + """ + return f.mean(axis=0) + + +def _compute(fun, fun_data, data, grid, num_pitch, simp=False, reduce=True): + """Compute ``fun`` for each α and ρ value iteratively to reduce memory usage. + + Parameters + ---------- + fun : callable + Function to compute. + fun_data : dict[str, jnp.ndarray] + Data to provide to ``fun``. This dict will be modified. + data : dict[str, jnp.ndarray] + DESC data dict. + simp : bool + Whether to use an open Simpson rule instead of uniform weights. + reduce : bool + Whether to compute mean over α and expand to grid. + Default is true. + """ + pitch_inv, pitch_inv_weight = Bounce1D.get_pitch_inv_quad( + grid.compress(data["min_tz |B|"]), + grid.compress(data["max_tz |B|"]), + num_pitch, + simp=simp, + ) + + def foreach_rho(x): + # using same λ values for every field line α on flux surface ρ + x["pitch_inv"] = pitch_inv + x["pitch_inv weight"] = pitch_inv_weight + return imap(fun, x) + + for name in Bounce1D.required_names: + fun_data[name] = data[name] + for name in fun_data: + fun_data[name] = Bounce1D.reshape_data(grid, fun_data[name]) + out = imap(foreach_rho, fun_data) + return grid.expand(_alpha_mean(out)) if reduce else out + + +@register_compute_fun( + name="fieldline length", + label="\\int_{\\zeta_{\\mathrm{min}}}^{\\zeta_{\\mathrm{max}}}" + " \\frac{d\\zeta}{|B^{\\zeta}|}", + units="m / T", + units_long="Meter / tesla", + description="(Mean) proper length of field line(s)", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["B^zeta"], + resolution_requirement="z", + source_grid_requirement={"coordinates": "raz", "is_meshgrid": True}, +) +def _fieldline_length(data, transforms, profiles, **kwargs): + grid = transforms["grid"].source_grid + L_ra = simpson( + y=grid.meshgrid_reshape(1 / data["B^zeta"], "arz"), + x=grid.compress(grid.nodes[:, 2], surface_label="zeta"), + axis=-1, + ) + data["fieldline length"] = grid.expand(jnp.abs(_alpha_mean(L_ra))) + return data + + +@register_compute_fun( + name="fieldline length/volume", + label="\\int_{\\zeta_{\\mathrm{min}}}^{\\zeta_{\\mathrm{max}}}" + " \\frac{d\\zeta}{|B^{\\zeta} \\sqrt g|}", + units="1 / Wb", + units_long="Inverse webers", + description="(Mean) proper length over volume of field line(s)", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["B^zeta", "sqrt(g)"], + resolution_requirement="z", + source_grid_requirement={"coordinates": "raz", "is_meshgrid": True}, +) +def _fieldline_length_over_volume(data, transforms, profiles, **kwargs): + grid = transforms["grid"].source_grid + G_ra = simpson( + y=grid.meshgrid_reshape(1 / (data["B^zeta"] * data["sqrt(g)"]), "arz"), + x=grid.compress(grid.nodes[:, 2], surface_label="zeta"), + axis=-1, + ) + data["fieldline length/volume"] = grid.expand(jnp.abs(_alpha_mean(G_ra))) + return data + + +@register_compute_fun( + name="effective ripple 3/2", + label=( + # ε¹ᐧ⁵ = π/(8√2) R₀²〈|∇ψ|〉⁻² B₀⁻¹ ∫dλ λ⁻² 〈 ∑ⱼ Hⱼ²/Iⱼ 〉 + "\\epsilon_{\\mathrm{eff}}^{3/2} = \\frac{\\pi}{8 \\sqrt{2}} " + "R_0^2 \\langle \\vert\\nabla \\psi\\vert \\rangle^{-2} " + "B_0^{-1} \\int d\\lambda \\lambda^{-2} " + "\\langle \\sum_j H_j^2 / I_j \\rangle" + ), + units="~", + units_long="None", + description="Effective ripple modulation amplitude to 3/2 power", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=[ + "min_tz |B|", + "max_tz |B|", + "kappa_g", + "R0", + "|grad(rho)|", + "<|grad(rho)|>", + "fieldline length", + ] + + Bounce1D.required_names, + resolution_requirement="z", + source_grid_requirement={"coordinates": "raz", "is_meshgrid": True}, + **_bounce_doc, + # Some notes on choosing the resolution hyperparameters: + # The default settings were chosen such that the effective ripple profile on + # the W7-X stellarator looks similar to the profile computed at higher resolution, + # indicating convergence. The parameters ``num_transit`` and ``knots_per_transit`` + # have a stronger effect on the result. As a reference for W7-X, when computing the + # effective ripple by tracing a single field line on each flux surface, a density of + # 100 knots per toroidal transit accurately reconstructs the ripples along the field + # line. After 10 toroidal transits convergence is apparent (after 15 the returns + # diminish). Dips in the resulting profile indicates insufficient ``num_transit``. + # Unreasonably high values indicates insufficient ``knots_per_transit``. + # One can plot the field line with ``Bounce1D.plot`` to see if the number of knots + # was sufficient to reconstruct the field line. + # TODO: Improve performance... see GitHub issue #1045. + # Need more efficient function approximation of |B|(α, ζ). +) +@partial(jit, static_argnames=["num_quad", "num_pitch", "num_well", "batch"]) +def _epsilon_32(params, transforms, profiles, data, **kwargs): + """https://doi.org/10.1063/1.873749. + + Evaluation of 1/ν neoclassical transport in stellarators. + V. V. Nemov, S. V. Kasilov, W. Kernbichler, M. F. Heyn. + Phys. Plasmas 1 December 1999; 6 (12): 4622–4632. + """ + # noqa: unused dependency + if "quad" in kwargs: + quad = kwargs["quad"] + else: + quad = chebgauss2(kwargs.get("num_quad", 32)) + num_well = kwargs.get("num_well", None) + batch = kwargs.get("batch", True) + grid = transforms["grid"].source_grid + + def dH(data, B, pitch): + # Integrand of Nemov eq. 30 with |∂ψ/∂ρ| (λB₀)¹ᐧ⁵ removed. + return ( + jnp.sqrt(jnp.abs(1 - pitch * B)) + * (4 / (pitch * B) - 1) + * data["|grad(rho)|*kappa_g"] + / B + ) + + def dI(data, B, pitch): + # Integrand of Nemov eq. 31. + return jnp.sqrt(jnp.abs(1 - pitch * B)) / B + + def eps_32(data): + """(∂ψ/∂ρ)⁻² B₀⁻² ∫ dλ λ⁻² ∑ⱼ Hⱼ²/Iⱼ.""" + # B₀ has units of λ⁻¹. + # Nemov's ∑ⱼ Hⱼ²/Iⱼ = (∂ψ/∂ρ)² (λB₀)³ ``(H**2 / I).sum(axis=-1)``. + # (λB₀)³ d(λB₀)⁻¹ = B₀² λ³ d(λ⁻¹) = -B₀² λ dλ. + bounce = Bounce1D(grid, data, quad, automorphism=None, is_reshaped=True) + H, I = bounce.integrate( + [dH, dI], + data["pitch_inv"], + data, + "|grad(rho)|*kappa_g", + points=bounce.points(data["pitch_inv"], num_well=num_well), + batch=batch, + ) + return ( + safediv(H**2, I).sum(axis=-1) + * data["pitch_inv"] ** (-3) + * data["pitch_inv weight"] + ).sum(axis=-1) + + # Interpolate |∇ρ| κ_g since it is smoother than κ_g alone. + B0 = data["max_tz |B|"] + data["effective ripple 3/2"] = ( + jnp.pi + / (8 * 2**0.5) + * (B0 * data["R0"] / data["<|grad(rho)|>"]) ** 2 + * _compute( + eps_32, + {"|grad(rho)|*kappa_g": data["|grad(rho)|"] * data["kappa_g"]}, + data, + grid, + kwargs.get("num_pitch", 50), + ) + / data["fieldline length"] + ) + return data + + +@register_compute_fun( + name="effective ripple", + label="\\epsilon_{\\mathrm{eff}}", + units="~", + units_long="None", + description="Effective ripple modulation amplitude", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="r", + data=["effective ripple 3/2"], +) +def _effective_ripple(params, transforms, profiles, data, **kwargs): + data["effective ripple"] = data["effective ripple 3/2"] ** (2 / 3) + return data diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index 2355a6fcb4..8ffe1f5a3d 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -608,10 +608,10 @@ def _B_omni(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["b", "grad(|B|)", "|B|", "grad(psi)"], + data=["b", "grad(|B|)", "|B|^2", "grad(psi)"], ) def _isodynamicity(params, transforms, profiles, data, **kwargs): data["isodynamicity"] = ( - dot(cross(data["b"], data["grad(|B|)"]), data["grad(psi)"]) / data["|B|"] ** 2 + dot(cross(data["b"], data["grad(|B|)"]), data["grad(psi)"]) / data["|B|^2"] ) return data diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index f2fde0c158..f622278c0a 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -9,12 +9,13 @@ expensive computations. """ +from quadax import cumulative_simpson from scipy.constants import elementary_charge, mu_0 from desc.backend import cond, jnp from ..integrals.surface_integral import surface_averages, surface_integrals -from ..utils import cumtrapz, dot, safediv +from ..utils import dot, safediv from .data_index import register_compute_fun @@ -142,8 +143,11 @@ def _chi_r(params, transforms, profiles, data, **kwargs): resolution_requirement="r", ) def _chi(params, transforms, profiles, data, **kwargs): - chi_r = transforms["grid"].compress(data["chi_r"]) - chi = cumtrapz(chi_r, transforms["grid"].compress(data["rho"]), initial=0) + chi = cumulative_simpson( + y=transforms["grid"].compress(data["chi_r"]), + x=transforms["grid"].compress(data["rho"]), + initial=0, + ) data["chi"] = transforms["grid"].expand(chi) return data diff --git a/desc/compute/_stability.py b/desc/compute/_stability.py index 5f7f4c60d1..5799982a32 100644 --- a/desc/compute/_stability.py +++ b/desc/compute/_stability.py @@ -1,4 +1,4 @@ -"""Compute functions for Mercier stability objectives. +"""Compute functions for stability objectives. Notes ----- diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index 20f5fe4604..5bf6cbb947 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -1240,7 +1240,7 @@ def map_coordinates( **kwargs, ) - def get_rtz_grid( + def _get_rtz_grid( self, radial, poloidal, diff --git a/desc/integrals/_bounce_utils.py b/desc/integrals/_bounce_utils.py index b49f18d615..3c3fa99e2b 100644 --- a/desc/integrals/_bounce_utils.py +++ b/desc/integrals/_bounce_utils.py @@ -16,12 +16,6 @@ polyroot_vec, polyval_vec, ) -from desc.integrals._quad_utils import ( - bijection_from_disc, - grad_bijection_from_disc, - simpson2, - uniform, -) from desc.integrals.basis import ( FourierChebyshevSeries, PiecewiseChebyshevSeries, @@ -29,6 +23,12 @@ _in_epigraph_and, _plot_intersect, ) +from desc.integrals.quad_utils import ( + bijection_from_disc, + grad_bijection_from_disc, + simpson2, + uniform, +) from desc.utils import ( atleast_nd, errorif, diff --git a/desc/integrals/_interp_utils.py b/desc/integrals/_interp_utils.py index 75112a56c4..e7bdfe98da 100644 --- a/desc/integrals/_interp_utils.py +++ b/desc/integrals/_interp_utils.py @@ -14,7 +14,7 @@ from orthax.chebyshev import chebroots from desc.backend import dct, jnp, rfft, rfft2, take -from desc.integrals._quad_utils import bijection_from_disc +from desc.integrals.quad_utils import bijection_from_disc from desc.utils import Index, errorif, safediv # TODO (#1154): diff --git a/desc/integrals/basis.py b/desc/integrals/basis.py index 171620f972..9d7e47529f 100644 --- a/desc/integrals/basis.py +++ b/desc/integrals/basis.py @@ -19,7 +19,7 @@ idct_non_uniform, irfft_non_uniform, ) -from desc.integrals._quad_utils import bijection_from_disc, bijection_to_disc +from desc.integrals.quad_utils import bijection_from_disc, bijection_to_disc from desc.io import IOAble from desc.utils import ( atleast_2d_end, diff --git a/desc/integrals/bounce_integral.py b/desc/integrals/bounce_integral.py index 07f6f80f01..b18d607152 100644 --- a/desc/integrals/bounce_integral.py +++ b/desc/integrals/bounce_integral.py @@ -26,14 +26,14 @@ irfft2_non_uniform, polyder_vec, ) -from desc.integrals._quad_utils import ( +from desc.integrals.basis import FourierChebyshevSeries, PiecewiseChebyshevSeries +from desc.integrals.quad_utils import ( automorphism_sin, bijection_from_disc, get_quadrature, grad_automorphism_sin, grad_bijection_from_disc, ) -from desc.integrals.basis import FourierChebyshevSeries, PiecewiseChebyshevSeries from desc.io import IOAble from desc.utils import atleast_nd, errorif, flatten_matrix, setdefault diff --git a/desc/integrals/_quad_utils.py b/desc/integrals/quad_utils.py similarity index 87% rename from desc/integrals/_quad_utils.py rename to desc/integrals/quad_utils.py index 9d563307bf..1ba4c86b91 100644 --- a/desc/integrals/_quad_utils.py +++ b/desc/integrals/quad_utils.py @@ -2,13 +2,15 @@ Notes ----- -Bounce integrals with bounce points where the derivative of |B| does -not vanish have 1/2 power law singularities. The strongly singular integrals -at the local extrema of |B| are not integrable. Hence, everywhere except the -extrema, a Chebyshev or Legendre quadrature under a change of variables works -because √(1−z²) / √(1−λ|B|(z)) ~ g(z, λ) where g(z, λ) is smooth in z. -Empirically, quadratic node clustering near the singularities is sufficient -for estimation of g(z). +Bounce integrals with bounce points where the derivative of |B| does not vanish +have 1/2 power law singularities. However, strongly singular integrals where the +domain of the integral ends at the local extrema of |B| are not integrable. + +Hence, everywhere except for the extrema, an implicit Chebyshev (``chebgauss1`` +or ``chebgauss2`` or modified Legendre quadrature (with ``automorphism_sin``) +captures the integral because √(1−ζ²) / √ (1−λ|B|) ∼ k(λ, ζ) is smooth in ζ. +The clustering of the nodes near the singularities is sufficient to estimate +k(ζ, λ). """ from orthax.chebyshev import chebgauss, chebweight @@ -33,10 +35,6 @@ def grad_bijection_from_disc(a, b): return 0.5 * (b - a) -# This map was tested as a change of variables to interpolate with -# Chebyshev series on a more uniform grid. Although it fit to -# oscillatory functions better, it gives small wiggles due to Runge -# effects near boundary. def automorphism_arcsin(x, gamma=jnp.cos(0.5)): """[-1, 1] ∋ x ↦ y ∈ [−1, 1]. @@ -74,18 +72,16 @@ def grad_automorphism_arcsin(x, gamma=jnp.cos(0.5)): grad_automorphism_arcsin.__doc__ += "\n" + automorphism_arcsin.__doc__ -def automorphism_sin(x, s=0, m=10): +def automorphism_sin(x, m=10): """[-1, 1] ∋ x ↦ y ∈ [−1, 1]. This map increases node density near the boundary by the asymptotic factor - 1/√(1−x²) and adds a cosine factor to the integrand when ``s=0``. + 1/√(1−x²) and adds a cosine factor to the integrand. Parameters ---------- x : jnp.ndarray Points to transform. - s : float - Strength of derivative suppression, s ∈ [0, 1]. m : float Number of machine epsilons used for floating point error buffer. @@ -95,23 +91,16 @@ def automorphism_sin(x, s=0, m=10): Transformed points. """ - errorif(not (0 <= s <= 1)) - # s = 0 -> derivative vanishes like cosine. - # s = 1 -> derivative vanishes like cosine^k. - y0 = jnp.sin(0.5 * jnp.pi * x) - y1 = x + jnp.sin(jnp.pi * x) / jnp.pi # k = 2 - y = (1 - s) * y0 + s * y1 + y = jnp.sin(0.5 * jnp.pi * x) # y is an expansion, so y(x) > x near x ∈ {−1, 1} and there is a tendency # for floating point error to overshoot the true value. eps = m * jnp.finfo(jnp.array(1.0).dtype).eps return jnp.clip(y, -1 + eps, 1 - eps) -def grad_automorphism_sin(x, s=0): +def grad_automorphism_sin(x): """Gradient of sin automorphism.""" - dy0_dx = 0.5 * jnp.pi * jnp.cos(0.5 * jnp.pi * x) - dy1_dx = 1.0 + jnp.cos(jnp.pi * x) - return (1 - s) * dy0_dx + s * dy1_dx + return 0.5 * jnp.pi * jnp.cos(0.5 * jnp.pi * x) grad_automorphism_sin.__doc__ += "\n" + automorphism_sin.__doc__ diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index d164284479..834f153bcf 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -35,6 +35,7 @@ PrincipalCurvature, Volume, ) +from ._neoclassical import EffectiveRipple from ._omnigenity import ( Isodynamicity, Omnigenity, diff --git a/desc/objectives/_neoclassical.py b/desc/objectives/_neoclassical.py new file mode 100644 index 0000000000..f171d81261 --- /dev/null +++ b/desc/objectives/_neoclassical.py @@ -0,0 +1,262 @@ +"""Objectives for targeting neoclassical transport.""" + +import numpy as np + +from desc.compute import get_profiles, get_transforms +from desc.compute.utils import _compute as compute_fun +from desc.grid import LinearGrid +from desc.utils import Timer + +from ..integrals.quad_utils import chebgauss2 +from .objective_funs import _Objective +from .utils import _parse_callable_target_bounds + + +class EffectiveRipple(_Objective): + """The effective ripple is a proxy for neoclassical transport. + + The 3D geometry of the magnetic field in stellarators produces local magnetic + wells that lead to bad confinement properties with enhanced radial drift, + especially for trapped particles. Neoclassical (thermal) transport can become the + dominant transport channel in stellarators which are not optimized to reduce it. + The effective ripple is a proxy, measuring the effective modulation amplitude of the + magnetic field averaged along a magnetic surface, which can be used to optimize for + stellarators with improved confinement. It is targeted as a flux surface function. + + References + ---------- + https://doi.org/10.1063/1.873749. + Evaluation of 1/ν neoclassical transport in stellarators. + V. V. Nemov, S. V. Kasilov, W. Kernbichler, M. F. Heyn. + Phys. Plasmas 1 December 1999; 6 (12): 4622–4632. + + Parameters + ---------- + 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 0. + bounds : tuple of {float, ndarray, callable}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable 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. + weight : {float, ndarray}, optional + Weighting to apply to the Objective, relative to other Objectives. + Must be broadcastable to Objective.dim_f + normalize : bool, optional + This quantity is already normalized so this parameter is ignored. + 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. + rho : ndarray + Unique coordinate values specifying flux surfaces to compute on. + alpha : ndarray + Unique coordinate values specifying field line labels to compute on. + Y_B : int + Number of points per toroidal transit at which to sample data along field + line. Default is 100. + num_transit : int + Number of toroidal transits to follow field line. + For axisymmetric devices, one poloidal transit is sufficient. Otherwise, + more transits will give more accurate result, with diminishing returns. + num_quad : int + Resolution for quadrature of bounce integrals. Default is 32. + num_pitch : int + Resolution for quadrature over velocity coordinate. Default 50. + batch : bool + Whether to vectorize part of the computation. Default is true. + num_well : int + Maximum number of wells to detect for each pitch and field line. + Default is to detect all wells, but due to limitations in JAX this option + may consume more memory. Specifying a number that tightly upper bounds + the number of wells will increase performance. + 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``. + + """ + + _coordinates = "r" + _units = "~" + _print_value_fmt = "Effective ripple ε: " + + def __init__( + self, + eq, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + loss_function=None, + deriv_mode="auto", + *, + rho=1.0, + alpha=0.0, + Y_B=100, + num_transit=10, + num_quad=32, + num_pitch=50, + batch=True, + num_well=None, + name="Effective ripple", + jac_chunk_size=None, + ): + if target is None and bounds is None: + target = 0.0 + + rho, alpha = np.atleast_1d(rho, alpha) + self._dim_f = rho.size + self._keys_1dr = [ + "iota", + "iota_r", + "<|grad(rho)|>", + "min_tz |B|", + "max_tz |B|", + "R0", # TODO: GitHub PR #1094 + ] + self._constants = { + "quad_weights": 1, + "rho": rho, + "alpha": alpha, + "zeta": np.linspace(0, 2 * np.pi * num_transit, Y_B * num_transit), + } + self._constants["quad x"], self._constants["quad w"] = chebgauss2(num_quad) + self._hyperparameters = { + "num_pitch": num_pitch, + "batch": batch, + "num_well": num_well, + } + + super().__init__( + things=eq, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + loss_function=loss_function, + deriv_mode=deriv_mode, + name=name, + jac_chunk_size=jac_chunk_size, + ) + + def build(self, use_jit=True, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + eq = self.things[0] + self._grid_1dr = LinearGrid( + rho=self._constants["rho"], M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym + ) + self._target, self._bounds = _parse_callable_target_bounds( + self._target, self._bounds, self._constants["rho"] + ) + + timer = Timer() + if verbose > 0: + print("Precomputing transforms") + timer.start("Precomputing transforms") + + self._constants["transforms_1dr"] = get_transforms( + self._keys_1dr, eq, self._grid_1dr + ) + self._constants["profiles"] = get_profiles( + self._keys_1dr + ["effective ripple"], eq, self._grid_1dr + ) + + timer.stop("Precomputing transforms") + if verbose > 1: + timer.disp("Precomputing transforms") + + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params, constants=None): + """Compute the effective ripple. + + Parameters + ---------- + params : dict + Dictionary of equilibrium degrees of freedom, e.g. + ``Equilibrium.params_dict`` + constants : dict + Dictionary of constant data, e.g. transforms, profiles etc. + Defaults to ``self.constants``. + + Returns + ------- + result : ndarray + Effective ripple as a function of the flux surface label. + + """ + if constants is None: + constants = self.constants + eq = self.things[0] + # TODO: compute all deps of effective ripple here + data = compute_fun( + eq, + self._keys_1dr, + params, + constants["transforms_1dr"], + constants["profiles"], + ) + # TODO: interpolate all deps to this grid with fft utilities from fourier bounce + grid = eq._get_rtz_grid( + constants["rho"], + constants["alpha"], + constants["zeta"], + coordinates="raz", + iota=self._grid_1dr.compress(data["iota"]), + params=params, + ) + data = { + key: ( + grid.copy_data_from_other(data[key], self._grid_1dr) + if key != "R0" + else data[key] + ) + for key in self._keys_1dr + } + data = compute_fun( + eq, + "effective ripple", + params, + get_transforms("effective ripple", eq, grid, jitable=True), + constants["profiles"], + data=data, + quad=(constants["quad x"], constants["quad w"]), + **self._hyperparameters, + ) + return grid.compress(data["effective ripple"]) diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index a643a6a585..3ff6a6aa1d 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -597,7 +597,7 @@ def compute(self, params, constants=None): "a": len_data["a"], } - grid = eq.get_rtz_grid( + grid = eq._get_rtz_grid( rho, alpha, zeta, diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 1c8817a880..537ede2f1e 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -624,11 +624,11 @@ class FixBoundaryR(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Rb_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Rb_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -705,11 +705,11 @@ class FixBoundaryZ(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Zb_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Zb_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -828,7 +828,7 @@ class FixThetaSFL(FixParameters): Equilibrium that will be optimized to satisfy the Objective. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Has no effect for this objective. normalize_target : bool, optional @@ -873,11 +873,11 @@ class FixAxisR(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Ra_n``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Ra_n``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -954,11 +954,11 @@ class FixAxisZ(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Za_n``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Za_n``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -1035,11 +1035,11 @@ class FixModeR(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.R_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.R_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -1116,11 +1116,11 @@ class FixModeZ(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Z_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Z_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -1198,7 +1198,7 @@ class FixModeLambda(FixParameters): Defaults to ``target=eq.L_lmn``. bounds : tuple, optional Lower and upper bounds on the objective. Overrides target. - Both bounds must be broadcastable to to Objective.dim_f. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.L_lmn``. weight : float, ndarray, optional Weighting to apply to the Objective, relative to other Objectives. @@ -1260,11 +1260,11 @@ class FixSumModesR(_FixedObjective): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.R_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.R_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -1426,11 +1426,11 @@ class FixSumModesZ(_FixedObjective): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Z_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Z_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -1593,11 +1593,11 @@ class FixSumModesLambda(_FixedObjective): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.L_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.L_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f. + Must be broadcastable to Objective.dim_f. normalize : bool, optional Has no effect for this objective. normalize_target : bool, optional @@ -1762,11 +1762,11 @@ class FixPressure(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.p_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.p_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -1844,11 +1844,11 @@ class FixAnisotropy(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.a_lmn``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.a_lmn``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Has no effect for this objective. normalize_target : bool, optional @@ -1921,11 +1921,11 @@ class FixIota(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.i_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.i_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Has no effect for this objective. normalize_target : bool, optional @@ -1998,11 +1998,11 @@ class FixCurrent(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.c_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.c_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -2080,11 +2080,11 @@ class FixElectronTemperature(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Te_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Te_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -2162,11 +2162,11 @@ class FixElectronDensity(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.ne_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.ne_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -2246,11 +2246,11 @@ class FixIonTemperature(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Ti_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Ti_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -2328,11 +2328,11 @@ class FixAtomicNumber(FixParameters): Must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Zeff_l``. 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. + Both bounds must be broadcastable to Objective.dim_f. Defaults to ``target=eq.Zeff_l``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Has no effect for this objective. normalize_target : bool, optional @@ -2405,11 +2405,11 @@ class FixPsi(FixParameters): Must be broadcastable to Objective.dim_f. Default is ``target=eq.Psi``. 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. + Both bounds must be broadcastable to Objective.dim_f. Default is ``target=eq.Psi``. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -2475,10 +2475,10 @@ class FixCurveShift(FixParameters): 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 + Both bounds must be broadcastable 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 + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional @@ -2527,10 +2527,10 @@ class FixCurveRotation(FixParameters): 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 + Both bounds must be broadcastable 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 + Must be broadcastable to Objective.dim_f normalize : bool, optional Has no effect for this objective. normalize_target : bool, optional @@ -3053,11 +3053,11 @@ class FixSheetCurrent(FixParameters): Defaults to the equilibrium sheet current parameters. 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. + Both bounds must be broadcastable to Objective.dim_f. Default is to use target. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. - Must be broadcastable to to Objective.dim_f + Must be broadcastable to Objective.dim_f normalize : bool, optional Whether to compute the error in physical units or non-dimensionalize. normalize_target : bool, optional diff --git a/desc/plotting.py b/desc/plotting.py index 6a59bc0345..096afd618b 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -495,13 +495,21 @@ def plot_1d(eq, name, grid=None, log=False, ax=None, return_data=False, **kwargs log=log, ax=ax, return_data=return_data, + grid=grid, **kwargs, ) rho = grid.nodes[:, 0] if not np.all(np.isclose(rho, rho[0])): # rho nodes are not constant, so user must be plotting against rho return plot_fsa( - eq, name, rho=rho, log=log, ax=ax, return_data=return_data, **kwargs + eq, + name, + rho=rho, + log=log, + ax=ax, + return_data=return_data, + grid=grid, + **kwargs, ) elif data_index[parameterization][name]["coordinates"] == "s": # curve qtys @@ -1111,6 +1119,7 @@ def plot_fsa( # noqa: C901 norm_F=False, ax=None, return_data=False, + grid=None, **kwargs, ): """Plot flux surface averages of quantities. @@ -1148,7 +1157,10 @@ def plot_fsa( # noqa: C901 ax : matplotlib AxesSubplot, optional Axis to plot on. return_data : bool - If True, return the data plotted as well as fig,ax + if True, return the data plotted as well as fig,ax + grid : _Grid + Grid to compute name on. If provided, the parameters + ``rho``, ``M``, and ``N`` are ignored. **kwargs : dict, optional Specify properties of the figure, axis, and plot appearance e.g.:: @@ -1186,22 +1198,24 @@ def plot_fsa( # noqa: C901 fig, ax = plot_fsa(eq, "B_theta", with_sqrt_g=False) """ - if np.isscalar(rho) and (int(rho) == rho): - rho = np.linspace(0, 1, rho + 1) - rho = np.atleast_1d(rho) if M is None: M = eq.M_grid if N is None: N = eq.N_grid + if grid is None: + if np.isscalar(rho) and (int(rho) == rho): + rho = np.linspace(0, 1, rho + 1) + rho = np.atleast_1d(rho) + grid = LinearGrid(M=M, N=N, NFP=eq.NFP, sym=eq.sym, rho=rho) + else: + rho = grid.compress(grid.nodes[:, 0]) + linecolor = kwargs.pop("linecolor", colorblind_colors[0]) ls = kwargs.pop("ls", "-") lw = kwargs.pop("lw", 1) fig, ax = _format_ax(ax, figsize=kwargs.pop("figsize", (4, 4))) label = kwargs.pop("label", None) - - grid = LinearGrid(M=M, N=N, NFP=eq.NFP, rho=rho) - p = "desc.equilibrium.equilibrium.Equilibrium" if "<" + name + ">" in data_index[p]: # If we identify the quantity to plot as something in data_index, then diff --git a/desc/utils.py b/desc/utils.py index 5b995f1658..134108c9c4 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -866,82 +866,6 @@ def safediv(a, b, fill=0, threshold=0): return num / den -def cumtrapz(y, x=None, dx=1.0, axis=-1, initial=None): - """Cumulatively integrate y(x) using the composite trapezoidal rule. - - Taken from SciPy, but changed NumPy references to JAX.NumPy: - https://github.com/scipy/scipy/blob/v1.10.1/scipy/integrate/_quadrature.py - - Parameters - ---------- - y : array_like - Values to integrate. - x : array_like, optional - The coordinate to integrate along. If None (default), use spacing `dx` - between consecutive elements in `y`. - dx : float, optional - Spacing between elements of `y`. Only used if `x` is None. - axis : int, optional - Specifies the axis to cumulate. Default is -1 (last axis). - initial : scalar, optional - If given, insert this value at the beginning of the returned result. - Typically, this value should be 0. Default is None, which means no - value at ``x[0]`` is returned and `res` has one element less than `y` - along the axis of integration. - - Returns - ------- - res : ndarray - The result of cumulative integration of `y` along `axis`. - If `initial` is None, the shape is such that the axis of integration - has one less value than `y`. If `initial` is given, the shape is equal - to that of `y`. - - """ - y = jnp.asarray(y) - if x is None: - d = dx - else: - x = jnp.asarray(x) - if x.ndim == 1: - d = jnp.diff(x) - # reshape to correct shape - shape = [1] * y.ndim - shape[axis] = -1 - d = d.reshape(shape) - elif len(x.shape) != len(y.shape): - raise ValueError("If given, shape of x must be 1-D or the " "same as y.") - else: - d = jnp.diff(x, axis=axis) - - if d.shape[axis] != y.shape[axis] - 1: - raise ValueError( - "If given, length of x along axis must be the " "same as y." - ) - - def tupleset(t, i, value): - l = list(t) - l[i] = value - return tuple(l) - - nd = len(y.shape) - slice1 = tupleset((slice(None),) * nd, axis, slice(1, None)) - slice2 = tupleset((slice(None),) * nd, axis, slice(None, -1)) - res = jnp.cumsum(d * (y[slice1] + y[slice2]) / 2.0, axis=axis) - - if initial is not None: - if not jnp.isscalar(initial): - raise ValueError("`initial` parameter should be a scalar.") - - shape = list(res.shape) - shape[axis] = 1 - res = jnp.concatenate( - [jnp.full(shape, initial, dtype=res.dtype), res], axis=axis - ) - - return res - - def ensure_tuple(x): """Returns x as a tuple of arrays.""" if isinstance(x, tuple): diff --git a/docs/adding_objectives.rst b/docs/adding_objectives.rst index 2fa3028bad..b7572410eb 100644 --- a/docs/adding_objectives.rst +++ b/docs/adding_objectives.rst @@ -192,6 +192,7 @@ A full example objective with comments describing the key points is given below: # and to make the objective value independent of grid resolution. return f + Converting to Cartesian coordinates ----------------------------------- diff --git a/tests/baseline/test_effective_ripple.png b/tests/baseline/test_effective_ripple.png new file mode 100644 index 0000000000..13a13596ff Binary files /dev/null and b/tests/baseline/test_effective_ripple.png differ diff --git a/tests/baseline/test_fsa_F_normalized.png b/tests/baseline/test_fsa_F_normalized.png index 023bb31f84..6a7fdaa3c7 100644 Binary files a/tests/baseline/test_fsa_F_normalized.png and b/tests/baseline/test_fsa_F_normalized.png differ diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 55c5be23a6..3fa0725f44 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/inputs/neo_out.w7x b/tests/inputs/neo_out.w7x new file mode 100644 index 0000000000..7ef20fe33d --- /dev/null +++ b/tests/inputs/neo_out.w7x @@ -0,0 +1,255 @@ + 2 0.5356233042E-03 -0.2775278088E-01 0.8561233823E+00 0.2833190119E+01 0.5624358558E+01 + 3 0.5378265865E-03 -0.5978662460E-01 0.8562776927E+00 0.2844342240E+01 0.5624358558E+01 + 4 0.5380758777E-03 -0.8459073834E-01 0.8564341254E+00 0.2852391793E+01 0.5624358558E+01 + 5 0.5373904091E-03 -0.1055463924E+00 0.8565926871E+00 0.2859155537E+01 0.5624358558E+01 + 6 0.5361352617E-03 -0.1240209068E+00 0.8567533842E+00 0.2865161902E+01 0.5624358558E+01 + 7 0.5337385233E-03 -0.1407260454E+00 0.8569162231E+00 0.2870652755E+01 0.5624358558E+01 + 8 0.5311366833E-03 -0.1560869773E+00 0.8570812100E+00 0.2875762845E+01 0.5624358558E+01 + 9 0.5281154419E-03 -0.1703822422E+00 0.8572483509E+00 0.2880576371E+01 0.5624358558E+01 + 10 0.5251391126E-03 -0.1838059447E+00 0.8574176519E+00 0.2885150123E+01 0.5624358558E+01 + 11 0.5219759630E-03 -0.1964989169E+00 0.8575891189E+00 0.2889524557E+01 0.5624358558E+01 + 12 0.5177741368E-03 -0.2085682680E+00 0.8577627577E+00 0.2893729707E+01 0.5624358558E+01 + 13 0.5141285674E-03 -0.2200970868E+00 0.8579385740E+00 0.2897788605E+01 0.5624358558E+01 + 14 0.5098369670E-03 -0.2311514760E+00 0.8581165733E+00 0.2901719379E+01 0.5624358558E+01 + 15 0.5054117219E-03 -0.2417850737E+00 0.8582967611E+00 0.2905536603E+01 0.5624358558E+01 + 16 0.5015357241E-03 -0.2520421053E+00 0.8584791428E+00 0.2909252209E+01 0.5624358558E+01 + 17 0.4963671801E-03 -0.2619595601E+00 0.8586637237E+00 0.2912876111E+01 0.5624358558E+01 + 18 0.4922156466E-03 -0.2715686782E+00 0.8588505087E+00 0.2916416654E+01 0.5624358558E+01 + 19 0.4882789406E-03 -0.2808962483E+00 0.8590395031E+00 0.2919880943E+01 0.5624358558E+01 + 20 0.4834425367E-03 -0.2899653516E+00 0.8592307118E+00 0.2923275086E+01 0.5624358558E+01 + 21 0.4789018456E-03 -0.2987960912E+00 0.8594241394E+00 0.2926604383E+01 0.5624358558E+01 + 22 0.4748425789E-03 -0.3074060510E+00 0.8596197909E+00 0.2929873463E+01 0.5624358558E+01 + 23 0.4707973427E-03 -0.3158107614E+00 0.8598176707E+00 0.2933086403E+01 0.5624358558E+01 + 24 0.4666681596E-03 -0.3240240643E+00 0.8600177833E+00 0.2936246811E+01 0.5624358558E+01 + 25 0.4627643486E-03 -0.3320583206E+00 0.8602201332E+00 0.2939357899E+01 0.5624358558E+01 + 26 0.4589237636E-03 -0.3399244124E+00 0.8604247246E+00 0.2942422544E+01 0.5624358558E+01 + 27 0.4556504343E-03 -0.3476323895E+00 0.8606315617E+00 0.2945443329E+01 0.5624358558E+01 + 28 0.4522927489E-03 -0.3551912573E+00 0.8608406485E+00 0.2948422589E+01 0.5624358558E+01 + 29 0.4478992816E-03 -0.3626092045E+00 0.8610519890E+00 0.2951362437E+01 0.5624358558E+01 + 30 0.4442751985E-03 -0.3698936718E+00 0.8612655871E+00 0.2954264796E+01 0.5624358558E+01 + 31 0.4406778763E-03 -0.3770514694E+00 0.8614814465E+00 0.2957131421E+01 0.5624358558E+01 + 32 0.4382169174E-03 -0.3840888674E+00 0.8616995708E+00 0.2959963919E+01 0.5624358558E+01 + 33 0.4341665514E-03 -0.3910115536E+00 0.8619199635E+00 0.2962763763E+01 0.5624358558E+01 + 34 0.4304384714E-03 -0.3978250253E+00 0.8621426281E+00 0.2965532314E+01 0.5624358558E+01 + 35 0.4283073823E-03 -0.4045342219E+00 0.8623675678E+00 0.2968270823E+01 0.5624358558E+01 + 36 0.4241997737E-03 -0.4111433741E+00 0.8625947858E+00 0.2970980453E+01 0.5624358558E+01 + 37 0.4208176604E-03 -0.4176568798E+00 0.8628242853E+00 0.2973662281E+01 0.5624358558E+01 + 38 0.4187106713E-03 -0.4240786787E+00 0.8630560692E+00 0.2976317307E+01 0.5624358558E+01 + 39 0.4149200882E-03 -0.4304124740E+00 0.8632901404E+00 0.2978946466E+01 0.5624358558E+01 + 40 0.4132104943E-03 -0.4366617115E+00 0.8635265017E+00 0.2981550630E+01 0.5624358558E+01 + 41 0.4090812424E-03 -0.4428295705E+00 0.8637651556E+00 0.2984130616E+01 0.5624358558E+01 + 42 0.4068495089E-03 -0.4489190938E+00 0.8640061048E+00 0.2986687192E+01 0.5624358558E+01 + 43 0.4036889011E-03 -0.4549331540E+00 0.8642493517E+00 0.2989221077E+01 0.5624358558E+01 + 44 0.4018150880E-03 -0.4608743565E+00 0.8644948986E+00 0.2991732952E+01 0.5624358558E+01 + 45 0.3988982082E-03 -0.4667452354E+00 0.8647427478E+00 0.2994223457E+01 0.5624358558E+01 + 46 0.3974684750E-03 -0.4725481796E+00 0.8649929013E+00 0.2996693201E+01 0.5624358558E+01 + 47 0.3943081016E-03 -0.4782853377E+00 0.8652453612E+00 0.2999142758E+01 0.5624358558E+01 + 48 0.3924560018E-03 -0.4839590644E+00 0.8655001293E+00 0.3001572674E+01 0.5624358558E+01 + 49 0.3891919985E-03 -0.4895711695E+00 0.8657572075E+00 0.3003983468E+01 0.5624358558E+01 + 50 0.3879586209E-03 -0.4951235945E+00 0.8660165975E+00 0.3006375636E+01 0.5624358558E+01 + 51 0.3853550624E-03 -0.5006181735E+00 0.8662783007E+00 0.3008749650E+01 0.5624358558E+01 + 52 0.3838310158E-03 -0.5060565330E+00 0.8665423187E+00 0.3011105959E+01 0.5624358558E+01 + 53 0.3831972104E-03 -0.5114404010E+00 0.8668086529E+00 0.3013444996E+01 0.5624358558E+01 + 54 0.3797540348E-03 -0.5167714133E+00 0.8670773044E+00 0.3015767172E+01 0.5624358558E+01 + 55 0.3787660229E-03 -0.5220509881E+00 0.8673482745E+00 0.3018072883E+01 0.5624358558E+01 + 56 0.3768645085E-03 -0.5272804776E+00 0.8676215641E+00 0.3020362508E+01 0.5624358558E+01 + 57 0.3748091098E-03 -0.5324612810E+00 0.8678971742E+00 0.3022636412E+01 0.5624358558E+01 + 58 0.3743022428E-03 -0.5375947132E+00 0.8681751055E+00 0.3024894943E+01 0.5624358558E+01 + 59 0.3721917347E-03 -0.5426820056E+00 0.8684553589E+00 0.3027138438E+01 0.5624358558E+01 + 60 0.3709924943E-03 -0.5477243136E+00 0.8687379349E+00 0.3029367221E+01 0.5624358558E+01 + 61 0.3704936263E-03 -0.5527227800E+00 0.8690228340E+00 0.3031581603E+01 0.5624358558E+01 + 62 0.3695075285E-03 -0.5576783420E+00 0.8693100565E+00 0.3033781883E+01 0.5624358558E+01 + 63 0.3673241652E-03 -0.5625923593E+00 0.8695996029E+00 0.3035968350E+01 0.5624358558E+01 + 64 0.3663927124E-03 -0.5674658823E+00 0.8698914731E+00 0.3038141283E+01 0.5624358558E+01 + 65 0.3669163148E-03 -0.5722994245E+00 0.8701856673E+00 0.3040300950E+01 0.5624358558E+01 + 66 0.3643615641E-03 -0.5770941553E+00 0.8704821855E+00 0.3042447611E+01 0.5624358558E+01 + 67 0.3638725081E-03 -0.5818509555E+00 0.8707810274E+00 0.3044581514E+01 0.5624358558E+01 + 68 0.3637155270E-03 -0.5865707132E+00 0.8710821929E+00 0.3046702903E+01 0.5624358558E+01 + 69 0.3624399768E-03 -0.5912541997E+00 0.8713856815E+00 0.3048812011E+01 0.5624358558E+01 + 70 0.3609036507E-03 -0.5959022078E+00 0.8716914927E+00 0.3050909061E+01 0.5624358558E+01 + 71 0.3608240656E-03 -0.6005154815E+00 0.8719996260E+00 0.3052994272E+01 0.5624358558E+01 + 72 0.3608141853E-03 -0.6050948371E+00 0.8723100807E+00 0.3055067855E+01 0.5624358558E+01 + 73 0.3591051612E-03 -0.6096409580E+00 0.8726228559E+00 0.3057130013E+01 0.5624358558E+01 + 74 0.3582191664E-03 -0.6141544925E+00 0.8729379508E+00 0.3059180942E+01 0.5624358558E+01 + 75 0.3577567504E-03 -0.6186361906E+00 0.8732553643E+00 0.3061220835E+01 0.5624358558E+01 + 76 0.3582227512E-03 -0.6230866394E+00 0.8735750954E+00 0.3063249875E+01 0.5624358558E+01 + 77 0.3573796359E-03 -0.6275064573E+00 0.8738971426E+00 0.3065268244E+01 0.5624358558E+01 + 78 0.3561820616E-03 -0.6318962415E+00 0.8742215048E+00 0.3067276114E+01 0.5624358558E+01 + 79 0.3559753002E-03 -0.6362566014E+00 0.8745481804E+00 0.3069273655E+01 0.5624358558E+01 + 80 0.3566872337E-03 -0.6405880028E+00 0.8748771680E+00 0.3071261030E+01 0.5624358558E+01 + 81 0.3561926951E-03 -0.6448912634E+00 0.8752084657E+00 0.3073238400E+01 0.5624358558E+01 + 82 0.3556456202E-03 -0.6491665944E+00 0.8755420719E+00 0.3075205919E+01 0.5624358558E+01 + 83 0.3549734067E-03 -0.6534146033E+00 0.8758779847E+00 0.3077163737E+01 0.5624358558E+01 + 84 0.3551126264E-03 -0.6576357940E+00 0.8762162020E+00 0.3079112000E+01 0.5624358558E+01 + 85 0.3545857299E-03 -0.6618307135E+00 0.8765567218E+00 0.3081050852E+01 0.5624358558E+01 + 86 0.3553357957E-03 -0.6659997637E+00 0.8768995418E+00 0.3082980430E+01 0.5624358558E+01 + 87 0.3559094263E-03 -0.6701434620E+00 0.8772446597E+00 0.3084900869E+01 0.5624358558E+01 + 88 0.3554084603E-03 -0.6742621718E+00 0.8775920731E+00 0.3086812301E+01 0.5624358558E+01 + 89 0.3549117246E-03 -0.6783563575E+00 0.8779417795E+00 0.3088714853E+01 0.5624358558E+01 + 90 0.3553223939E-03 -0.6824263935E+00 0.8782937762E+00 0.3090608651E+01 0.5624358558E+01 + 91 0.3554592231E-03 -0.6864727591E+00 0.8786480605E+00 0.3092493817E+01 0.5624358558E+01 + 92 0.3556951502E-03 -0.6904958246E+00 0.8790046294E+00 0.3094370468E+01 0.5624358558E+01 + 93 0.3566066727E-03 -0.6944959398E+00 0.8793634801E+00 0.3096238721E+01 0.5624358558E+01 + 94 0.3573775084E-03 -0.6984735115E+00 0.8797246094E+00 0.3098098688E+01 0.5624358558E+01 + 95 0.3568009062E-03 -0.7024288864E+00 0.8800880142E+00 0.3099950480E+01 0.5624358558E+01 + 96 0.3569803376E-03 -0.7063624006E+00 0.8804536912E+00 0.3101794204E+01 0.5624358558E+01 + 97 0.3571244371E-03 -0.7102744416E+00 0.8808216370E+00 0.3103629965E+01 0.5624358558E+01 + 98 0.3576032718E-03 -0.7141653068E+00 0.8811918480E+00 0.3105457866E+01 0.5624358558E+01 + 99 0.3584656126E-03 -0.7180353475E+00 0.8815643206E+00 0.3107278006E+01 0.5624358558E+01 + 100 0.3589762422E-03 -0.7218845886E+00 0.8819390512E+00 0.3109090483E+01 0.5624358558E+01 + 101 0.3603525044E-03 -0.7257139097E+00 0.8823160358E+00 0.3110895392E+01 0.5624358558E+01 + 102 0.3614623113E-03 -0.7295235710E+00 0.8826952706E+00 0.3112692827E+01 0.5624358558E+01 + 103 0.3606437046E-03 -0.7333133781E+00 0.8830767515E+00 0.3114482878E+01 0.5624358558E+01 + 104 0.3612289154E-03 -0.7370838635E+00 0.8834604743E+00 0.3116265634E+01 0.5624358558E+01 + 105 0.3617325738E-03 -0.7408353164E+00 0.8838464348E+00 0.3118041181E+01 0.5624358558E+01 + 106 0.3627296406E-03 -0.7445679954E+00 0.8842346285E+00 0.3119809604E+01 0.5624358558E+01 + 107 0.3636781394E-03 -0.7482821854E+00 0.8846250511E+00 0.3121570985E+01 0.5624358558E+01 + 108 0.3645861215E-03 -0.7519781653E+00 0.8850176978E+00 0.3123325406E+01 0.5624358558E+01 + 109 0.3660372381E-03 -0.7556561504E+00 0.8854125641E+00 0.3125072944E+01 0.5624358558E+01 + 110 0.3673810312E-03 -0.7593163796E+00 0.8858096451E+00 0.3126813676E+01 0.5624358558E+01 + 111 0.3686426742E-03 -0.7629591275E+00 0.8862089359E+00 0.3128547678E+01 0.5624358558E+01 + 112 0.3682077515E-03 -0.7665846220E+00 0.8866104314E+00 0.3130275022E+01 0.5624358558E+01 + 113 0.3692150443E-03 -0.7701930821E+00 0.8870141266E+00 0.3131995781E+01 0.5624358558E+01 + 114 0.3701971195E-03 -0.7737847699E+00 0.8874200161E+00 0.3133710023E+01 0.5624358558E+01 + 115 0.3711684254E-03 -0.7773598977E+00 0.8878280947E+00 0.3135417818E+01 0.5624358558E+01 + 116 0.3723769756E-03 -0.7809186822E+00 0.8882383569E+00 0.3137119231E+01 0.5624358558E+01 + 117 0.3737720685E-03 -0.7844612812E+00 0.8886507971E+00 0.3138814329E+01 0.5624358558E+01 + 118 0.3747879323E-03 -0.7879881212E+00 0.8890654096E+00 0.3140503174E+01 0.5624358558E+01 + 119 0.3767781682E-03 -0.7914991352E+00 0.8894821887E+00 0.3142185829E+01 0.5624358558E+01 + 120 0.3777747813E-03 -0.7949946684E+00 0.8899011285E+00 0.3143862355E+01 0.5624358558E+01 + 121 0.3796587687E-03 -0.7984748194E+00 0.8903222229E+00 0.3145532810E+01 0.5624358558E+01 + 122 0.3798054048E-03 -0.8019398699E+00 0.8907454659E+00 0.3147197254E+01 0.5624358558E+01 + 123 0.3802716899E-03 -0.8053900151E+00 0.8911708513E+00 0.3148855741E+01 0.5624358558E+01 + 124 0.3822902226E-03 -0.8088253590E+00 0.8915983727E+00 0.3150508329E+01 0.5624358558E+01 + 125 0.3838595498E-03 -0.8122461333E+00 0.8920280237E+00 0.3152155070E+01 0.5624358558E+01 + 126 0.3847126325E-03 -0.8156525096E+00 0.8924597978E+00 0.3153796019E+01 0.5624358558E+01 + 127 0.3867526686E-03 -0.8190446578E+00 0.8928936883E+00 0.3155431226E+01 0.5624358558E+01 + 128 0.3884499806E-03 -0.8224227461E+00 0.8933296884E+00 0.3157060742E+01 0.5624358558E+01 + 129 0.3899816280E-03 -0.8257869406E+00 0.8937677914E+00 0.3158684617E+01 0.5624358558E+01 + 130 0.3919113564E-03 -0.8291374843E+00 0.8942079901E+00 0.3160302899E+01 0.5624358558E+01 + 131 0.3934145947E-03 -0.8324743718E+00 0.8946502777E+00 0.3161915636E+01 0.5624358558E+01 + 132 0.3952829602E-03 -0.8357978904E+00 0.8950946468E+00 0.3163522873E+01 0.5624358558E+01 + 133 0.3974268481E-03 -0.8391081517E+00 0.8955410901E+00 0.3165124656E+01 0.5624358558E+01 + 134 0.3972871834E-03 -0.8424053174E+00 0.8959896003E+00 0.3166721029E+01 0.5624358558E+01 + 135 0.3987657387E-03 -0.8456895256E+00 0.8964401699E+00 0.3168312037E+01 0.5624358558E+01 + 136 0.4008021684E-03 -0.8489609104E+00 0.8968927912E+00 0.3169897720E+01 0.5624358558E+01 + 137 0.4027721175E-03 -0.8522195737E+00 0.8973474564E+00 0.3171478122E+01 0.5624358558E+01 + 138 0.4041412863E-03 -0.8554657844E+00 0.8978041579E+00 0.3173053282E+01 0.5624358558E+01 + 139 0.4062959375E-03 -0.8586996027E+00 0.8982628876E+00 0.3174623242E+01 0.5624358558E+01 + 140 0.4082679605E-03 -0.8619211708E+00 0.8987236374E+00 0.3176188039E+01 0.5624358558E+01 + 141 0.4104778940E-03 -0.8651306292E+00 0.8991863992E+00 0.3177747712E+01 0.5624358558E+01 + 142 0.4124517163E-03 -0.8683280773E+00 0.8996511648E+00 0.3179302298E+01 0.5624358558E+01 + 143 0.4140767243E-03 -0.8715137795E+00 0.9001179257E+00 0.3180851834E+01 0.5624358558E+01 + 144 0.4162220581E-03 -0.8746876572E+00 0.9005866735E+00 0.3182396356E+01 0.5624358558E+01 + 145 0.4187761988E-03 -0.8778499342E+00 0.9010573996E+00 0.3183935898E+01 0.5624358558E+01 + 146 0.4210669326E-03 -0.8810007249E+00 0.9015300952E+00 0.3185470496E+01 0.5624358558E+01 + 147 0.4228995579E-03 -0.8841401657E+00 0.9020047517E+00 0.3187000182E+01 0.5624358558E+01 + 148 0.4241207015E-03 -0.8872683572E+00 0.9024813599E+00 0.3188524988E+01 0.5624358558E+01 + 149 0.4263793416E-03 -0.8903854219E+00 0.9029599110E+00 0.3190044948E+01 0.5624358558E+01 + 150 0.4279659591E-03 -0.8934914825E+00 0.9034403958E+00 0.3191560093E+01 0.5624358558E+01 + 151 0.4299169993E-03 -0.8965866201E+00 0.9039228050E+00 0.3193070452E+01 0.5624358558E+01 + 152 0.4320017504E-03 -0.8996709885E+00 0.9044071293E+00 0.3194576055E+01 0.5624358558E+01 + 153 0.4346783163E-03 -0.9027446574E+00 0.9048933592E+00 0.3196076933E+01 0.5624358558E+01 + 154 0.4372527954E-03 -0.9058077845E+00 0.9053814852E+00 0.3197573112E+01 0.5624358558E+01 + 155 0.4392481637E-03 -0.9088604068E+00 0.9058714976E+00 0.3199064622E+01 0.5624358558E+01 + 156 0.4415898156E-03 -0.9119026863E+00 0.9063633866E+00 0.3200551488E+01 0.5624358558E+01 + 157 0.4448743558E-03 -0.9149346784E+00 0.9068571424E+00 0.3202033738E+01 0.5624358558E+01 + 158 0.4472083682E-03 -0.9179565262E+00 0.9073527548E+00 0.3203511397E+01 0.5624358558E+01 + 159 0.4496999602E-03 -0.9209682771E+00 0.9078502138E+00 0.3204984489E+01 0.5624358558E+01 + 160 0.4505319881E-03 -0.9239692540E+00 0.9083495093E+00 0.3206453041E+01 0.5624358558E+01 + 161 0.4544594233E-03 -0.9269609089E+00 0.9088506308E+00 0.3207917075E+01 0.5624358558E+01 + 162 0.4583879881E-03 -0.9299433552E+00 0.9093535679E+00 0.3209376615E+01 0.5624358558E+01 + 163 0.4610490526E-03 -0.9329167272E+00 0.9098583101E+00 0.3210831682E+01 0.5624358558E+01 + 164 0.4625361161E-03 -0.9358796200E+00 0.9103648468E+00 0.3212282300E+01 0.5624358558E+01 + 165 0.4647084157E-03 -0.9388329504E+00 0.9108731671E+00 0.3213728490E+01 0.5624358558E+01 + 166 0.4672011185E-03 -0.9417768875E+00 0.9113832603E+00 0.3215170272E+01 0.5624358558E+01 + 167 0.4697908127E-03 -0.9447114815E+00 0.9118951153E+00 0.3216607667E+01 0.5624358558E+01 + 168 0.4729648670E-03 -0.9476368358E+00 0.9124087211E+00 0.3218040695E+01 0.5624358558E+01 + 169 0.4758414663E-03 -0.9505530292E+00 0.9129240665E+00 0.3219469375E+01 0.5624358558E+01 + 170 0.4785972725E-03 -0.9534601548E+00 0.9134411401E+00 0.3220893725E+01 0.5624358558E+01 + 171 0.4815977141E-03 -0.9563582847E+00 0.9139599306E+00 0.3222313765E+01 0.5624358558E+01 + 172 0.4842961457E-03 -0.9592474967E+00 0.9144804265E+00 0.3223729512E+01 0.5624358558E+01 + 173 0.4870214865E-03 -0.9621278819E+00 0.9150026161E+00 0.3225140983E+01 0.5624358558E+01 + 174 0.4906080843E-03 -0.9649994892E+00 0.9155264877E+00 0.3226548195E+01 0.5624358558E+01 + 175 0.4935061742E-03 -0.9678624282E+00 0.9160520295E+00 0.3227951166E+01 0.5624358558E+01 + 176 0.4962874545E-03 -0.9707168038E+00 0.9165792295E+00 0.3229349911E+01 0.5624358558E+01 + 177 0.4985589053E-03 -0.9735626343E+00 0.9171080757E+00 0.3230744447E+01 0.5624358558E+01 + 178 0.5032819962E-03 -0.9763999994E+00 0.9176385559E+00 0.3232134788E+01 0.5624358558E+01 + 179 0.5066961255E-03 -0.9792290969E+00 0.9181706578E+00 0.3233520951E+01 0.5624358558E+01 + 180 0.5090257778E-03 -0.9820497822E+00 0.9187043690E+00 0.3234902950E+01 0.5624358558E+01 + 181 0.5125869437E-03 -0.9848622215E+00 0.9192396771E+00 0.3236280800E+01 0.5624358558E+01 + 182 0.5159541313E-03 -0.9876665063E+00 0.9197765695E+00 0.3237654517E+01 0.5624358558E+01 + 183 0.5187009734E-03 -0.9904626940E+00 0.9203150334E+00 0.3239024114E+01 0.5624358558E+01 + 184 0.5222613153E-03 -0.9932509707E+00 0.9208550561E+00 0.3240389607E+01 0.5624358558E+01 + 185 0.5254607296E-03 -0.9960311723E+00 0.9213966247E+00 0.3241751008E+01 0.5624358558E+01 + 186 0.5281250998E-03 -0.9988034953E+00 0.9219397260E+00 0.3243108334E+01 0.5624358558E+01 + 187 0.5319364330E-03 -0.1001567982E+01 0.9224843470E+00 0.3244461598E+01 0.5624358558E+01 + 188 0.5361050736E-03 -0.1004324695E+01 0.9230304744E+00 0.3245810814E+01 0.5624358558E+01 + 189 0.5396822931E-03 -0.1007073711E+01 0.9235780949E+00 0.3247155998E+01 0.5624358558E+01 + 190 0.5427395729E-03 -0.1009815099E+01 0.9241271949E+00 0.3248497163E+01 0.5624358558E+01 + 191 0.5463932260E-03 -0.1012548910E+01 0.9246777611E+00 0.3249834324E+01 0.5624358558E+01 + 192 0.5490555266E-03 -0.1015275206E+01 0.9252297795E+00 0.3251167497E+01 0.5624358558E+01 + 193 0.5544139873E-03 -0.1017994188E+01 0.9257832366E+00 0.3252496695E+01 0.5624358558E+01 + 194 0.5583218996E-03 -0.1020705697E+01 0.9263381183E+00 0.3253821936E+01 0.5624358558E+01 + 195 0.5623483713E-03 -0.1023409847E+01 0.9268944107E+00 0.3255143235E+01 0.5624358558E+01 + 196 0.5659722353E-03 -0.1026106688E+01 0.9274520997E+00 0.3256460607E+01 0.5624358558E+01 + 197 0.5693420391E-03 -0.1028796338E+01 0.9280111710E+00 0.3257774070E+01 0.5624358558E+01 + 198 0.5743803978E-03 -0.1031478823E+01 0.9285716104E+00 0.3259083640E+01 0.5624358558E+01 + 199 0.5770145140E-03 -0.1034154241E+01 0.9291334033E+00 0.3260389334E+01 0.5624358558E+01 + 200 0.5818618277E-03 -0.1036822579E+01 0.9296965353E+00 0.3261691172E+01 0.5624358558E+01 + 201 0.5864779780E-03 -0.1039483932E+01 0.9302609916E+00 0.3262989171E+01 0.5624358558E+01 + 202 0.5902832635E-03 -0.1042138350E+01 0.9308267576E+00 0.3264283350E+01 0.5624358558E+01 + 203 0.5949353714E-03 -0.1044785872E+01 0.9313938184E+00 0.3265573729E+01 0.5624358558E+01 + 204 0.5989998948E-03 -0.1047426562E+01 0.9319621590E+00 0.3266860329E+01 0.5624358558E+01 + 205 0.6035123337E-03 -0.1050060478E+01 0.9325317642E+00 0.3268143171E+01 0.5624358558E+01 + 206 0.6076848906E-03 -0.1052687665E+01 0.9331026190E+00 0.3269422275E+01 0.5624358558E+01 + 207 0.6110659536E-03 -0.1055308173E+01 0.9336747080E+00 0.3270697666E+01 0.5624358558E+01 + 208 0.6161950129E-03 -0.1057922061E+01 0.9342480158E+00 0.3271969366E+01 0.5624358558E+01 + 209 0.6206170105E-03 -0.1060529373E+01 0.9348225269E+00 0.3273237398E+01 0.5624358558E+01 + 210 0.6243581601E-03 -0.1063130159E+01 0.9353982257E+00 0.3274501789E+01 0.5624358558E+01 + 211 0.6301878435E-03 -0.1065724463E+01 0.9359750964E+00 0.3275762562E+01 0.5624358558E+01 + 212 0.6348564298E-03 -0.1068312325E+01 0.9365531232E+00 0.3277019745E+01 0.5624358558E+01 + 213 0.6409693709E-03 -0.1070894379E+01 0.9371322901E+00 0.3278273365E+01 0.5624358558E+01 + 214 0.6445945085E-03 -0.1073469843E+01 0.9377125811E+00 0.3279523448E+01 0.5624358558E+01 + 215 0.6487362879E-03 -0.1076038732E+01 0.9382939801E+00 0.3280770024E+01 0.5624358558E+01 + 216 0.6543845189E-03 -0.1078601336E+01 0.9388764707E+00 0.3282013122E+01 0.5624358558E+01 + 217 0.6594707420E-03 -0.1081157760E+01 0.9394600366E+00 0.3283252772E+01 0.5624358558E+01 + 218 0.6637734702E-03 -0.1083708047E+01 0.9400446612E+00 0.3284489004E+01 0.5624358558E+01 + 219 0.6688813885E-03 -0.1086252186E+01 0.9406303280E+00 0.3285721850E+01 0.5624358558E+01 + 220 0.6742251757E-03 -0.1088790263E+01 0.9412170203E+00 0.3286951341E+01 0.5624358558E+01 + 221 0.6783745738E-03 -0.1091322318E+01 0.9418047212E+00 0.3288177509E+01 0.5624358558E+01 + 222 0.6860294147E-03 -0.1093848433E+01 0.9423934139E+00 0.3289400388E+01 0.5624358558E+01 + 223 0.6917248485E-03 -0.1096368613E+01 0.9429830813E+00 0.3290620011E+01 0.5624358558E+01 + 224 0.6952970086E-03 -0.1098882575E+01 0.9435737063E+00 0.3291836412E+01 0.5624358558E+01 + 225 0.7009861191E-03 -0.1101390977E+01 0.9441652716E+00 0.3293049625E+01 0.5624358558E+01 + 226 0.7079907726E-03 -0.1103893599E+01 0.9447577598E+00 0.3294259685E+01 0.5624358558E+01 + 227 0.7137939815E-03 -0.1106390467E+01 0.9453511535E+00 0.3295466627E+01 0.5624358558E+01 + 228 0.7204173126E-03 -0.1108881604E+01 0.9459454352E+00 0.3296670485E+01 0.5624358558E+01 + 229 0.7259309361E-03 -0.1111367034E+01 0.9465405871E+00 0.3297871296E+01 0.5624358558E+01 + 230 0.7333074199E-03 -0.1113846912E+01 0.9471365915E+00 0.3299069095E+01 0.5624358558E+01 + 231 0.7389183411E-03 -0.1116321048E+01 0.9477334305E+00 0.3300263918E+01 0.5624358558E+01 + 232 0.7457036761E-03 -0.1118789522E+01 0.9483310860E+00 0.3301455802E+01 0.5624358558E+01 + 233 0.7521987445E-03 -0.1121252548E+01 0.9489295399E+00 0.3302644782E+01 0.5624358558E+01 + 234 0.7591639453E-03 -0.1123710089E+01 0.9495287741E+00 0.3303830896E+01 0.5624358558E+01 + 235 0.7651524095E-03 -0.1126162174E+01 0.9501287702E+00 0.3305014180E+01 0.5624358558E+01 + 236 0.7726802801E-03 -0.1128608821E+01 0.9507295098E+00 0.3306194671E+01 0.5624358558E+01 + 237 0.7812531282E-03 -0.1131050221E+01 0.9513309742E+00 0.3307372406E+01 0.5624358558E+01 + 238 0.7900038521E-03 -0.1133487370E+01 0.9519331450E+00 0.3308547424E+01 0.5624358558E+01 + 239 0.7960242066E-03 -0.1135918182E+01 0.9525360032E+00 0.3309719761E+01 0.5624358558E+01 + 240 0.8027242871E-03 -0.1138343528E+01 0.9531395301E+00 0.3310889458E+01 0.5624358558E+01 + 241 0.8126073430E-03 -0.1140763674E+01 0.9537437066E+00 0.3312056552E+01 0.5624358558E+01 + 242 0.8207166622E-03 -0.1143178616E+01 0.9543485137E+00 0.3313221083E+01 0.5624358558E+01 + 243 0.8287285425E-03 -0.1145588345E+01 0.9549539322E+00 0.3314383092E+01 0.5624358558E+01 + 244 0.8387832768E-03 -0.1147992966E+01 0.9555599428E+00 0.3315542620E+01 0.5624358558E+01 + 245 0.8483394325E-03 -0.1150392511E+01 0.9561665260E+00 0.3316699711E+01 0.5624358558E+01 + 246 0.8572967496E-03 -0.1152787016E+01 0.9567736624E+00 0.3317854407E+01 0.5624358558E+01 + 247 0.8676906403E-03 -0.1155176647E+01 0.9573813323E+00 0.3319006755E+01 0.5624358558E+01 + 248 0.8772297135E-03 -0.1157561125E+01 0.9579895160E+00 0.3320156803E+01 0.5624358558E+01 + 249 0.8883086668E-03 -0.1159940628E+01 0.9585981936E+00 0.3321304601E+01 0.5624358558E+01 + 250 0.8978604521E-03 -0.1162315133E+01 0.9592073452E+00 0.3322450201E+01 0.5624358558E+01 + 251 0.9109672393E-03 -0.1164684813E+01 0.9598169507E+00 0.3323593661E+01 0.5624358558E+01 + 252 0.9213844125E-03 -0.1167049689E+01 0.9604269900E+00 0.3324735040E+01 0.5624358558E+01 + 253 0.9344291096E-03 -0.1169409609E+01 0.9610374428E+00 0.3325874404E+01 0.5624358558E+01 + 254 0.9461184145E-03 -0.1171764863E+01 0.9616482887E+00 0.3327011824E+01 0.5624358558E+01 + 255 0.9598834330E-03 -0.1174115311E+01 0.9622595073E+00 0.3328147378E+01 0.5624358558E+01 + 256 0.9724107611E-03 -0.1176461036E+01 0.9628710778E+00 0.3329281150E+01 0.5624358558E+01 diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py index 9619166042..e745188be0 100644 --- a/tests/test_compute_everything.py +++ b/tests/test_compute_everything.py @@ -8,6 +8,7 @@ from desc.coils import FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, SplineXYZCoil from desc.compute import data_index, xyz2rpz, xyz2rpz_vec +from desc.compute.utils import _grow_seeds from desc.examples import get from desc.geometry import ( FourierPlanarCurve, @@ -217,12 +218,14 @@ def test_compute_everything(): # size cap at 100 mb, so can't hit suggested resolution for some things. warnings.filterwarnings("ignore", category=ResolutionWarning) for p in things: - names = { - name - for name in data_index[p] - # Skip these quantities as they should be covered in other tests. - if not data_index[p][name]["source_grid_requirement"] - } + + names = set(data_index[p].keys()) + + def need_src(name): + return bool(data_index[p][name]["source_grid_requirement"]) + + names -= _grow_seeds(p, set(filter(need_src, names)), names) + this_branch_data_rpz[p] = things[p].compute( list(names), **grid.get(p, {}), basis="rpz" ) diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 752cfdcc8a..52bd5b13e4 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -6,6 +6,7 @@ from desc.compute import rpz2xyz_vec from desc.equilibrium import Equilibrium +from desc.equilibrium.coords import get_rtz_grid from desc.examples import get from desc.geometry import FourierRZToroidalSurface from desc.grid import LinearGrid @@ -1575,3 +1576,50 @@ def test_parallel_grad(): np.testing.assert_allclose( data["gbdrift (secular)"], data["gbdrift (secular)/phi"] * data["phi"] ) + + +@pytest.mark.unit +def test_parallel_grad_fd(DummyStellarator): + """Test that the parallel gradients match with numerical gradients.""" + eq = load(load_from=str(DummyStellarator["output_path"]), file_format="hdf5") + grid = get_rtz_grid( + eq, + 0.5, + 0, + np.linspace(0, 2 * np.pi, 50), + coordinates="raz", + period=(np.inf, 2 * np.pi, np.inf), + ) + data = eq.compute( + [ + "|B|", + "|B|_z|r,a", + "|e_zeta|r,a|", + "|e_zeta|r,a|_z|r,a", + "B^zeta", + "B^zeta_z|r,a", + ], + grid=grid, + ) + dz = grid.source_grid.spacing[:, 2] + fd = np.convolve(data["|B|"], FD_COEF_1_4, "same") / dz + np.testing.assert_allclose( + data["|B|_z|r,a"][2:-2], + fd[2:-2], + rtol=1e-2, + atol=1e-2 * np.mean(np.abs(data["|B|_z|r,a"])), + ) + fd = np.convolve(data["|e_zeta|r,a|"], FD_COEF_1_4, "same") / dz + np.testing.assert_allclose( + data["|e_zeta|r,a|_z|r,a"][2:-2], + fd[2:-2], + rtol=1e-2, + atol=1e-2 * np.mean(np.abs(data["|e_zeta|r,a|_z|r,a"])), + ) + fd = np.convolve(data["B^zeta"], FD_COEF_1_4, "same") / dz + np.testing.assert_allclose( + data["B^zeta_z|r,a"][2:-2], + fd[2:-2], + rtol=1e-2, + atol=1e-2 * np.mean(np.abs(data["B^zeta_z|r,a"])), + ) diff --git a/tests/test_examples.py b/tests/test_examples.py index c4d423c664..b43245e28b 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1965,7 +1965,7 @@ def test_ballooning_stability_opt(): for i in range(len(surfaces)): rho = surfaces[i] - grid = eq.get_rtz_grid( + grid = eq._get_rtz_grid( rho, alpha, zeta, @@ -2045,7 +2045,7 @@ def test_ballooning_stability_opt(): for i in range(len(surfaces)): rho = surfaces[i] - grid = eq.get_rtz_grid( + grid = eq._get_rtz_grid( rho, alpha, zeta, diff --git a/tests/test_integrals.py b/tests/test_integrals.py index aaa557a0f9..62e364946b 100644 --- a/tests/test_integrals.py +++ b/tests/test_integrals.py @@ -42,7 +42,8 @@ interp_to_argmin_hard, ) from desc.integrals._interp_utils import fourier_pts -from desc.integrals._quad_utils import ( +from desc.integrals.basis import FourierChebyshevSeries +from desc.integrals.quad_utils import ( automorphism_sin, bijection_from_disc, chebgauss1, @@ -53,7 +54,6 @@ leggauss_lob, tanh_sinh, ) -from desc.integrals.basis import FourierChebyshevSeries from desc.integrals.singularities import _get_quadrature_nodes from desc.integrals.surface_integral import _get_grid_surface from desc.transform import Transform diff --git a/tests/test_interp_utils.py b/tests/test_interp_utils.py index 4c7df8d005..265de4b821 100644 --- a/tests/test_interp_utils.py +++ b/tests/test_interp_utils.py @@ -27,8 +27,8 @@ polyroot_vec, polyval_vec, ) -from desc.integrals._quad_utils import bijection_to_disc from desc.integrals.basis import FourierChebyshevSeries +from desc.integrals.quad_utils import bijection_to_disc class TestPolyUtils: diff --git a/tests/test_neoclassical.py b/tests/test_neoclassical.py new file mode 100644 index 0000000000..d517678cb9 --- /dev/null +++ b/tests/test_neoclassical.py @@ -0,0 +1,182 @@ +"""Test for neoclassical transport compute functions.""" + +from datetime import datetime + +import matplotlib.pyplot as plt +import numpy as np +import pytest +from tests.test_plotting import tol_1d + +from desc.equilibrium.coords import get_rtz_grid +from desc.examples import get +from desc.grid import LinearGrid +from desc.utils import setdefault +from desc.vmec import VMECIO + + +@pytest.mark.unit +def test_fieldline_average(): + """Test that fieldline average converges to surface average.""" + rho = np.array([1]) + alpha = np.array([0]) + eq = get("DSHAPE") + iota_grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym) + iota = iota_grid.compress(eq.compute("iota", grid=iota_grid)["iota"]).item() + # For axisymmetric devices, one poloidal transit must be exact. + zeta = np.linspace(0, 2 * np.pi / iota, 25) + grid = get_rtz_grid(eq, rho, alpha, zeta, coordinates="raz") + data = eq.compute( + ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid + ) + np.testing.assert_allclose( + data["fieldline length"] / data["fieldline length/volume"], + data["V_r(r)"] / (4 * np.pi**2), + rtol=1e-3, + ) + assert np.all(data["fieldline length"] > 0) + assert np.all(data["fieldline length/volume"] > 0) + + # Otherwise, many toroidal transits are necessary to sample surface. + eq = get("W7-X") + zeta = np.linspace(0, 40 * np.pi, 300) + grid = get_rtz_grid(eq, rho, alpha, zeta, coordinates="raz") + data = eq.compute( + ["fieldline length", "fieldline length/volume", "V_r(r)"], grid=grid + ) + np.testing.assert_allclose( + data["fieldline length"] / data["fieldline length/volume"], + data["V_r(r)"] / (4 * np.pi**2), + rtol=1e-3, + ) + assert np.all(data["fieldline length"] > 0) + assert np.all(data["fieldline length/volume"] > 0) + + +@pytest.mark.unit +@pytest.mark.mpl_image_compare(remove_text=True, tolerance=tol_1d) +def test_effective_ripple(): + """Test effective ripple with W7-X.""" + eq = get("W7-X") + rho = np.linspace(0, 1, 10) + alpha = np.array([0]) + zeta = np.linspace(0, 20 * np.pi, 1000) + grid = get_rtz_grid(eq, rho, alpha, zeta, coordinates="raz") + data = eq.compute("effective ripple", grid=grid) + assert np.isfinite(data["effective ripple"]).all() + np.testing.assert_allclose( + data["effective ripple 3/2"] ** (2 / 3), + data["effective ripple"], + err_msg="Bug in source grid logic in eq.compute.", + ) + eps_32 = grid.compress(data["effective ripple 3/2"]) + fig, ax = plt.subplots() + ax.plot(rho, eps_32, marker="o") + + neo_rho, neo_eps_32 = NeoIO.read("tests/inputs/neo_out.w7x") + np.testing.assert_allclose(eps_32, np.interp(rho, neo_rho, neo_eps_32), rtol=0.16) + return fig + + +class NeoIO: + """Class to interface with NEO.""" + + def __init__(self, name, eq, ns=256, M_booz=None, N_booz=None): + self.name = name + self.vmec_file = f"wout_{name}.nc" + self.booz_file = f"boozmn.{name}" + self.neo_in_file = f"neo_in.{name}" + self.neo_out_file = f"neo_out.{name}" + + self.eq = eq + self.ns = ns # number of surfaces + self.M_booz = setdefault(M_booz, 3 * eq.M + 1) + self.N_booz = setdefault(N_booz, 3 * eq.N) + + @staticmethod + def read(name): + """Return ρ and ε¹ᐧ⁵ from NEO output with given name.""" + neo_eps = np.loadtxt(name)[:, 1] + neo_rho = np.sqrt(np.linspace(1 / (neo_eps.size + 1), 1, neo_eps.size)) + # replace bad values with linear interpolation + good = np.isfinite(neo_eps) + neo_eps[~good] = np.interp(neo_rho[~good], neo_rho[good], neo_eps[good]) + return neo_rho, neo_eps + + def write(self): + """Write neo input file.""" + print(f"Writing VMEC wout to {self.vmec_file}") + VMECIO.save(self.eq, self.vmec_file, surfs=self.ns, verbose=0) + self._write_booz() + self._write_neo() + + def _write_booz(self): + print(f"Writing boozer output file to {self.booz_file}") + import booz_xform as bx + + b = bx.Booz_xform() + b.read_wout(self.vmec_file) + b.mboz = self.M_booz + b.nboz = self.N_booz + b.run() + b.write_boozmn(self.booz_file) + + def _write_neo( + self, + theta_n=200, + phi_n=200, + num_pitch=50, + multra=1, + acc_req=0.01, + nbins=100, + nstep_per=75, + nstep_min=500, + nstep_max=2000, + verbose=2, + ): + print(f"Writing NEO input file to {self.neo_in_file}") + f = open(self.neo_in_file, "w") + + def writeln(s): + f.write(str(s)) + f.write("\n") + + # https://princetonuniversity.github.io/STELLOPT/NEO + writeln(f"'#' {datetime.now()}") + writeln(f"'#' {self.vmec_file}") + writeln(f"'#' M_booz={self.M_booz}. N_booz={self.N_booz}.") + writeln(self.booz_file) + writeln(self.neo_out_file) + # Neo computes things on the so-called "half grid" between the full grid. + # There are only ns - 1 surfaces there. + writeln(self.ns - 1) + # NEO starts indexing at 1 and does not compute on axis (index 1). + surface_indices = " ".join(str(i) for i in range(2, self.ns + 1)) + writeln(surface_indices) + writeln(theta_n) + writeln(phi_n) + writeln(0) + writeln(0) + writeln(num_pitch) + writeln(multra) + writeln(acc_req) + writeln(nbins) + writeln(nstep_per) + writeln(nstep_min) + writeln(nstep_max) + writeln(0) + writeln(verbose) + writeln(0) + writeln(0) + writeln(2) + writeln(0) + writeln(0) + writeln(0) + writeln(0) + writeln(0) + writeln("'#'\n'#'\n'#'") + writeln(0) + writeln(f"neo_cur.{self.name}") + writeln(200) + writeln(2) + writeln(0) + f.close() diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 0476e9def1..4ebcece58c 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -49,6 +49,7 @@ CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, + EffectiveRipple, Elongation, Energy, ForceBalance, @@ -2422,6 +2423,16 @@ def test_loss_function_asserts(): RotationalTransform(eq=eq, loss_function=fun) +def _reduced_resolution_objective(eq, objective): + """Speed up testing suite by defining rules to reduce objective resolution.""" + kwargs = {} + if objective in {EffectiveRipple}: + kwargs["Y_B"] = 50 + kwargs["num_transit"] = 2 + kwargs["num_pitch"] = 25 + return objective(eq=eq, **kwargs) + + class TestComputeScalarResolution: """Test that compute_scalar values are roughly independent of grid resolution.""" @@ -2846,8 +2857,9 @@ def test_compute_scalar_resolution_others(self, objective): M_grid=int(self.eq.M * res), N_grid=int(self.eq.N * res), ) - - obj = ObjectiveFunction(objective(eq=self.eq), use_jit=False) + obj = ObjectiveFunction( + _reduced_resolution_objective(self.eq, objective), use_jit=False + ) obj.build(verbose=0) f[i] = obj.compute_scalar(obj.x()) np.testing.assert_allclose(f, f[-1], rtol=5e-2) @@ -2903,6 +2915,7 @@ class TestObjectiveNaNGrad: CoilSetLinkingNumber, CoilSetMinDistance, CoilTorsion, + EffectiveRipple, ForceBalanceAnisotropic, FusionPower, HeatingPowerISS04, @@ -3170,6 +3183,17 @@ def test_objective_no_nangrad_omnigenity(self, helicity): g = obj.grad(obj.x()) assert not np.any(np.isnan(g)), str(helicity) + @pytest.mark.unit + def test_objective_no_nangrad_effective_ripple(self): + """Effective ripple.""" + eq = get("ESTELL") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(2, 2, 2, 4, 4, 4) + obj = ObjectiveFunction(_reduced_resolution_objective(eq, EffectiveRipple)) + obj.build(verbose=0) + g = obj.grad(obj.x()) + assert not np.any(np.isnan(g)) + @pytest.mark.unit def test_objective_no_nangrad_ballooning(self): """BallooningStability.""" diff --git a/tests/test_quad_utils.py b/tests/test_quad_utils.py index 81a2d41197..9ee62b1a1d 100644 --- a/tests/test_quad_utils.py +++ b/tests/test_quad_utils.py @@ -8,7 +8,7 @@ from scipy.special import roots_chebyu from desc.backend import jnp -from desc.integrals._quad_utils import ( +from desc.integrals.quad_utils import ( automorphism_arcsin, automorphism_sin, bijection_from_disc,