Implement from_input_file for Equilibrium, update docs with this …

…info (#1327)

resolves #1298

CHANGELOG.md

@@ -18,6 +18,8 @@ to compute the toroidal flux when possible, as opposed to a 2D surface integral
   - if an ``int``, this is the chunk size to be used.
   - if ``"auto"`` for the ``ObjectiveFunction``, will use a heuristic for the maximum ``jac_chunk_size`` needed to fit the jacobian calculation on the available device memory, according to the formula: ``max_jac_chunk_size = (desc_config.get("avail_mem") / estimated_memory_usage - 0.22)  / 0.85  * self.dim_x`` with ``estimated_memory_usage = 2.4e-7 * self.dim_f * self.dim_x + 1``
 - the ``ObjectiveFunction`` ``jac_chunk_size`` is used if ``deriv_mode="batched"``, and the ``_Objective`` ``jac_chunk_size`` will be used if ``deriv_mode="blocked"``
+- Add ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file
+
 
 
 Bug Fixes

desc/equilibrium/equilibrium.py

@@ -2,6 +2,7 @@
 
 import copy
 import numbers
+import warnings
 from collections.abc import MutableSequence
 
 import numpy as np
@@ -26,6 +27,7 @@
     ZernikeRZToroidalSection,
 )
 from desc.grid import Grid, LinearGrid, QuadratureGrid, _Grid
+from desc.input_reader import InputReader
 from desc.io import IOAble
 from desc.objectives import (
     ForceBalance,
@@ -2035,6 +2037,55 @@ def from_near_axis(
 
         return eq
 
+    @classmethod
+    def from_input_file(cls, path, **kwargs):
+        """Create an Equilibrium from information in a DESC or VMEC input file.
+
+        Parameters
+        ----------
+        path : Path-like or str
+            Path to DESC or VMEC input file.
+        **kwargs : dict, optional
+            keyword arguments to pass to the constructor of the
+            Equilibrium being created.
+
+        Returns
+        -------
+        Equilibrium : Equilibrium
+            Equilibrium generated from the given input file.
+
+        """
+        inputs = InputReader().parse_inputs(path)[-1]
+        if (inputs["bdry_ratio"] is not None) and (inputs["bdry_ratio"] != 1):
+            warnings.warn(
+                "`bdry_ratio` is intended as an input for the continuation method."
+                "`bdry_ratio`=1 uses the given surface modes as is, any other scalar "
+                "value will scale the non-axisymmetric  modes by that value. The "
+                "final value of `bdry_ratio` in the input file is "
+                f"{inputs['bdry_ratio']}, this means the created Equilibrium won't "
+                "have the given surface but a scaled version instead."
+            )
+        inputs["surface"][:, 1:3] = inputs["surface"][:, 1:3].astype(int)
+        # remove the keys (pertaining to continuation and solver tols)
+        # that an Equilibrium does not need
+        unused_keys = [
+            "pres_ratio",
+            "bdry_ratio",
+            "pert_order",
+            "ftol",
+            "xtol",
+            "gtol",
+            "maxiter",
+            "objective",
+            "optimizer",
+            "bdry_mode",
+            "output_path",
+            "verbose",
+        ]
+        [inputs.pop(key) for key in unused_keys]
+        inputs.update(kwargs)
+        return cls(**inputs)
+
     def solve(
         self,
         objective="force",

desc/geometry/surface.py

@@ -320,9 +320,13 @@ def from_input_file(cls, path, **kwargs):
             inputs = InputReader().parse_inputs(path)[-1]
         if (inputs["bdry_ratio"] is not None) and (inputs["bdry_ratio"] != 1):
             warnings.warn(
-                "boundary_ratio = {} != 1, surface may not be as expected".format(
-                    inputs["bdry_ratio"]
-                )
+                "`bdry_ratio` is intended as an input for the continuation method."
+                "`bdry_ratio`=1 uses the given surface modes as is, any other  "
+                "scalar value will scale the non-axisymmetric  modes by that "
+                "value. The final value of `bdry_ratio` in the input file is "
+                f"{inputs['bdry_ratio']}, this means the created "
+                "FourierRZToroidalSurface will be a scaled version of the "
+                "input file boundary."
             )
         surf = cls(
             inputs["surface"][:, 3],

desc/input_reader.py

@@ -695,7 +695,7 @@ def write_desc_input(filename, inputs, header=""):  # noqa: C901 - fxn too compl
         f.write("# global parameters\n")
         f.write("sym = {:1d} \n".format(inputs[0]["sym"]))
         f.write("NFP = {:3d} \n".format(int(inputs[0]["NFP"])))
-        f.write("Psi = {:.8f} \n".format(inputs[0]["Psi"]))
+        f.write("Psi = {:.8e} \n".format(inputs[0]["Psi"]))
 
         f.write("\n# spectral resolution\n")
         for key, val in {
@@ -796,7 +796,7 @@ def desc_output_to_input(  # noqa: C901 - fxn too complex
         xtol=1e-6,
         gtol=1e-6,
         maxiter=100,
-        threshold=1e-10,
+        threshold=0,
     ):
         """Generate a DESC input file from a DESC output file.
 

docs/api_equilibrium.rst

@@ -16,6 +16,10 @@ An ``EquilibriaFamily`` is a ``list`` like object for storing multiple equilibri
     desc.equilibrium.Equilibrium
     desc.equilibrium.EquilibriaFamily
 
+The ``Equilibrium`` class may be instantiated in a couple of ways in addition to providing inputs to its constructor.
+- from an existing DESC or VMEC input file with its ``from_input_file`` method
+- from a ``pyQSC`` ``Qsc``  or ``pyQIC`` ``Qic`` near-axis equilibrium with the ``Equilibrium``'s ``from_near_axis`` method
+
 
 Geometry
 ********
@@ -35,6 +39,8 @@ the magnetic axis, cross section, and various space curves.
     desc.geometry.SplineXYZCurve
     desc.geometry.ZernikeRZToroidalSection
 
+The ``FourierRZToroidalSurface`` and the ``FourierRZCurve`` classes may be instantiated from an existing DESC or VMEC input file with their ``from_input_file`` method.
+
 
 Profiles
 ********

tests/inputs/input.DSHAPE_desc

@@ -1,11 +1,12 @@
 # This DESC input file was auto generated from the VMEC input file
-# /DESC/tests/inputs/input.DSHAPE
+# .//tests//inputs//input.DSHAPE
+# on 10/30/2024 at 14:22:23.
 # For details on the various options see https://desc-docs.readthedocs.io/en/stable/input.html
 
 # global parameters
 sym = 1
 NFP =   1
-Psi = 1.00000000
+Psi = 1.00000000e+00
 
 # spectral resolution
 L_rad = 11

tests/test_input_output.py

@@ -57,10 +57,10 @@ def test_vmec_input(tmpdir_factory):
         lines_direct = f.readlines()
     with open(path_converted_file) as f:
         lines_converted = f.readlines()
-    # skip first 3 lines as they have date and pwd info
-    for line1, line2 in zip(lines_correct[3:], lines_converted[4:]):
+    # skip first 4 lines as they have date and pwd info
+    for line1, line2 in zip(lines_correct[4:], lines_converted[4:]):
         assert line1.strip() == line2.strip()
-    for line1, line2 in zip(lines_correct[3:], lines_direct):
+    for line1, line2 in zip(lines_correct[4:], lines_direct):
         assert line1.strip() == line2.strip()
 
 
@@ -226,6 +226,72 @@ def test_near_axis_input_files():
         os.remove(".//tests//inputs//input.QSC_r2_5.5_vmec_desc")
 
 
+@pytest.mark.unit
+def test_from_input_file_equilibrium_desc_vmec_DSHAPE():
+    """Test that from_input_file works for DESC input files."""
+    vmec_path = ".//tests//inputs//input.DSHAPE"
+    desc_path = ".//tests//inputs//input.DSHAPE_desc"
+    kwargs = {"spectral_indexing": "fringe"}
+    with pytest.warns(UserWarning, match="Left handed"):
+        eq = Equilibrium.from_input_file(desc_path, **kwargs)
+    with pytest.warns(UserWarning):
+        eq_VMEC = Equilibrium.from_input_file(vmec_path, **kwargs)
+
+    # make sure the loaded eqs are equivalent
+    np.testing.assert_allclose(eq.R_lmn, eq_VMEC.R_lmn)
+    np.testing.assert_allclose(eq.Z_lmn, eq_VMEC.Z_lmn)
+    np.testing.assert_allclose(eq.L_lmn, eq_VMEC.L_lmn)
+    np.testing.assert_allclose(eq.Rb_lmn, eq_VMEC.Rb_lmn)
+    np.testing.assert_allclose(eq.Zb_lmn, eq_VMEC.Zb_lmn)
+    np.testing.assert_allclose(eq.Ra_n, eq_VMEC.Ra_n)
+    np.testing.assert_allclose(eq.Za_n, eq_VMEC.Za_n)
+    np.testing.assert_allclose(eq.Psi, eq_VMEC.Psi)
+    assert eq.pressure.equiv(eq_VMEC.pressure)
+    assert eq.iota.equiv(eq_VMEC.iota)
+    assert eq.current is None
+    assert eq_VMEC.current is None
+    assert eq.sym == eq_VMEC.sym
+
+    # check against the DSHAPE bdry and profiles
+    eq_example = desc.examples.get("DSHAPE")
+    eq.change_resolution(L=eq_example.L, M=eq_example.M)
+    np.testing.assert_allclose(eq.Rb_lmn, eq_example.Rb_lmn)
+    np.testing.assert_allclose(eq.Zb_lmn, eq_example.Zb_lmn)
+    np.testing.assert_allclose(eq.Psi, eq_example.Psi)
+    np.testing.assert_allclose(eq.p_l, eq_example.p_l)
+    # our example's iota is negative of this input files's
+    np.testing.assert_allclose(eq.i_l, -eq_example.i_l)
+    assert eq.sym == eq_example.sym
+
+
+@pytest.mark.unit
+def test_from_input_file_equilibrium_desc_vmec():
+    """Test that from_input_file gives same eq for DESC and VMEC input files."""
+    vmec_path = ".//tests//inputs//input.QSC_r2_5.5_vmec"
+    desc_path = ".//tests//inputs//input.QSC_r2_5.5_desc"
+    kwargs = {"L": 10, "M": 10, "N": 14}
+    eq = Equilibrium.from_input_file(desc_path, **kwargs)
+    with pytest.warns(UserWarning):
+        eq_VMEC = Equilibrium.from_input_file(vmec_path, **kwargs)
+
+    np.testing.assert_allclose(eq.R_lmn, eq_VMEC.R_lmn)
+    np.testing.assert_allclose(eq.Z_lmn, eq_VMEC.Z_lmn)
+    np.testing.assert_allclose(eq.L_lmn, eq_VMEC.L_lmn)
+    np.testing.assert_allclose(eq.Rb_lmn, eq_VMEC.Rb_lmn)
+    np.testing.assert_allclose(eq.Zb_lmn, eq_VMEC.Zb_lmn)
+    np.testing.assert_allclose(eq.Ra_n, eq_VMEC.Ra_n)
+    np.testing.assert_allclose(eq.Za_n, eq_VMEC.Za_n)
+    np.testing.assert_allclose(eq.Psi, eq_VMEC.Psi)
+    assert eq.pressure.equiv(eq_VMEC.pressure)
+    assert eq.current.equiv(eq_VMEC.current)
+    assert eq.iota is None
+    assert eq_VMEC.iota is None
+    assert eq.sym == eq_VMEC.sym
+
+    if os.path.exists(".//tests//inputs//input.QSC_r2_5.5_vmec_desc"):
+        os.remove(".//tests//inputs//input.QSC_r2_5.5_vmec_desc")
+
+
 @pytest.mark.unit
 def test_vmec_input_surface_threshold():
     """Test ."""