Add pass through kwargs from plotting to compute (#1915)

- Only adds `compute_kwargs` option to "generic" plotting functions
`plot_{1d,2d,3d,section,fsa}`, the other ones generally already accept
kwargs relevant to their function (eg `plot_boozer_modes`)
- No tests yet, as I can't think of any that would be meaningful until
#1899 is resolved.

Resolves #1352

---------

Co-authored-by: Dario Panici <37969854+dpanici@users.noreply.github.com>
Co-authored-by: Yigit Gunsur Elmacioglu <102380275+YigitElma@users.noreply.github.com>

Author: 3 people

desc/plotting.py

@@ -266,7 +266,14 @@ def _get_plot_axes(grid):
     return tuple(plot_axes)
 
 
-def _compute(eq, name, grid, component=None, reshape=True):
+def _compute(
+    eq,
+    name,
+    grid,
+    component=None,
+    reshape=True,
+    compute_kwargs=None,
+):
     """Compute quantity specified by name on grid for Equilibrium eq.
 
     Parameters
@@ -279,6 +286,8 @@ def _compute(eq, name, grid, component=None, reshape=True):
         Grid of coordinates to evaluate at.
     component : str, optional
         For vector variables, which element to plot. Default is the norm of the vector.
+    compute_kwargs : dict, optional
+        Additional keyword arguments to pass to ``eq.compute``
 
     Returns
     -------
@@ -302,7 +311,9 @@ def _compute(eq, name, grid, component=None, reshape=True):
 
     label = data_index[parameterization][name]["label"]
 
-    data = eq.compute(name, grid=grid)[name]
+    compute_kwargs = setdefault(compute_kwargs, {})
+
+    data = eq.compute(name, grid=grid, **compute_kwargs)[name]
 
     if data_index[parameterization][name]["dim"] > 1:
         if component is None:
@@ -505,7 +516,15 @@ def plot_coefficients(eq, L=True, M=True, N=True, ax=None, **kwargs):
 
 
 def plot_1d(  # noqa : C901
-    eq, name, grid=None, log=False, normalize=None, ax=None, return_data=False, **kwargs
+    eq,
+    name,
+    grid=None,
+    log=False,
+    normalize=None,
+    ax=None,
+    return_data=False,
+    compute_kwargs=None,
+    **kwargs,
 ):
     """Plot 1D profiles.
 
@@ -525,6 +544,8 @@ def plot_1d(  # noqa : C901
         Axis to plot on.
     return_data : bool
         If True, return the data plotted as well as fig,ax
+    compute_kwargs : dict, optional
+        Additional keyword arguments to pass to ``eq.compute``.
     **kwargs : dict, optional
         Specify properties of the figure, axis, and plot appearance e.g.::
 
@@ -587,6 +608,7 @@ def plot_1d(  # noqa : C901
                 ax=ax,
                 return_data=return_data,
                 grid=grid,
+                compute_kwargs=compute_kwargs,
                 **kwargs,
             )
         rho = grid.nodes[:, 0]
@@ -601,6 +623,7 @@ def plot_1d(  # noqa : C901
                 ax=ax,
                 return_data=return_data,
                 grid=grid,
+                compute_kwargs=compute_kwargs,
                 **kwargs,
             )
 
@@ -614,11 +637,18 @@ def plot_1d(  # noqa : C901
     plot_axes = _get_plot_axes(grid)
 
     data, ylabel = _compute(
-        eq, name, grid, kwargs.pop("component", None), reshape=False
+        eq,
+        name,
+        grid,
+        kwargs.pop("component", None),
+        reshape=False,
+        compute_kwargs=compute_kwargs,
     )
 
     if normalize:
-        norm_data, _ = _compute(eq, normalize, grid, reshape=False)
+        norm_data, _ = _compute(
+            eq, normalize, grid, compute_kwargs=compute_kwargs, reshape=False
+        )
         data = data / np.nanmean(np.abs(norm_data))  # normalize
 
     # reshape data to 1D
@@ -688,7 +718,15 @@ def plot_1d(  # noqa : C901
 
 
 def plot_2d(  # noqa : C901
-    eq, name, grid=None, log=False, normalize=None, ax=None, return_data=False, **kwargs
+    eq,
+    name,
+    grid=None,
+    log=False,
+    normalize=None,
+    ax=None,
+    return_data=False,
+    compute_kwargs=None,
+    **kwargs,
 ):
     """Plot 2D cross-sections.
 
@@ -708,6 +746,8 @@ def plot_2d(  # noqa : C901
         Axis to plot on.
     return_data : bool
         If True, return the data plotted as well as fig,ax
+    compute_kwargs : dict, optional
+        Additional keyword arguments to pass to ``eq.compute``.
     **kwargs : dict, optional
         Specify properties of the figure, axis, and plot appearance e.g.::
 
@@ -787,6 +827,7 @@ def plot_2d(  # noqa : C901
             name,
             grid,
             component=component,
+            compute_kwargs=compute_kwargs,
         )
     else:
         data, label = _compute_Bn(
@@ -802,7 +843,13 @@ def plot_2d(  # noqa : C901
     divider = make_axes_locatable(ax)
 
     if normalize:
-        norm_data, _ = _compute(eq, normalize, grid, reshape=False)
+        norm_data, _ = _compute(
+            eq,
+            normalize,
+            grid,
+            reshape=False,
+            compute_kwargs=compute_kwargs,
+        )
         data = data / np.nanmean(np.abs(norm_data))  # normalize
 
     # reshape data to 2D
@@ -954,6 +1001,7 @@ def plot_3d(  # noqa : C901
     normalize=None,
     fig=None,
     return_data=False,
+    compute_kwargs=None,
     **kwargs,
 ):
     """Plot 3D surfaces.
@@ -974,6 +1022,8 @@ def plot_3d(  # noqa : C901
         Figure to plot on.
     return_data : bool
         If True, return the data plotted as well as fig,ax
+    compute_kwargs : dict, optional
+        Additional keyword arguments to pass to ``eq.compute``.
     **kwargs : dict, optional
         Specify properties of the figure, axis, and plot appearance e.g.::
 
@@ -1061,6 +1111,7 @@ def plot_3d(  # noqa : C901
             name,
             grid,
             component=component,
+            compute_kwargs=compute_kwargs,
         )
     else:
         data, label = _compute_Bn(
@@ -1224,6 +1275,7 @@ def plot_fsa(  # noqa: C901
     ax=None,
     return_data=False,
     grid=None,
+    compute_kwargs=None,
     **kwargs,
 ):
     """Plot flux surface averages of quantities.
@@ -1262,6 +1314,8 @@ def plot_fsa(  # noqa: C901
     grid : _Grid
         Grid to compute name on. If provided, the parameters
         ``rho``, ``M``, and ``N`` are ignored.
+    compute_kwargs : dict, optional
+        Additional keyword arguments to pass to ``eq.compute``.
     **kwargs : dict, optional
         Specify properties of the figure, axis, and plot appearance e.g.::
 
@@ -1354,7 +1408,12 @@ def plot_fsa(  # noqa: C901
             # desired surface average.
             name = "<" + name + ">"
     values, ylabel = _compute(
-        eq, name, grid, kwargs.pop("component", None), reshape=False
+        eq,
+        name,
+        grid,
+        kwargs.pop("component", None),
+        reshape=False,
+        compute_kwargs=compute_kwargs,
     )
     ylabel = ylabel.split("~")
     if (
@@ -1371,7 +1430,13 @@ def plot_fsa(  # noqa: C901
     else:
         compute_surface_averages = surface_averages_map(grid, expand_out=False)
         if with_sqrt_g:  # flux surface average
-            sqrt_g = _compute(eq, "sqrt(g)", grid, reshape=False)[0]
+            sqrt_g = _compute(
+                eq,
+                "sqrt(g)",
+                grid,
+                reshape=False,
+                compute_kwargs=compute_kwargs,
+            )[0]
             # Attempt to compute the magnetic axis limit.
             # Compute derivative depending on various naming schemes.
             # e.g. B -> B_r, V(r) -> V_r(r), S_r(r) -> S_rr(r)
@@ -1384,7 +1449,13 @@ def plot_fsa(  # noqa: C901
             )
             values_r = next(
                 (
-                    _compute(eq, x, grid, reshape=False)[0]
+                    _compute(
+                        eq,
+                        x,
+                        grid,
+                        reshape=False,
+                        compute_kwargs=compute_kwargs,
+                    )[0]
                     for x in schemes
                     if x in data_index[p]
                 ),
@@ -1393,7 +1464,10 @@ def plot_fsa(  # noqa: C901
             if (np.isfinite(values) & np.isfinite(values_r))[grid.axis].all():
                 # Otherwise cannot compute axis limit in this agnostic manner.
                 sqrt_g = grid.replace_at_axis(
-                    sqrt_g, _compute(eq, "sqrt(g)_r", grid, reshape=False)[0], copy=True
+                    sqrt_g,
+                    _compute(eq, "sqrt(g)_r", grid, reshape=False)[0],
+                    copy=True,
+                    compute_kwargs=compute_kwargs,
                 )
             averages = compute_surface_averages(values, sqrt_g=sqrt_g)
             ylabel = r"$\langle " + ylabel[0][1:] + r" \rangle~" + "~".join(ylabel[1:])
@@ -1413,7 +1487,13 @@ def plot_fsa(  # noqa: C901
         plot_data_ylabel_key = f"<{name}>_fsa"
 
     if normalize:
-        norm_data = _compute(eq, normalize, grid, reshape=False)[0]
+        norm_data = _compute(
+            eq,
+            normalize,
+            grid,
+            reshape=False,
+            compute_kwargs=compute_kwargs,
+        )[0]
         values = values / np.nanmean(np.abs(norm_data))  # normalize
     if log:
         values = np.abs(values)  # ensure data is positive for log plot
@@ -1455,7 +1535,15 @@ def plot_fsa(  # noqa: C901
 
 
 def plot_section(
-    eq, name, grid=None, log=False, normalize=None, ax=None, return_data=False, **kwargs
+    eq,
+    name,
+    grid=None,
+    log=False,
+    normalize=None,
+    ax=None,
+    return_data=False,
+    compute_kwargs=None,
+    **kwargs,
 ):
     """Plot Poincare sections.
 
@@ -1475,6 +1563,8 @@ def plot_section(
         Axis to plot on.
     return_data : bool
         If True, return the data plotted as well as fig,ax
+    compute_kwargs : dict, optional
+        Additional keyword arguments to pass to ``eq.compute``.
     **kwargs : dict, optional
         Specify properties of the figure, axis, and plot appearance e.g.::
 
@@ -1584,9 +1674,22 @@ def plot_section(
     rows = np.floor(np.sqrt(nphi)).astype(int)
     cols = np.ceil(nphi / rows).astype(int)
 
-    data, _ = _compute(eq, name, grid, kwargs.pop("component", None), reshape=False)
+    data, _ = _compute(
+        eq,
+        name,
+        grid,
+        kwargs.pop("component", None),
+        reshape=False,
+        compute_kwargs=compute_kwargs,
+    )
     if normalize:
-        norm_data, _ = _compute(eq, normalize, grid, reshape=False)
+        norm_data, _ = _compute(
+            eq,
+            normalize,
+            grid,
+            reshape=False,
+            compute_kwargs=compute_kwargs,
+        )
         data = data / np.nanmean(np.abs(norm_data))  # normalize
 
     figw = 5 * cols