correctly read in power_series_I vmec input files

Author: dpanici

desc/input_reader.py

@@ -1101,7 +1101,7 @@ def parse_vmec_inputs(vmec_fname, threshold=0):  # noqa: C901
                 )
             )
         ctype = vmec_indata.get("PCURR_TYPE", "power_series")
-        if ctype.lower() != "power_series":
+        if not ctype.lower() in ["power_series", "power_series_i"]:
             warnings.warn(
                 colored(
                     "current is not a power series! DESC can only read power series"
@@ -1135,7 +1135,11 @@ def parse_vmec_inputs(vmec_fname, threshold=0):  # noqa: C901
         # read current
         curr_tor = vmec_indata.get("CURTOR", None)
         AC = np.atleast_1d(vmec_indata.get("AC", np.array([0.0]))).astype(float)
-        ls = np.arange(0, AC.size) * 2
+        ls = (
+            np.arange(0, AC.size) * 2
+            if ctype.lower == "power_series"
+            else np.arange(1, AC.size + 1) * 2
+        )
         inputs["current"] = np.vstack([ls, AC]).T
 
         # axis
@@ -1511,16 +1515,17 @@ def parse_vmec_inputs(vmec_fname, threshold=0):  # noqa: C901
         # scale pressure profile
         inputs["pressure"][:, 1] *= pres_scale
         if not iota_flag:
-            # integrate current profile wrt s=rho^2
-            inputs["current"] = np.pad(
-                np.vstack(
-                    (
-                        inputs["current"][:, 0] + 2,
-                        inputs["current"][:, 1] * 2 / (inputs["current"][:, 0] + 2),
-                    )
-                ).T,
-                ((1, 0), (0, 0)),
-            )
+            if ctype == "power_series":
+                # integrate current derivative profile wrt s=rho^2
+                inputs["current"] = np.pad(
+                    np.vstack(
+                        (
+                            inputs["current"][:, 0] + 2,
+                            inputs["current"][:, 1] * 2 / (inputs["current"][:, 0] + 2),
+                        )
+                    ).T,
+                    ((1, 0), (0, 0)),
+                )
             # scale current profile
             if curr_tor is not None:
                 inputs["current"][:, 1] *= curr_tor / (

tests/test_vmec.py

@@ -1697,3 +1697,30 @@ def test_write_vmec_input(TmpDir):
     np.testing.assert_allclose(W0, W1)
     np.testing.assert_allclose(rho_err, 0, atol=1e-4)
     np.testing.assert_allclose(theta_err, 0, atol=1e-4)
+
+
+@pytest.mark.unit
+def test_write_vmec_input_current_power_series(TmpDir):
+    """Test generated VMEC input file gives the original equilibrium profiles."""
+    # write VMEC input file
+    fname = str(TmpDir.join("input.SOLOVEV_current"))
+    eq0 = get("SOLOVEV")
+    current_profile = eq0.get_profile("current", kind="power_series")
+    current_profile.params[0] = 0.0  # zero out the axis current
+    with pytest.warns(UserWarning, match="Setting toroidal"):
+        eq0.current = current_profile
+    VMECIO.write_vmec_input(eq0, fname)
+    eq1 = Equilibrium.from_input_file(fname, spectral_indexing="fringe")
+    # to match the original eq's L, which is not default
+    eq1.change_resolution(L=2 * eq1.M, L_grid=2 * eq1.M_grid)
+
+    # check that loaded eq matches original equilibrium
+    np.testing.assert_allclose(eq0.L, eq1.L)
+    np.testing.assert_allclose(eq0.M, eq1.M)
+    np.testing.assert_allclose(eq0.N, eq1.N)
+    np.testing.assert_allclose(eq0.NFP, eq1.NFP)
+    np.testing.assert_allclose(eq0.Psi, eq1.Psi)
+    np.testing.assert_allclose(eq0.p_l, eq1.p_l)
+    np.testing.assert_allclose(eq0.c_l, eq1.c_l, atol=1e-5)
+    np.testing.assert_allclose(eq0.Rb_lmn, eq1.Rb_lmn)
+    np.testing.assert_allclose(eq0.Zb_lmn, eq1.Zb_lmn)