Add a method to load a VmecWout object from a NetCDF file.

Author: jons-pf

src/vmecpp/__init__.py

@@ -830,7 +830,31 @@ def _to_cpp_wout(self) -> _vmecpp.WOutFileContents:
 
         return cpp_wout
 
-    # TODO(eguiraud): implement from_wout_file
+    @staticmethod
+    def from_wout_file(wout_filename: str | Path) -> VmecWOut:
+        """Load wout contents in NetCDF format.
+
+        This is the format used by Fortran VMEC implementations and the one expected by
+        SIMSOPT.
+        """
+        with netCDF4.Dataset(wout_filename, "r") as fnc:
+            fnc.set_auto_mask(False)
+            attrs = {}
+            for key in fnc.variables:
+                if key.endswith("__logical__"):
+                    attrs[key[:-11]] = fnc[key][()] != 0
+                elif key == "volume_p":
+                    attrs["volume"] = fnc[key][()]
+                elif key in ["xm", "xn", "xm_nyq", "xn_nyq"]:
+                    attrs[key] = np.array(fnc[key][()], dtype=int)
+                elif key in ["pmass_type", "piota_type", "pcurr_type", "mgrid_file"]:
+                    attrs[key] = fnc[key][()].tobytes().decode("ascii")
+                else:
+                    attrs[key] = fnc[key][()]
+            if "lmns_full" not in fnc.variables:
+                attrs["lmns_full"] = None
+            return VmecWOut(**attrs)
+        raise RuntimeError("Failed to load NetCDF wout file " + str(wout_filename))
 
 
 class Threed1Volumetrics(pydantic.BaseModel):

tests/test_init.py

@@ -86,12 +86,16 @@ def cma_output() -> vmecpp.VmecOutput:
     return out
 
 
-def test_vmecwout_save(cma_output):
+def test_vmecwout_io(cma_output):
     with tempfile.NamedTemporaryFile() as tmp_file:
         cma_output.wout.save(tmp_file.name)
 
         assert Path(tmp_file.name).exists()
 
+        # check that from_wout_file can load the file as well
+        loaded_wout = vmecpp.VmecWOut.from_wout_file(tmp_file.name)
+        assert loaded_wout is not None
+
         test_dataset = netCDF4.Dataset(tmp_file.name, "r")
 
     expected_dataset = netCDF4.Dataset(TEST_DATA_DIR / "wout_cma.nc", "r")