Fix norm bug in curvature computation

desc/compute/_basis_vectors.py

@@ -655,20 +655,20 @@ def _b(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="n",
-    label="\\hat{n}",
+    name="n_rho",
+    label="\\hat{\\mathbf{n}}_{\\rho}",
     units="~",
     units_long="None",
-    description="Unit vector normal to flux surface",
+    description="Unit vector normal to constant rho surface (direction of e^rho)",
     dim=3,
     params=[],
     transforms={},
     profiles=[],
     coordinates="rtz",
     data=["e^rho"],
 )
-def _n(params, transforms, profiles, data, **kwargs):
-    data["n"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"], axis=-1)).T
+def _n_rho(params, transforms, profiles, data, **kwargs):
+    data["n_rho"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"], axis=-1)).T
     return data
 
 

desc/compute/_field.py

@@ -2985,10 +2985,10 @@ def _kappa(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["kappa", "n"],
+    data=["kappa", "n_rho"],
 )
 def _kappa_n(params, transforms, profiles, data, **kwargs):
-    data["kappa_n"] = dot(data["kappa"], data["n"])
+    data["kappa_n"] = dot(data["kappa"], data["n_rho"])
     return data
 
 
@@ -3003,10 +3003,10 @@ def _kappa_n(params, transforms, profiles, data, **kwargs):
     transforms={},
     profiles=[],
     coordinates="rtz",
-    data=["kappa", "n", "b"],
+    data=["kappa", "n_rho", "b"],
 )
 def _kappa_g(params, transforms, profiles, data, **kwargs):
-    data["kappa_g"] = dot(data["kappa"], cross(data["n"], data["b"]))
+    data["kappa_g"] = dot(data["kappa"], cross(data["n_rho"], data["b"]))
     return data
 
 

desc/compute/_geometry.py

@@ -233,24 +233,6 @@ def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs):
     return data
 
 
-@register_compute_fun(
-    name="n_rho",
-    label="\\hat{\\mathbf{n}}_{\\rho}",
-    units="~",
-    units_long="None",
-    description="Unit normal vector to constant rho surface",
-    dim=3,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="rtz",
-    data=["e^rho"],
-)
-def _n_rho(params, transforms, profiles, data, **kwargs):
-    data["n_rho"] = (data["e^rho"].T / jnp.linalg.norm(data["e^rho"])).T
-    return data
-
-
 @register_compute_fun(
     name="L_sff",
     label="L_{sff}",

desc/plotting.py

@@ -2798,10 +2798,7 @@ def plot_basis(basis, return_data=False, **kwargs):
         fig, ax = plt.subplots(figsize=kwargs.get("figsize", (6, 4)))
 
         f = basis.evaluate(grid.nodes)
-        plot_data = {}
-        plot_data["n"] = basis.modes[:, 2]
-        plot_data["amplitude"] = []
-        plot_data["zeta"] = z
+        plot_data = {"n": basis.modes[:, 2], "amplitude": [], "zeta": z}
 
         for fi, n in zip(f.T, basis.modes[:, 2]):
             ax.plot(z, fi, label="$n={:d}$".format(int(n)))
@@ -2841,12 +2838,13 @@ def plot_basis(basis, return_data=False, **kwargs):
         )
         ax = np.empty((2 * mmax + 1, 2 * nmax + 1), dtype=object)
         f = basis.evaluate(grid.nodes)
-        plot_data = {}
-        plot_data["m"] = basis.modes[:, 1]
-        plot_data["n"] = basis.modes[:, 2]
-        plot_data["amplitude"] = []
-        plot_data["zeta"] = z
-        plot_data["theta"] = t
+        plot_data = {
+            "m": basis.modes[:, 1],
+            "n": basis.modes[:, 2],
+            "amplitude": [],
+            "zeta": z,
+            "theta": t,
+        }
 
         for fi, m, n in zip(f.T, basis.modes[:, 1], basis.modes[:, 2]):
             ax[mmax + m, nmax + n] = plt.subplot(gs[mmax + m + 1, n + nmax])

docs/variables.csv

@@ -34,7 +34,7 @@ Name,Label,Units,Description,Module
 ``e^theta``,:math:`\mathbf{e}^{\theta}`,inverse meters,Contravariant poloidal basis vector,``desc.compute._basis_vectors``
 ``e^zeta``,:math:`\mathbf{e}^{\zeta}`,inverse meters,Contravariant toroidal basis vector,``desc.compute._basis_vectors``
 ``b``,:math:`\hat{b}`,None,Unit vector along magnetic field,``desc.compute._basis_vectors``
-``n``,:math:`\hat{n}`,None,Unit vector normal to flux surface,``desc.compute._basis_vectors``
+``n_rho``,:math:`\hat{\mathbf{n}}_{\rho}`,None,Unit vector normal to constant rho surface (direction of e^rho),``desc.compute._basis_vectors``
 ``grad(alpha)``,:math:`\nabla \alpha`,Inverse meters,Unit vector normal to flux surface,``desc.compute._basis_vectors``
 ``trapped fraction``,:math:`1 - \frac{3}{4} \langle B^2 \rangle \int_0^{1/Bmax} \frac{\lambda\; d\lambda}{\langle \sqrt{1 - \lambda B} \rangle}`,None,Neoclassical effective trapped particle fraction,``desc.compute._bootstrap``
 ``<J*B> Redl``,:math:`\langle\mathbf{J}\cdot\mathbf{B}\rangle_{Redl}`,Tesla Ampere / meter^2,"Bootstrap current profile, Redl model for quasisymmetry",``desc.compute._bootstrap``
@@ -55,7 +55,7 @@ Name,Label,Units,Description,Module
 ``theta_sfl_r``,:math:`\partial_{\rho} \vartheta`,radians,"PEST straight field line poloidal angular coordinate, derivative wrt radial coordinate",``desc.compute._core``
 ``theta_sfl_t``,:math:`\partial_{\theta} \vartheta`,radians,"PEST straight field line poloidal angular coordinate, derivative wrt poloidal coordinate",``desc.compute._core``
 ``theta_sfl_z``,:math:`\partial_{\zeta} \vartheta`,radians,"PEST straight field line poloidal angular coordinate, derivative wrt toroidal coordinate",``desc.compute._core``
-``alpha``,:math:`\alpha`,None,Field line label,``desc.compute._core``
+``alpha``,:math:`\alpha`,None,"Field line label, defined on [0, 2pi)",``desc.compute._core``
 ``alpha_r``,:math:`\partial_\rho \alpha`,None,"Field line label, derivative wrt radial coordinate",``desc.compute._core``
 ``alpha_t``,:math:`\partial_\theta \alpha`,None,"Field line label, derivative wrt poloidal coordinate",``desc.compute._core``
 ``alpha_z``,:math:`\partial_\zeta \alpha`,None,"Field line label, derivative wrt toroidal coordinate",``desc.compute._core``
@@ -298,7 +298,6 @@ Name,Label,Units,Description,Module
 ``a``,:math:`a`,meters,Average minor radius,``desc.compute._geometry``
 ``R0/a``,:math:`R_{0} / a`,None,Aspect ratio,``desc.compute._geometry``
 ``a_major/a_minor``,:math:`a_{major} / a_{minor}`,None,Maximum elongation,``desc.compute._geometry``
-``n_rho``,:math:`\hat{\mathbf{n}}_{\rho}`,None,Unit normal vector to constant rho surface,``desc.compute._geometry``
 ``L_sff``,:math:`L_{sff}`,meters,L coefficient of second fundamental form,``desc.compute._geometry``
 ``M_sff``,:math:`M_{sff}`,meters,M coefficient of second fundamental form,``desc.compute._geometry``
 ``N_sff``,:math:`N_{sff}`,meters,N coefficient of second fundamental form,``desc.compute._geometry``

tests/test_surfaces.py

@@ -208,7 +208,7 @@ def test_misc(self):
     def test_curvature(self):
         """Tests for gaussian, mean, principle curvatures.
 
-        (kind of pointless since its a flat surface so its always 0)
+        (kind of pointless since it's a flat surface so its always 0)
         """
         s = ZernikeRZToroidalSection()
         grid = LinearGrid(theta=np.pi / 2, rho=0.5)