Add aliases for data_index

desc/compute/_metric.py

@@ -1294,7 +1294,8 @@ def _g_sup_zz(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="g^rt",
+    # name="g^rt"
+    name=["g^" + i + j for i in ['r', 't', 'z'] for j in ['r', 't', 'z']],
     label="g^{\\rho\\theta}",
     units="m^{-2}",
     units_long="inverse square meters",
@@ -1307,7 +1308,12 @@ def _g_sup_zz(params, transforms, profiles, data, **kwargs):
     data=["e^rho", "e^theta"],
 )
 def _g_sup_rt(params, transforms, profiles, data, **kwargs):
-    data["g^rt"] = dot(data["e^rho"], data["e^theta"])
+    # data["g^rt"] = dot(data["e^rho"], data["e^theta"])
+    # reruns each time but looks cleaner this way i think
+    aliases = ["g^" + i + j for i in ['r', 't', 'z'] for j in ['r', 't', 'z']]
+    for alias in aliases:
+        data[alias] = dot(data["e^rho"], data["e^theta"])
+
     return data
 
 

desc/compute/data_index.py

@@ -88,7 +88,12 @@ def _decorator(func):
                         raise ValueError(
                             f"Already registered function with parameterization {p} and name {name}."
                         )
-                    data_index[base_class][name] = d.copy()
+                    if isinstance(name, list):
+                        data_index[base_class][name[0]] = d.copy()
+                        for alias in name:
+                            data_index[base_class][alias] = data_index[base_class][name[0]]
+                    else:
+                        data_index[base_class][name] = d.copy()
                     flag = True
             if not flag:
                 raise ValueError(

tests/test_compute_funs.py

@@ -43,6 +43,16 @@ def myconvolve_2d(arr_1d, stencil, shape):
     return conv
 
 
+@pytest.mark.unit
+def test_aliases():
+    eq = Equilibrium()  # torus
+    rho = np.linspace(0, 1, 64)
+    grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym, rho=rho)
+    aliases = ["g^" + i + j for i in ['r', 't', 'z'] for j in ['r', 't', 'z']]
+    data = eq.compute(aliases, grid=grid)              
+    np.testing.assert_allclose([data[alias] for alias in aliases])
+
+
 @pytest.mark.unit
 def test_total_volume(DummyStellarator):
     """Test that the volume enclosed by the LCFS is equal to the total volume."""