Integrate on boundary to compute length scale quantities

desc/compute/_omnigenity.py

@@ -345,7 +345,7 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs):
     N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
 )
 def _B_mn(params, transforms, profiles, data, **kwargs):
-    norm = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1))
+    norm = 2 ** (3 - jnp.sum(transforms["B"].basis.modes == 0, axis=1))
     grid = transforms["grid"]
 
     def fun(rho, theta_B, zeta_B, sqrtg_B, B):

desc/equilibrium/equilibrium.py

@@ -982,7 +982,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."
@@ -991,23 +990,34 @@ 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") + f" got L_grid={grid.L} < {self._L_grid}.",
                 )
                 warnif(
                     # 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") + f" got M_grid={grid.M} < {self._M_grid}.",
                 )
                 warnif(
                     # 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") + f" got N_grid={grid.N} < {self._N_grid}.",
                 )

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
@@ -218,7 +218,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
@@ -945,7 +951,11 @@ class LinearGrid(_Grid):
         Change this only if your nodes are placed within one field period
         or should be interpreted as spanning one field period.
     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
@@ -1401,7 +1411,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``}

desc/integrals/bounce_integral.py

@@ -281,7 +281,7 @@ class Bounce2D(Bounce):
     # However, the basis for the latter are trigonometric functions with
     # irrational frequencies, courtesy of the irrational rotational transform.
     # Globally convergent root-finding schemes for that basis (at fixed α) are
-    # not known. The denominator of a close rational could be absorbed into the
+    # not efficient. The denominator of a close rational could be absorbed into the
     # coordinate ϕ, but this balloons the frequency, and hence degree of the series.
     #
     # Quadrature is chosen over Runge-Kutta methods of the form
@@ -1211,10 +1211,9 @@ def integrate(
             as the indices that correspond to that field line.
         data : dict[str, jnp.ndarray]
             Shape (num rho, num alpha, num zeta).
-            Real scalar-valued periodic functions in (θ, ζ) ∈ [0, 2π) × [0, 2π/NFP)
-            evaluated on the ``grid`` supplied to construct this object.
-            Use the method ``Bounce1D.reshape`` to reshape the data into the
-            expected shape.
+            Real scalar-valued functions evaluated on the ``grid`` supplied to
+            construct this object. Use the method ``Bounce1D.reshape`` to
+            reshape the data into the expected shape.
         names : str or list[str]
             Names in ``data`` to interpolate. Default is all keys in ``data``.
         points : tuple[jnp.ndarray]

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

desc/objectives/_neoclassical.py

@@ -251,7 +251,6 @@ def compute(self, params, constants=None):
             Effective ripple as a function of the flux surface label.
 
         """
-        # TODO (#1094)
         if constants is None:
             constants = self.constants
         eq = self.things[0]

desc/objectives/_stability.py

@@ -471,6 +471,9 @@ def build(self, eq=None, use_jit=True, verbose=1):
         iota_profiles = get_profiles(self._iota_keys, obj=eq, grid=iota_grid)
         iota_transforms = get_transforms(self._iota_keys, obj=eq, grid=iota_grid)
 
+        # TODO: Generalize balloning stabilty funs to multiple flux surfaces,
+        #       include last closed flux surface requirement, and remove quadrature
+        #       transforms.
         # Separate grid to calculate the right length scale for normalization
         len_grid = QuadratureGrid(L=eq.L, M=eq.M, N=eq.N, NFP=eq.NFP)
 

desc/plotting.py

@@ -12,7 +12,6 @@
 from matplotlib import cycler, rcParams
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from pylatexenc.latex2text import LatexNodes2Text
-from termcolor import colored
 
 from desc.backend import sign
 from desc.basis import fourier, zernike_radial_poly
@@ -187,15 +186,12 @@
         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,
+            "ax argument must be None or an axis instance or array of axes",
+        )
+        return plt.gcf(), ax
 
 
 def _get_grid(**kwargs):
@@ -278,11 +274,10 @@
 
     """
     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",
@@ -294,9 +289,7 @@
 
     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:
@@ -520,47 +513,50 @@
         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="Grid must be 1D",
+        )
+        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()
@@ -641,8 +637,7 @@
         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="Grid must be 2D")
     component = kwargs.pop("component", None)
     if name != "B*n":
         data, label = _compute(
@@ -655,14 +650,12 @@
         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.",
         )
 
@@ -947,14 +940,12 @@
         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.",
         )
 
@@ -985,8 +976,7 @@
 
     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")
@@ -2497,8 +2487,7 @@
     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),
@@ -2925,7 +2914,7 @@
         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(