Ku/available energy for turbulence transport

CHANGELOG.md

@@ -3,6 +3,7 @@ Changelog
 
 New Features
 
+- Adds Available Energy for nonlinear measure of turbelent transport [#2215](https://github.com/PlasmaControl/DESC/pull/2215) and local potential well plotting utitlity for analysis.
 - Adds ``desc.objectives.DeflationOperator``, a new objective class which can be used to apply deflation techniques to equilibrium and optimization problems to find multiple local minima or multiple solutions from a single initial point, either by wrapping an existing ``desc.objectives._Objective`` object or by including as an additional penalty or constraint. Also adds a tutorial showing this functionality.
 - Sub-objectives of an `ObjectiveFunction` can now have different `use_jit` values than the `ObjectiveFunction`. These objectives have to be built before building the `ObjectiveFunction`.
 - Adds ``num_neighbors`` parameter to ``CoilSetMinDistance`` that limits the pairwise distance computation to the nearest neighbors per coil, reducing memory useage for large coilsets.
@@ -30,6 +31,7 @@ Performance Improvements
 - Reduces import time of `desc` modules.
     - Now, `desc.compute._build_data_index` uses depth-first search algorithm to construct the dependency tree.
     - Some of the default value computations at import time are removed (i.e. `desc.integrals.bounce_integral.default_quad`)
+- Significantly improves convergence of bounce point computations. Now lower ``Y_B`` can be used.
 - [Significantly improves convergence of inverse stream maps](https://github.com/PlasmaControl/DESC/pull/1919).
 - Resolves a JAX memory regression in bounce integrals by avoiding materialization of a large tensor in memory. Previously, we had closed the issue by adding nuffts as a workaround. This update actually solves the issue for the case when a user specifies to not use nuffts as well.
 - ``ObjectiveFunction.print_value`` can now use the previously computed ``compute_scaled_error`` values to print. For bounded objectives, we fall back to computing ``compute_unscaled``. Additionally, ``compute_scaled_error`` and array splitting are used in other parts of the code to prevent recompilation for one-time tasks, which makes initialization faster.

desc/compute/__init__.py

@@ -42,6 +42,7 @@
     _profiles,
     _stability,
     _surface,
+    _turbulence,
 )
 from .data_index import all_kwargs, allowed_kwargs, data_index
 from .utils import (

desc/compute/_fast_ion.py

@@ -33,7 +33,7 @@ def _gamma_c_data(data):
     # units Tesla meters. Smoothness is determined by positive lower bound of
     # log argument, and hence behaves as ∂log(|B|/B₀)/∂ρ |B| = ∂|B|/∂ρ.
     return {
-        "|grad(psi)|*kappa_g": data["|grad(psi)|"] * data["kappa_g"],
+        "|grad(psi)|*kappa_g": data["|grad(psi)|*kappa_g"],
         "|grad(rho)|*|e_alpha|r,p|": data["|grad(rho)|"] * data["|e_alpha|r,p|"],
         "|B|_r|v,p": data["|B|_r|v,p"],
         "K": data["iota_r"]
@@ -48,7 +48,7 @@ def _v_tau(data, B, pitch):
     return safediv(2.0, jnp.sqrt(jnp.abs(1 - pitch * B)))
 
 
-def _radial_drift_1(data, B, pitch):
+def _radial_drift(data, B, pitch):
     return (
         safediv(1 - 0.5 * pitch * B, jnp.sqrt(jnp.abs(1 - pitch * B)))
         * data["|grad(psi)|*kappa_g"]
@@ -57,21 +57,18 @@ def _radial_drift_1(data, B, pitch):
 
 
 def _poloidal_drift_periodic(data, B, pitch):
-    return (
-        safediv(1 - 0.5 * pitch * B, jnp.sqrt(jnp.abs(1 - pitch * B)))
-        * data["|B|_r|v,p"]
-        + jnp.sqrt(jnp.abs(1 - pitch * B)) * data["K"]
-    ) / B
+    g = jnp.sqrt(jnp.abs(1 - pitch * B))
+    return (safediv(1 - 0.5 * pitch * B, g) * data["|B|_r|v,p"] + g * data["K"]) / B
 
 
-def _radial_drift_2(data, B, pitch):
+def _radial_drift_wb_inverse(data, B, pitch):
     return safediv(
         data["cvdrift0"] * (1 - 0.5 * pitch * B),
         jnp.sqrt(jnp.abs(1 - pitch * B)),
     )
 
 
-def _poloidal_drift_secular(data, B, pitch):
+def _poloidal_drift_secular_wb_inverse(data, B, pitch):
     # TODO (#465), multiply by (omega + zeta) instead of zeta
     return safediv(
         (data["gbdrift (periodic)"] + data["gbdrift (secular)/phi"] * data["zeta"])
@@ -104,10 +101,9 @@ def _poloidal_drift_secular(data, B, pitch):
         "b",
         "grad(phi)",
         "grad(psi)",
-        "|grad(psi)|",
         "|grad(rho)|",
         "|e_alpha|r,p|",
-        "kappa_g",
+        "|grad(psi)|*kappa_g",
         "iota_r",
         "V_psi",
     ]
@@ -153,7 +149,7 @@ def Gamma_c(data):
         def fun(pitch_inv):
             points = bounce.points(pitch_inv, opts.num_well)
             v_tau, radial_drift, poloidal_drift = bounce.integrate(
-                [_v_tau, _radial_drift_1, _poloidal_drift_periodic],
+                [_v_tau, _radial_drift, _poloidal_drift_periodic],
                 pitch_inv,
                 data,
                 ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
@@ -182,9 +178,8 @@ def fun(pitch_inv):
         Gamma_c, _gamma_c_data(data), data, angle, grid, opts.surf_batch_size
     )
     assert out.ndim == 1
-    data["Gamma_c"] = (
-        grid.expand(out) / data["V_psi"] / (opts.num_field_periods / grid.NFP * 2**0.5)
-    )
+    scalar = opts.num_field_periods / grid.NFP * 2**0.5
+    data["Gamma_c"] = grid.expand(out / scalar) / data["V_psi"]
     return data
 
 
@@ -208,10 +203,9 @@ def fun(pitch_inv):
         "b",
         "grad(phi)",
         "grad(psi)",
-        "|grad(psi)|",
         "|grad(rho)|",
         "|e_alpha|r,p|",
-        "kappa_g",
+        "|grad(psi)|*kappa_g",
         "iota_r",
     ]
     + Bounce2D.required_names,
@@ -243,7 +237,7 @@ def gamma_c0(data):
         bounce = Bounce2D(grid, data, data["angle"], **opts)
         points = bounce.points(pitch_inv, opts.num_well)
         radial_drift, poloidal_drift = bounce.integrate(
-            [_radial_drift_1, _poloidal_drift_periodic],
+            [_radial_drift, _poloidal_drift_periodic],
             pitch_inv,
             data,
             ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
@@ -324,7 +318,7 @@ def Gamma_c(data):
 
         def fun(pitch_inv):
             v_tau, radial_drift, poloidal_drift = bounce.integrate(
-                [_v_tau, _radial_drift_2, _poloidal_drift_secular],
+                [_v_tau, _radial_drift_wb_inverse, _poloidal_drift_secular_wb_inverse],
                 pitch_inv,
                 data,
                 ["cvdrift0", "gbdrift (periodic)", "gbdrift (secular)/phi"],
@@ -354,7 +348,7 @@ def fun(pitch_inv):
         opts.surf_batch_size,
     )
     assert out.ndim == 1
-    data["Gamma_c Velasco"] = (
-        grid.expand(out) / data["V_psi"] / (opts.num_field_periods / grid.NFP * 2**0.5)
-    )
+
+    scalar = opts.num_field_periods / grid.NFP * 2**0.5
+    data["Gamma_c Velasco"] = grid.expand(out / scalar) / data["V_psi"]
     return data

desc/compute/_metric.py

@@ -1797,6 +1797,24 @@ def _gradrho(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="|grad(rho)|*kappa_g",
+    label="|\\nabla\\rho| \\kappa_g",
+    units="m^{-2}",
+    units_long="Inverse square meters",
+    description="Radial coordinate gradient magnitude times geodesic curvature.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|grad(rho)|", "kappa_g"],
+)
+def _gradrho_mag_times_kappa_g(params, transforms, profiles, data, **kwargs):
+    data["|grad(rho)|*kappa_g"] = data["|grad(rho)|"] * data["kappa_g"]
+    return data
+
+
 @register_compute_fun(
     name="<|grad(rho)|>",  # same as S(r) / V_r(r)
     label="\\langle \\vert \\nabla \\rho \\vert \\rangle",
@@ -1841,6 +1859,24 @@ def _gradpsi_mag(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="|grad(psi)|*kappa_g",
+    label="|\\nabla\\psi| \\kappa_g",
+    units="T",
+    units_long="Tesla",
+    description="Toroidal flux gradient magnitude times geodesic curvature.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["|grad(psi)|", "kappa_g"],
+)
+def _gradpsi_mag_times_kappa_g(params, transforms, profiles, data, **kwargs):
+    data["|grad(psi)|*kappa_g"] = data["|grad(psi)|"] * data["kappa_g"]
+    return data
+
+
 @register_compute_fun(
     name="|grad(psi)|^2",
     label="|\\nabla\\psi|^{2}",

desc/compute/_neoclassical.py

@@ -7,7 +7,7 @@
 from ..batching import batch_map
 from ..integrals.bounce_integral import Bounce2D, Options
 from ..integrals.surface_integral import surface_integrals
-from ..utils import parse_argname_change, safediv
+from ..utils import apply, parse_argname_change, safediv
 from .data_index import register_compute_fun
 
 
@@ -64,9 +64,8 @@ def _dI_2(data, B, pitch):
     data=[
         "min_tz |B|",
         "max_tz |B|",
-        "kappa_g",
+        "|grad(rho)|*kappa_g",
         "R0",
-        "|grad(rho)|",
         "<|grad(rho)|>",
         "V_psi",
     ]
@@ -126,15 +125,15 @@ def fun(pitch_inv):
     )
     out = Bounce2D.batch(
         eps_32,
-        {"|grad(rho)|*kappa_g": data["|grad(rho)|"] * data["kappa_g"]},
+        apply(data, subset=("|grad(rho)|*kappa_g",)),
         data,
         angle,
         grid,
         opts.surf_batch_size,
     )
     assert out.ndim == 1
-    data["effective ripple 3/2"] = scalar * (
-        (B0 / data["<|grad(rho)|>"]) ** 2 * grid.expand(out) / data["V_psi"]
+    data["effective ripple 3/2"] = (
+        (B0 / data["<|grad(rho)|>"]) ** 2 * grid.expand(out * scalar) / data["V_psi"]
     )
     return data
 

desc/compute/_old.py

@@ -9,12 +9,12 @@
 from desc.backend import jit, jnp
 
 from ..integrals.bounce_integral import Bounce1D, Options
-from ..utils import safediv
+from ..utils import apply, safediv
 from ._fast_ion import (
     _gamma_c_data,
     _poloidal_drift_periodic,
-    _radial_drift_1,
-    _radial_drift_2,
+    _radial_drift,
+    _radial_drift_wb_inverse,
     _v_tau,
 )
 from ._neoclassical import _dI_1, _dI_2
@@ -43,9 +43,8 @@
     data=[
         "min_tz |B|",
         "max_tz |B|",
-        "kappa_g",
         "R0",
-        "|grad(rho)|",
+        "|grad(rho)|*kappa_g",
         "<|grad(rho)|>",
         "fieldline length",
     ]
@@ -85,16 +84,15 @@ def eps_32(data):
     scalar = jnp.pi / (8 * 2**0.5) * data["R0"] ** 2
     out = Bounce1D.batch(
         eps_32,
-        {"|grad(rho)|*kappa_g": data["|grad(rho)|"] * data["kappa_g"]},
+        apply(data, subset=("|grad(rho)|*kappa_g",)),
         data,
         grid,
         opts.surf_batch_size,
     )
     assert out.ndim == 1
     data["old effective ripple 3/2"] = (
         (B0 / data["<|grad(rho)|>"]) ** 2
-        * scalar
-        * grid.expand(out)
+        * grid.expand(out * scalar)
         / data["fieldline length"]
     )
     return data
@@ -151,10 +149,9 @@ def _effective_ripple_1D(params, transforms, profiles, data, **kwargs):
         "b",
         "grad(phi)",
         "grad(psi)",
-        "|grad(psi)|",
+        "|grad(psi)|*kappa_g",
         "|grad(rho)|",
         "|e_alpha|r,p|",
-        "kappa_g",
         "iota_r",
         "fieldline length",
     ]
@@ -196,7 +193,7 @@ def Gamma_c(data):
         bounce = Bounce1D(grid, data, opts.quad)
         points = bounce.points(pitch_inv, num_well)
         v_tau, radial_drift, poloidal_drift = bounce.integrate(
-            [_v_tau, _radial_drift_1, _poloidal_drift_periodic],
+            [_v_tau, _radial_drift, _poloidal_drift_periodic],
             pitch_inv,
             data,
             ["|grad(psi)|*kappa_g", "|B|_r|v,p", "K"],
@@ -217,7 +214,7 @@ def Gamma_c(data):
     out = Bounce1D.batch(Gamma_c, _gamma_c_data(data), data, grid, opts.surf_batch_size)
     assert out.ndim == 1
     data["old Gamma_c"] = (
-        grid.expand(out) / data["fieldline length"] / (2**1.5 * jnp.pi)
+        grid.expand(out / (2**1.5 * jnp.pi)) / data["fieldline length"]
     )
     return data
 
@@ -259,7 +256,7 @@ def _Gamma_c_Velasco_1D(params, transforms, profiles, data, **kwargs):
     opts = Options.guess(eta=-1, grid=grid, Y_B=grid.num_zeta, **kwargs)
     num_well = kwargs.get("num_well", -1)
 
-    def _poloidal_drift_secular(data, B, pitch):
+    def _poloidal_drift_secular_wb_inverse(data, B, pitch):
         return safediv(
             data["gbdrift"] * (1 - 0.5 * pitch * B),
             jnp.sqrt(jnp.abs(1 - pitch * B)),
@@ -270,7 +267,7 @@ def Gamma_c(data):
             data["min_tz |B|"], data["max_tz |B|"], opts.pitch_quad
         )
         v_tau, radial_drift, poloidal_drift = Bounce1D(grid, data, opts.quad).integrate(
-            [_v_tau, _radial_drift_2, _poloidal_drift_secular],
+            [_v_tau, _radial_drift_wb_inverse, _poloidal_drift_secular_wb_inverse],
             pitch_inv,
             data,
             ["cvdrift0", "gbdrift"],
@@ -291,6 +288,6 @@ def Gamma_c(data):
     )
     assert out.ndim == 1
     data["old Gamma_c Velasco"] = (
-        grid.expand(out) / data["fieldline length"] / (2**1.5 * jnp.pi)
+        grid.expand(out / (2**1.5 * jnp.pi)) / data["fieldline length"]
     )
     return data