diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 662413501..cacb0095f 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -9,8 +9,9 @@ expensive computations. """ -from desc.backend import jnp +from desc.backend import jnp, trapezoid +from ..grid import QuadratureGrid from ..integrals.surface_integral import line_integrals, surface_integrals from ..utils import cross, dot, safenorm from .data_index import register_compute_fun @@ -27,11 +28,16 @@ transforms={"grid": []}, profiles=[], coordinates="", - data=["sqrt(g)"], + data=["sqrt(g)", "V(r)", "rho"], resolution_requirement="rtz", ) def _V(params, transforms, profiles, data, **kwargs): - data["V"] = jnp.sum(data["sqrt(g)"] * transforms["grid"].weights) + if isinstance(transforms["grid"], QuadratureGrid): + data["V"] = jnp.sum(data["sqrt(g)"] * transforms["grid"].weights) + else: + # To approximate volume at ρ ~ 1, we scale by ρ⁻², assuming the integrand + # varies little from max ρ to ρ = 1 and a roughly circular cross-section. + data["V"] = jnp.max(data["V(r)"]) / jnp.max(data["rho"]) ** 2 return data @@ -40,27 +46,19 @@ def _V(params, transforms, profiles, data, **kwargs): label="V", units="m^{3}", units_long="cubic meters", - description="Volume", + description="Volume enclosed by surface, scaled by max(ρ)⁻²", dim=1, params=[], transforms={"grid": []}, profiles=[], coordinates="", - data=["e_theta", "e_zeta", "x"], + data=["V(r)", "rho"], parameterization="desc.geometry.surface.FourierRZToroidalSurface", - resolution_requirement="tz", ) def _V_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - # divergence theorem: integral(dV div [0, 0, Z]) = integral(dS dot [0, 0, Z]) - data["V"] = jnp.max( # take max in case there are multiple surfaces for some reason - jnp.abs( - surface_integrals( - transforms["grid"], - cross(data["e_theta"], data["e_zeta"])[:, 2] * data["x"][:, 2], - expand_out=False, - ) - ) - ) + # To approximate volume at ρ ~ 1, we scale by ρ⁻², assuming the integrand + # varies little from max ρ to ρ = 1 and a roughly circular cross-section. + data["V"] = jnp.max(data["V(r)"]) / jnp.max(data["rho"]) ** 2 return data @@ -76,6 +74,10 @@ def _V_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="r", data=["e_theta", "e_zeta", "Z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], resolution_requirement="tz", ) def _V_of_r(params, transforms, profiles, data, **kwargs): @@ -156,50 +158,31 @@ def _V_rrr_of_r(params, transforms, profiles, data, **kwargs): label="A(\\zeta)", units="m^{2}", units_long="square meters", - description="Cross-sectional area as function of zeta", + description="Area of enclosed cross-section (enclosed constant phi surface), " + "scaled by max(ρ)⁻², as function of zeta", dim=1, params=[], transforms={"grid": []}, profiles=[], coordinates="z", - data=["|e_rho x e_theta|"], + data=["Z", "n_rho", "e_theta|r,p", "rho"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.ZernikeRZToroidalSection", + "desc.geometry.surface.FourierRZToroidalSurface", ], - resolution_requirement="rt", + resolution_requirement="t", + grid_requirement={"sym": False}, + # FIXME: For nonzero omega we need to integrate over theta at constant phi. + # Add source_grid_requirement={"coordinates": "rtp", "is_meshgrid": True} + # TODO: Recognize when omega = 0 and ignore all source grid requirements + # if the given grid satisfies them with phi replaced by zeta. ) def _A_of_z(params, transforms, profiles, data, **kwargs): - data["A(z)"] = surface_integrals( - transforms["grid"], - data["|e_rho x e_theta|"], - surface_label="zeta", - expand_out=True, - ) - return data - - -@register_compute_fun( - name="A(z)", - label="A(\\zeta)", - units="m^{2}", - units_long="square meters", - description="Area of enclosed cross-section (enclosed constant phi surface), " - "scaled by max(ρ)⁻², as function of zeta", - dim=1, - params=[], - transforms={"grid": []}, - profiles=[], - coordinates="z", - data=["Z", "n_rho", "e_theta|r,p", "rho"], - parameterization=["desc.geometry.surface.FourierRZToroidalSurface"], - resolution_requirement="rt", # just need max(rho) near 1 - # FIXME: Add source grid requirement once omega is nonzero. -) -def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - # Denote any vector v = [vᴿ, v^ϕ, vᶻ] with a tuple of its contravariant components. + # Denote any vector v = v¹ R̂ + v² ϕ̂ + v³ Ẑ by v = [v¹, v², v³] where R̂, ϕ̂, Ẑ + # are the normalized basis vectors of the cylindrical coordinates R, ϕ, Z. # We use a 2D divergence theorem over constant ϕ toroidal surface (i.e. R, Z plane). - # In this geometry, the divergence operator on a polar basis vector is + # In this geometry, the divergence operator in this coordinate system is # div = ([∂_R, ∂_ϕ, ∂_Z] ⊗ [1, 0, 1]) dot . # ∫ dA div v = ∫ dℓ n dot v # where n is the unit normal such that n dot e_θ|ρ,ϕ = 0 and n dot e_ϕ|R,Z = 0, @@ -213,16 +196,12 @@ def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg line_integrals( transforms["grid"], data["Z"] * n[:, 2] * safenorm(data["e_theta|r,p"], axis=-1), - # FIXME: Works currently for omega = zero, but for nonzero omega - # we need to integrate over theta at constant phi. - # Should be simple once we have coordinate mapping and source grid - # logic from GitHub pull request #1024. line_label="theta", fix_surface=("rho", max_rho), expand_out=True, ) # To approximate area at ρ ~ 1, we scale by ρ⁻², assuming the integrand - # varies little from ρ = max_rho to ρ = 1 and a roughly circular cross-section. + # varies little from max ρ to ρ = 1 and a roughly circular cross-section. / max_rho**2 ) return data @@ -233,23 +212,53 @@ def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg label="A", units="m^{2}", units_long="square meters", - description="Average cross-sectional area", + description="Average enclosed cross-sectional area, scaled by max(ρ)⁻²", + # Simple toroidal average A₀ = ∫ A(ζ) dζ / (2π) matches the convention for the + # average major radius R₀ = ∫ R(ρ=0) dζ / (2π). dim=0, params=[], transforms={"grid": []}, profiles=[], coordinates="", - data=["A(z)"], + data=["Z", "n_rho", "e_theta|r,p", "rho", "phi"], parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.core.Surface", + "desc.equilibrium.equilibrium.Equilibrium", ], - resolution_requirement="z", + resolution_requirement="tz", + # FIXME: For nonzero omega we need to integrate over theta at constant phi. + # Add source_grid_requirement={"coordinates": "rtp", "is_meshgrid": True} + # TODO: Recognize when omega = 0 and ignore all source grid requirements + # if the given grid satisfies them with phi replaced by zeta. ) def _A(params, transforms, profiles, data, **kwargs): - data["A"] = jnp.mean( - transforms["grid"].compress(data["A(z)"], surface_label="zeta") + # Denote any vector v = v¹ R̂ + v² ϕ̂ + v³ Ẑ by v = [v¹, v², v³] where R̂, ϕ̂, Ẑ + # are the normalized basis vectors of the cylindrical coordinates R, ϕ, Z. + # We use a 2D divergence theorem over constant ϕ toroidal surface (i.e. R, Z plane). + # In this geometry, the divergence operator in this coordinate system is + # div = ([∂_R, ∂_ϕ, ∂_Z] ⊗ [1, 0, 1]) dot . + # ∫ dA div v = ∫ dℓ n dot v + # where n is the unit normal such that n dot e_θ|ρ,ϕ = 0 and n dot e_ϕ|R,Z = 0, + # and the labels following | denote those coordinates are fixed. + # Now choose v = [0, 0, Z], and n in the direction (e_θ|ρ,ζ × e_ζ|ρ,θ) ⊗ [1, 0, 1]. + n = data["n_rho"] + n = n.at[:, 1].set(0) + n = n / jnp.linalg.norm(n, axis=-1)[:, jnp.newaxis] + max_rho = jnp.max(data["rho"]) + A = jnp.abs( + line_integrals( + transforms["grid"], + data["Z"] * n[:, 2] * safenorm(data["e_theta|r,p"], axis=-1), + line_label="theta", + fix_surface=("rho", max_rho), + expand_out=False, + ) + # To approximate area at ρ ~ 1, we scale by ρ⁻², assuming the integrand + # varies little from max ρ to ρ = 1 and a roughly circular cross-section. + / max_rho**2 ) + phi = transforms["grid"].compress(data["phi"], "zeta") + data["A"] = jnp.squeeze(trapezoid(A, phi) / jnp.ptp(phi) if A.size > 1 else A) return data @@ -276,25 +285,22 @@ def _A_of_r(params, transforms, profiles, data, **kwargs): label="S", units="m^{2}", units_long="square meters", - description="Surface area of outermost flux surface", + description="Surface area of outermost flux surface, scaled by max(ρ)⁻¹", dim=0, params=[], transforms={"grid": []}, profiles=[], coordinates="", - data=["|e_theta x e_zeta|"], + data=["S(r)", "rho"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.FourierRZToroidalSurface", ], - resolution_requirement="tz", ) def _S(params, transforms, profiles, data, **kwargs): - data["S"] = jnp.max( - surface_integrals( - transforms["grid"], data["|e_theta x e_zeta|"], expand_out=False - ) - ) + # To approximate surface are at ρ ~ 1, we scale by ρ⁻¹, assuming the integrand + # varies little from max ρ to ρ = 1. + data["S"] = jnp.max(data["S(r)"]) / jnp.max(data["rho"]) return data @@ -310,6 +316,10 @@ def _S(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="r", data=["|e_theta x e_zeta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], resolution_requirement="tz", ) def _S_of_r(params, transforms, profiles, data, **kwargs): @@ -360,10 +370,12 @@ def _S_rr_of_r(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="R0", - label="R_{0}", + label="R_{0} = V / (2\\pi A) = \\int R(\\rho=0) d\\zeta / (2\\pi)", units="m", units_long="meters", description="Average major radius", + # This differs from the average value of R on the magnetic axis. + # R₀ ≠ 〈 R(ρ=0) 〉 = ∫ (R ‖e_ζ‖)(ρ=0) dζ / ∫ ‖e_ζ‖(ρ=0) dζ. dim=0, params=[], transforms={}, @@ -441,7 +453,12 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs): "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.core.Surface", ], - resolution_requirement="rt", # just need max(rho) near 1 + resolution_requirement="t", + grid_requirement={"sym": False}, + # FIXME: For nonzero omega we need to integrate over theta at constant phi. + # Add source_grid_requirement={"coordinates": "rtp", "is_meshgrid": True} + # TODO: Recognize when omega = 0 and ignore all source grid requirements + # if the given grid satisfies them with phi replaced by zeta. ) def _perimeter_of_z(params, transforms, profiles, data, **kwargs): max_rho = jnp.max(data["rho"]) @@ -449,16 +466,12 @@ def _perimeter_of_z(params, transforms, profiles, data, **kwargs): line_integrals( transforms["grid"], safenorm(data["e_theta|r,p"], axis=-1), - # FIXME: Works currently for omega = zero, but for nonzero omega - # we need to integrate over theta at constant phi. - # Should be simple once we have coordinate mapping and source grid - # logic from GitHub pull request #1024. line_label="theta", fix_surface=("rho", max_rho), expand_out=True, ) # To approximate perimeter at ρ ~ 1, we scale by ρ⁻¹, assuming the integrand - # varies little from ρ = max_rho to ρ = 1. + # varies little from max ρ to ρ = 1. / max_rho ) return data diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index f5ce69b3b..26168bc19 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -979,7 +979,6 @@ def need_src(name): # Warn if best way to compute accurately is increasing resolution. for dep in deps: req = data_index[p][dep]["resolution_requirement"] - coords = data_index[p][dep]["coordinates"] msg = lambda direction: ( f"Dependency {dep} may require more {direction}" " resolution to compute accurately." @@ -988,7 +987,9 @@ def need_src(name): # if need more radial resolution "r" in req and grid.L < self.L_grid # and won't override grid to one with more radial resolution - and not (override_grid and coords in {"z", ""}), + and not ( + override_grid and (is_1dz_tor_grid(dep) or is_0d_vol_grid(dep)) + ), ResolutionWarning, msg("radial"), ) @@ -996,7 +997,14 @@ def need_src(name): # if need more poloidal resolution "t" in req and grid.M < self.M_grid # and won't override grid to one with more poloidal resolution - and not (override_grid and coords in {"r", "z", ""}), + and not ( + override_grid + and ( + is_1dr_rad_grid(dep) + or is_1dz_tor_grid(dep) + or is_0d_vol_grid(dep) + ) + ), ResolutionWarning, msg("poloidal"), ) @@ -1004,7 +1012,9 @@ def need_src(name): # if need more toroidal resolution "z" in req and grid.N < self.N_grid # and won't override grid to one with more toroidal resolution - and not (override_grid and coords in {"r", ""}), + and not ( + override_grid and (is_1dr_rad_grid(dep) or is_0d_vol_grid(dep)) + ), ResolutionWarning, msg("toroidal"), ) diff --git a/desc/grid.py b/desc/grid.py index c4b8a46e7..357e26f63 100644 --- a/desc/grid.py +++ b/desc/grid.py @@ -66,19 +66,19 @@ def _set_up(self): del self._unique_theta_idx def _enforce_symmetry(self): - """Enforce stellarator symmetry. + """Remove unnecessary nodes assuming poloidal symmetry. - 1. Remove nodes with theta > pi. - 2. Rescale theta spacing to preserve dtheta weight. - Need to rescale on each theta coordinate curve by a different factor. - dtheta should = 2π / number of nodes remaining on that theta curve. - Nodes on the symmetry line should not be rescaled. + 1. Remove nodes with θ > π. + 2. Rescale θ spacing to preserve dθ weight. + Need to rescale on each θ coordinate curve by a different factor. + dθ = 2π / number of nodes remaining on that θ curve. + Nodes on the symmetry line should not be rescaled. """ if not self.sym: return # indices where poloidal coordinate is off the symmetry line of - # poloidal coord=0 or pi + # poloidal coord=0 or π off_sym_line_idx = self.nodes[:, 1] % np.pi != 0 __, inverse, off_sym_line_per_rho_surf_count = np.unique( self.nodes[off_sym_line_idx, 0], return_inverse=True, return_counts=True @@ -109,7 +109,7 @@ def _enforce_symmetry(self): # The first two assumptions let _per_poloidal_curve = _per_rho_surf. # The third assumption lets the scale factor be constant over a # particular theta curve, so that each node in the open interval - # (0, pi) has its spacing scaled up by the same factor. + # (0, π) has its spacing scaled up by the same factor. # Nodes at endpoints 0, π should not be scaled. scale = off_sym_line_per_rho_surf_count / ( off_sym_line_per_rho_surf_count - to_delete_per_rho_surf_count @@ -206,7 +206,13 @@ def NFP(self): @property def sym(self): - """bool: True for stellarator symmetry, False otherwise.""" + """bool: True for poloidal symmetry, False otherwise. + + If true, the poloidal domain of this grid is [0, π] ⊂ [0, 2π). + Note that this is distinct from stellarator symmetry. + Still, when stellarator symmetry exists, flux surface integrals and + volume integrals are invariant to this truncation. + """ return self.__dict__.setdefault("_sym", False) @property @@ -917,7 +923,11 @@ class LinearGrid(_Grid): NFP : int Number of field periods (Default = 1). sym : bool - True for stellarator symmetry, False otherwise (Default = False). + Whether to truncate the poloidal domain to [0, π] ⊂ [0, 2π). + Note that this is distinct from stellarator symmetry. + Still, when stellarator symmetry exists, flux surface integrals and + volume integrals are invariant to this truncation, so setting this flag + to true will reduce memory consumption. (Default = False). axis : bool True to include a point at rho=0 (default), False for rho[0] = rho[1]/2. endpoint : bool @@ -1377,7 +1387,11 @@ class ConcentricGrid(_Grid): NFP : int number of field periods (Default = 1) sym : bool - True for stellarator symmetry, False otherwise (Default = False) + Whether to truncate the poloidal domain to [0, π] ⊂ [0, 2π). + Note that this is distinct from stellarator symmetry. + Still, when stellarator symmetry exists, flux surface integrals and + volume integrals are invariant to this truncation, so setting this flag + to true will reduce memory consumption. (Default = False). axis : bool True to include the magnetic axis, False otherwise (Default = False) node_pattern : {``'cheb1'``, ``'cheb2'``, ``'jacobi'``, ``linear``} diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 65fa1cf82..e21c5352f 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -99,6 +99,7 @@ def build(self, use_jit=True, verbose=1): M=eq.M * 2, N=eq.N * 2, NFP=eq.NFP, + sym=False, ) else: grid = self._grid @@ -242,6 +243,7 @@ def build(self, use_jit=True, verbose=1): M=eq.M * 2, N=eq.N * 2, NFP=eq.NFP, + sym=False, ) else: grid = self._grid diff --git a/desc/plotting.py b/desc/plotting.py index 12900ff41..59a22e591 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -186,15 +186,14 @@ def _format_ax(ax, is3d=False, rows=1, cols=1, figsize=None, equal=False): return plt.gcf(), ax else: ax = np.atleast_1d(ax) - if isinstance(ax.flatten()[0], matplotlib.axes.Axes): - return plt.gcf(), ax - else: - raise TypeError( - colored( - "ax argument must be None or an axis instance or array of axes", - "red", - ) - ) + errorif( + not isinstance(ax.flatten()[0], matplotlib.axes.Axes), + TypeError, + colored( + "ax argument must be None or an axis instance or array of axes", "red" + ), + ) + return plt.gcf(), ax def _get_grid(**kwargs): @@ -277,11 +276,10 @@ def _compute(eq, name, grid, component=None, reshape=True): """ parameterization = _parse_parameterization(eq) - if name not in data_index[parameterization]: - raise ValueError( - f"Unrecognized value '{name}' for " - + f"parameterization {parameterization}." - ) + errorif( + name not in data_index[parameterization], + msg=f"Unrecognized value '{name}' for parameterization {parameterization}.", + ) assert component in [ None, "R", @@ -293,9 +291,7 @@ def _compute(eq, name, grid, component=None, reshape=True): label = data_index[parameterization][name]["label"] - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - data = eq.compute(name, grid=grid)[name] + data = eq.compute(name, grid=grid)[name] if data_index[parameterization][name]["dim"] > 1: if component is None: @@ -511,47 +507,50 @@ def plot_1d(eq, name, grid=None, log=False, ax=None, return_data=False, **kwargs grid_kwargs = {"L": default_L, "N": default_N, "NFP": NFP} grid = _get_grid(**grid_kwargs) plot_axes = _get_plot_axes(grid) + + data, ylabel = _compute( + eq, name, grid, kwargs.pop("component", None), reshape=False + ) + + # reshape data to 1D if len(plot_axes) != 1: - return ValueError(colored("Grid must be 1D", "red")) + surface_label = {"r": "rho", "t": "theta", "z": "zeta"}.get( + data_index[parameterization][name]["coordinates"], None + ) + axis = {"r": 0, "t": 1, "z": 2}.get( + data_index[parameterization][name]["coordinates"], None + ) + errorif( + surface_label is None or axis is None, + NotImplementedError, + msg=colored("Grid must be 1D", "red"), + ) + data = grid.compress(data, surface_label=surface_label) + nodes = grid.compress(grid.nodes[:, axis], surface_label=surface_label) + else: + axis = plot_axes[0] + data = data.ravel() + nodes = grid.nodes[:, axis] - data, ylabel = _compute(eq, name, grid, kwargs.pop("component", None)) label = kwargs.pop("label", None) - fig, ax = _format_ax(ax, figsize=kwargs.pop("figsize", None)) - - # reshape data to 1D - data = data.flatten() linecolor = kwargs.pop("linecolor", colorblind_colors[0]) ls = kwargs.pop("ls", "-") lw = kwargs.pop("lw", 1) if log: data = np.abs(data) # ensure data is positive for log plot - ax.semilogy( - grid.nodes[:, plot_axes[0]], - data, - label=label, - color=linecolor, - ls=ls, - lw=lw, - ) + ax.semilogy(nodes, data, label=label, color=linecolor, ls=ls, lw=lw) else: - ax.plot( - grid.nodes[:, plot_axes[0]], - data, - label=label, - color=linecolor, - ls=ls, - lw=lw, - ) + ax.plot(nodes, data, label=label, color=linecolor, ls=ls, lw=lw) xlabel_fontsize = kwargs.pop("xlabel_fontsize", None) ylabel_fontsize = kwargs.pop("ylabel_fontsize", None) assert len(kwargs) == 0, f"plot_1d got unexpected keyword argument: {kwargs.keys()}" - xlabel = _AXIS_LABELS_RTZ[plot_axes[0]] + xlabel = _AXIS_LABELS_RTZ[axis] ax.set_xlabel(xlabel, fontsize=xlabel_fontsize) ax.set_ylabel(ylabel, fontsize=ylabel_fontsize) _set_tight_layout(fig) - plot_data = {xlabel.strip("$").strip("\\"): grid.nodes[:, plot_axes[0]], name: data} + plot_data = {xlabel.strip("$").strip("\\"): nodes, name: data} if label is not None: ax.legend() @@ -630,8 +629,7 @@ def plot_2d( grid_kwargs = {"M": 33, "N": 33, "NFP": eq.NFP, "axis": False} grid = _get_grid(**grid_kwargs) plot_axes = _get_plot_axes(grid) - if len(plot_axes) != 2: - return ValueError(colored("Grid must be 2D", "red")) + errorif(len(plot_axes) != 2, msg=colored("Grid must be 2D", "red")) component = kwargs.pop("component", None) if name != "B*n": data, label = _compute( @@ -644,14 +642,12 @@ def plot_2d( field = kwargs.pop("field", None) errorif( field is None, - ValueError, - "If B*n is entered as the variable to plot, a magnetic field" + msg="If B*n is entered as the variable to plot, a magnetic field" " must be provided.", ) errorif( not np.all(np.isclose(grid.nodes[:, 0], 1)), - ValueError, - "If B*n is entered as the variable to plot, " + msg="If B*n is entered as the variable to plot, " "the grid nodes must be at rho=1.", ) @@ -929,14 +925,12 @@ def plot_3d( field = kwargs.pop("field", None) errorif( field is None, - ValueError, - "If B*n is entered as the variable to plot, a magnetic field" + msg="If B*n is entered as the variable to plot, a magnetic field" " must be provided.", ) errorif( not np.all(np.isclose(grid.nodes[:, 0], 1)), - ValueError, - "If B*n is entered as the variable to plot, " + msg="If B*n is entered as the variable to plot, " "the grid nodes must be at rho=1.", ) @@ -970,8 +964,7 @@ def plot_3d( showaxislabels = kwargs.pop("showaxislabels", True) errorif( len(kwargs) != 0, - ValueError, - f"plot_3d got unexpected keyword argument: {kwargs.keys()}", + msg=f"plot_3d got unexpected keyword argument: {kwargs.keys()}", ) with warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -2472,8 +2465,7 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): showaxislabels = kwargs.pop("showaxislabels", True) errorif( len(kwargs) != 0, - ValueError, - f"plot_coils got unexpected keyword argument: {kwargs.keys()}", + msg=f"plot_coils got unexpected keyword argument: {kwargs.keys()}", ) errorif( not isinstance(coils, _Coil), @@ -2900,7 +2892,7 @@ def plot_boozer_surface( iota = grid_compute.compress(data["iota"]) else: # OmnigenousField iota = kwargs.pop("iota", None) - errorif(iota is None, ValueError, "iota must be supplied for OmnigenousField") + errorif(iota is None, msg="iota must be supplied for OmnigenousField") with warnings.catch_warnings(): warnings.simplefilter("ignore") data = thing.compute( diff --git a/tests/baseline/test_1d_elongation.png b/tests/baseline/test_1d_elongation.png index ccf04b533..65f760839 100644 Binary files a/tests/baseline/test_1d_elongation.png and b/tests/baseline/test_1d_elongation.png differ diff --git a/tests/baseline/test_plot_b_mag.png b/tests/baseline/test_plot_b_mag.png index c81f89f03..1a7db591f 100644 Binary files a/tests/baseline/test_plot_b_mag.png and b/tests/baseline/test_plot_b_mag.png differ diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 7a887be3d..490d89bc4 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/test_axis_limits.py b/tests/test_axis_limits.py index b8d173f7e..1c1a8e1de 100644 --- a/tests/test_axis_limits.py +++ b/tests/test_axis_limits.py @@ -189,6 +189,7 @@ def assert_is_continuous( } """ + p = "desc.equilibrium.equilibrium.Equilibrium" if kwargs is None: kwargs = {} # TODO: remove when boozer transform works with multiple surfaces @@ -200,6 +201,8 @@ def assert_is_continuous( or "_mn" in name or name == "B modes" or _skip_this(eq, name) + # skip if require full grid (sym false) + or not data_index[p][name]["grid_requirement"].get("sym", True) ) ] @@ -211,7 +214,6 @@ def assert_is_continuous( integrate = surface_integrals_map(grid, expand_out=False) data = eq.compute(names=names, grid=grid) - p = "desc.equilibrium.equilibrium.Equilibrium" for name in names: if name in not_continuous_limits: continue diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 9a9216cc8..21082b3a3 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -161,6 +161,12 @@ def test_surface_areas_2(): @pytest.mark.unit def test_elongation(): """Test that elongation approximation is correct.""" + surf1 = FourierRZToroidalSurface( + R_lmn=[10, 1, 0.2], + Z_lmn=[-1, -0.2], + modes_R=[[0, 0], [1, 0], [0, 1]], + modes_Z=[[-1, 0], [0, -1]], + ) surf2 = FourierRZToroidalSurface( R_lmn=[10, 1, 0.2], Z_lmn=[-2, -0.2], @@ -173,17 +179,15 @@ def test_elongation(): modes_R=[[0, 0], [1, 0], [0, 1]], modes_Z=[[-1, 0], [0, -1]], ) - eq1 = Equilibrium() # elongation = 1 - eq2 = Equilibrium(surface=surf2) # elongation = 2 - eq3 = Equilibrium(surface=surf3) # elongation = 3 - grid = LinearGrid(L=5, M=2 * eq3.M_grid, N=eq3.N_grid, NFP=eq3.NFP, sym=eq3.sym) - data1 = eq1.compute(["a_major/a_minor"], grid=grid) - data2 = eq2.compute(["a_major/a_minor"], grid=grid) - data3 = eq3.compute(["a_major/a_minor"], grid=grid) + assert surf3.sym + grid = LinearGrid(rho=1, M=3 * surf3.M, N=surf3.N, NFP=surf3.NFP, sym=False) + data1 = surf1.compute(["a_major/a_minor"], grid=grid) + data2 = surf2.compute(["a_major/a_minor"], grid=grid) + data3 = surf3.compute(["a_major/a_minor"], grid=grid) # elongation approximation is less accurate as elongation increases np.testing.assert_allclose(1.0, data1["a_major/a_minor"]) - np.testing.assert_allclose(2.0, data2["a_major/a_minor"], rtol=1e-3) - np.testing.assert_allclose(3.0, data3["a_major/a_minor"], rtol=1e-2) + np.testing.assert_allclose(2.0, data2["a_major/a_minor"], rtol=1e-4) + np.testing.assert_allclose(3.0, data3["a_major/a_minor"], rtol=1e-3) @pytest.mark.slow @@ -1502,11 +1506,13 @@ def test_surface_equilibrium_geometry(): for key in data: x = eq.compute(key)[key].max() # max needed for elongation broadcasting y = eq.surface.compute(key)[key].max() - if key == "a_major/a_minor": - rtol, atol = 1e-2, 0 # need looser tol here bc of different grids + if key in ("A", "a", "R0", "R0/a", "a_major/a_minor"): + rtol, atol = 1e-3, 0 # need looser tol here bc of different grids else: rtol, atol = 1e-8, 0 - np.testing.assert_allclose(x, y, rtol=rtol, atol=atol, err_msg=name + key) + np.testing.assert_allclose( + x, y, rtol=rtol, atol=atol, err_msg=name + " " + key + ) # compare at rho=1, where we expect the eq.compute and the # surface.compute to agree for these surface basis vectors grid = LinearGrid(rho=np.array(1.0), M=10, N=10, NFP=eq.NFP) diff --git a/tests/test_configuration.py b/tests/test_configuration.py index fed3c0235..59c533ecc 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -462,9 +462,12 @@ def test_get_rho_surface(self): rho = 0.5 surf = eq.get_surface_at(rho=rho) assert surf.rho == rho - np.testing.assert_allclose(surf.compute("S")["S"], 4 * np.pi**2 * R0 * rho) + np.testing.assert_allclose( + surf.compute("S")["S"] * rho, 4 * np.pi**2 * R0 * rho + ) @pytest.mark.unit + @pytest.mark.xfail(reason="GitHub issue 1127.") def test_get_zeta_surface(self): """Test getting a constant zeta surface.""" eq = Equilibrium() diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 8488c91a8..08722e45c 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -89,7 +89,7 @@ def test_1d_p(self): def test_1d_elongation(self): """Test plotting 1d elongation as a function of toroidal angle.""" eq = get("precise_QA") - grid = LinearGrid(N=20, NFP=eq.NFP) + grid = LinearGrid(M=eq.M_grid, N=20, NFP=eq.NFP) fig, ax, data = plot_1d( eq, "a_major/a_minor", grid=grid, figsize=(4, 4), return_data=True ) diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index c6d3a2348..ae7b592d6 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -405,6 +405,7 @@ class TestZernikeRZToroidalSection: """Tests for ZernikeRZToroidalSection class.""" @pytest.mark.unit + @pytest.mark.xfail(reason="GitHub issue 1127.") def test_area(self): """Test calculation of surface area.""" s = ZernikeRZToroidalSection() diff --git a/tests/test_vmec.py b/tests/test_vmec.py index 0fef594b3..45dd8590f 100644 --- a/tests/test_vmec.py +++ b/tests/test_vmec.py @@ -475,12 +475,16 @@ def test_vmec_save_1(VMEC_save): ) np.testing.assert_allclose(vmec.variables["chipf"][:], desc.variables["chipf"][:]) np.testing.assert_allclose( - vmec.variables["Rmajor_p"][:], desc.variables["Rmajor_p"][:] + vmec.variables["Rmajor_p"][:], desc.variables["Rmajor_p"][:], rtol=3e-4 ) np.testing.assert_allclose( - vmec.variables["Aminor_p"][:], desc.variables["Aminor_p"][:] + vmec.variables["Aminor_p"][:], + desc.variables["Aminor_p"][:], + rtol=2e-4, + ) + np.testing.assert_allclose( + vmec.variables["aspect"][:], desc.variables["aspect"][:], rtol=4e-4 ) - np.testing.assert_allclose(vmec.variables["aspect"][:], desc.variables["aspect"][:]) np.testing.assert_allclose( vmec.variables["volume_p"][:], desc.variables["volume_p"][:], rtol=1e-5 )