Skip to content

Commit 99a2d11

Browse files
committed
Reordering and renaming (compatibility issues of WOut)
1 parent b9683e1 commit 99a2d11

File tree

1 file changed

+33
-31
lines changed

1 file changed

+33
-31
lines changed

src/vmecpp/__init__.py

Lines changed: 33 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -604,14 +604,14 @@ def save(self, out_path: str | Path) -> None:
604604
with netCDF4.Dataset(out_path, "w", format="NETCDF3_CLASSIC") as fnc:
605605
# scalar ints
606606
for varname in [
607-
"ier_flag",
608-
"niter",
609607
"nfp",
610608
"ns",
611609
"mpol",
612610
"ntor",
613611
"mnmax",
614612
"mnmax_nyq",
613+
"niter",
614+
"ier_flag",
615615
"signgs",
616616
"nextcur",
617617
]:
@@ -626,6 +626,7 @@ def save(self, out_path: str | Path) -> None:
626626
for varname in [
627627
"wb",
628628
"wp",
629+
"gamma",
629630
"rmax_surf",
630631
"rmin_surf",
631632
"zmax_surf",
@@ -643,25 +644,40 @@ def save(self, out_path: str | Path) -> None:
643644
"Aminor_p",
644645
"Rmajor_p",
645646
"volume_p",
647+
"ftolv",
648+
"fsql",
646649
"fsqr",
647650
"fsqz",
648-
"fsql",
649-
"ftolv",
650-
"gamma",
651651
]:
652652
fnc.createVariable(varname, np.float64)
653653
fnc[varname][:] = getattr(self, varname)
654654

655-
# create dimensions
655+
# create dimensions (in the same order as VMEC2000)
656+
# For the input extension
657+
fnc.createDimension("dim_00100", 100)
658+
# dim_00200 = mgrid_file_max_string_length
659+
fnc.createDimension("dim_00200", 200)
660+
# dim_00020 = profile_strings_max_len
661+
fnc.createDimension("dim_00020", 20)
662+
if self.lfreeb:
663+
assert self.nextcur > 0
664+
fnc.createDimension("ext_current", self.nextcur)
665+
# and a single character for mgrid_mode
666+
fnc.createDimension("dim_00001", 1)
656667
fnc.createDimension("mn_mode", self.mnmax)
657-
fnc.createDimension("radius", self.ns)
658-
fnc.createDimension("n_tor", self.ntor + 1) # Fortran quirk
659668
fnc.createDimension("mn_mode_nyq", self.mnmax_nyq)
669+
fnc.createDimension("n_tor", self.ntor + 1) # Fortran quirk
660670
fnc.createDimension("preset", _PRESET_DIM)
661671
fnc.createDimension("ndfmax", _NDF_MAX_DIM)
662-
if self.lfreeb:
663-
assert self.nextcur > 0
664-
fnc.createDimension("ext_current", self.nextcur)
672+
fnc.createDimension("radius", self.ns)
673+
674+
# Dimensions that are not in use yet
675+
_TIME = 100
676+
_MN_MAX_POT = 100 # TODO(jurasic) self.mnmaxpot
677+
fnc.createDimension("time", _TIME)
678+
fnc.createDimension("mn_mode_pot", _MN_MAX_POT)
679+
fnc.createDimension("current_label", 30)
680+
fnc.createDimension("dim_00006", 6)
665681

666682
# radial profiles
667683
for varname in [
@@ -727,8 +743,8 @@ def save(self, out_path: str | Path) -> None:
727743
fnc[varname][:] = getattr(self, varname)[:]
728744

729745
for varname in [
730-
"bmnc",
731746
"gmnc",
747+
"bmnc",
732748
"bsubumnc",
733749
"bsubvmnc",
734750
"bsubsmns",
@@ -763,12 +779,10 @@ def save(self, out_path: str | Path) -> None:
763779
fnc["version_"][:] = self.version_
764780

765781
# strings
766-
# maximum length of the string, copied from wout_cma.nc
767-
max_string_length = 20
768-
769782
def create_string_variable(
770-
varname: str, max_string_length: int, dimension_name: str
783+
varname: str, dimension_name: str
771784
) -> netCDF4.Variable:
785+
max_string_length = fnc.dimensions[dimension_name].size
772786
string_variable = fnc.createVariable(varname, "S1", (dimension_name,))
773787

774788
# Put the string in the format netCDF3 requires. Don't know what to say.
@@ -782,23 +796,11 @@ def create_string_variable(
782796
string_variable[:] = padded_value_as_netcdf3_compatible_chararray
783797
return string_variable
784798

785-
fnc.createDimension("profile_strings_max_len", max_string_length)
786799
for varname in ["pcurr_type", "pmass_type", "piota_type"]:
787-
create_string_variable(
788-
varname, max_string_length, "profile_strings_max_len"
789-
)
800+
create_string_variable(varname, "dim_00020")
790801

791-
# now mgrid_file
792-
max_string_length = 200 # value copied from wout_cma.nc
793-
fnc.createDimension("mgrid_file_max_string_length", max_string_length)
794-
create_string_variable(
795-
"mgrid_file", max_string_length, "mgrid_file_max_string_length"
796-
)
797-
798-
# and a single character for mgrid_mode
799-
max_string_length = 1 # value copied from wout_cma.nc
800-
fnc.createDimension("mgrid_mode_length", max_string_length)
801-
create_string_variable("mgrid_mode", max_string_length, "mgrid_mode_length")
802+
create_string_variable("mgrid_file", "dim_00200")
803+
create_string_variable("mgrid_mode", "dim_00001")
802804

803805
@staticmethod
804806
def _from_cpp_wout(cpp_wout: _vmecpp.VmecppWOut) -> VmecWOut:

0 commit comments

Comments
 (0)