formatting

PlasmaControl · Nov 13, 2023 · 8b8559a · 8b8559a
1 parent c2372ea
commit 8b8559a
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 8 deletions.
diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py
@@ -1853,6 +1853,7 @@ def _e_sub_rho_rtt(params, transforms, profiles, data, **kwargs):
         "omega_tz",
         "omega_z",
     ],
+    aliases=["x_rrt", "x_trr", "x_rtr"],
 )
 def _e_sub_rho_rtz(params, transforms, profiles, data, **kwargs):
     data["e_rho_rtz"] = jnp.array(

diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py
@@ -24,7 +24,21 @@ def find_permutations(name, separator="_"):
 def assign_alias_data(
     alias, primary, base_class, data_index, params, profiles, transforms, data, **kwargs
 ):
-    """Assigns primary data to alias."""
+    """Assigns primary data to alias.
+
+    Parameters
+    ----------
+    alias : `str`
+        data_index key for alias of primary
+    primary : `str`
+        key defined in compute function
+
+    Returns
+    -------
+    data : `dict`
+        computed data dictionary (includes both alias and primary)
+
+    """
     data = data_index[base_class][primary]["fun"](
         params, transforms, profiles, data, **kwargs
     )

diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py
@@ -49,17 +49,24 @@ def myconvolve_2d(arr_1d, stencil, shape):
 @pytest.mark.unit
 def test_aliases():
     """Tests that data_index aliases are equal."""
-    eq = Equilibrium()
+    surface = FourierRZToroidalSurface(
+        R_lmn=[10, 1, 0.2],
+        Z_lmn=[-2, -0.2],
+        modes_R=[[0, 0], [1, 0], [0, 1]],
+        modes_Z=[[-1, 0], [0, -1]],
+    )
+
+    eq = Equilibrium(surface=surface)
 
     # automatic case
-    data_1 = eq.compute("R_tz")
-    data_2 = eq.compute("R_zt")
-    np.testing.assert_allclose(data_1["R_tz"], data_2["R_tz"])
+    primary_data = eq.compute("R_tz")
+    alias_data = eq.compute("R_zt")
+    np.testing.assert_allclose(primary_data["R_tz"], alias_data["R_zt"])
 
     # manual case
-    data_1 = eq.compute("x_rrt")
-    data_2 = eq.compute("e_rho_rt")
-    np.testing.assert_allclose(data_1["e_rho_rt"], data_2["e_rho_rt"])
+    primary_data = eq.compute("e_rho_rt")
+    alias_data = eq.compute("x_rrt")
+    np.testing.assert_allclose(primary_data["e_rho_rt"], alias_data["x_rrt"])
 
 
 @pytest.mark.unit