Skip to content

Commit

Permalink
Add a method to load a VmecWout object from a NetCDF file.
Browse files Browse the repository at this point in the history
  • Loading branch information
jons-pf committed Feb 14, 2025
1 parent 97142b3 commit ff24e32
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 2 deletions.
26 changes: 25 additions & 1 deletion src/vmecpp/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
6 changes: 5 additions & 1 deletion tests/test_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ 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)

Expand Down Expand Up @@ -121,6 +121,10 @@ def test_vmecwout_save(cma_output):
test_value[:], expected_value[:], err_msg=error_msg, rtol=1e-6, atol=1e-7
)

# 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


def test_jxbout_bindings(cma_output):
for varname in [
Expand Down

0 comments on commit ff24e32

Please sign in to comment.