def test_read_input(self):
"""
Try reading a VMEC input file.
"""
ictrl[:] = 0
ictrl[0] = restart_flag + readin_flag
logger.info("Calling runvmec. ictrl={} comm={}".format(ictrl, fcomm))
vmec.runvmec(ictrl, filename, verbose, fcomm, reset_file)
self.assertEqual(ictrl[1], 0)
self.assertEqual(vmec.vmec_input.nfp, 3)
self.assertEqual(vmec.vmec_input.mpol, 4)
self.assertEqual(vmec.vmec_input.ntor, 3)
logger.info('rbc.shape: {}'.format(vmec.vmec_input.rbc.shape))
logger.info('rbc: {}'.format(vmec.vmec_input.rbc[101:103, 0:4]))
# n = 0, m = 0:
self.assertAlmostEqual(vmec.vmec_input.rbc[101, 0], 1.3782)
# n = 0, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[101, 1], 4.6465E-01)
# n = 1, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[102, 1], 1.6516E-01)
logger.info("About to cleanup(False)")
vmec.cleanup(False)
def compare_to_vmec(name, r=0.005, nphi=151):
"""
Check that VMEC can run the input file outputed by pyQSC
and check that the resulting VMEC output file has
the expected parameters
"""
# Add the directory of this file to the specified filename:
inputFile="input."+str(name).replace(" ","")
abs_filename = os.path.join(os.path.dirname(__file__), inputFile)
# Run pyQsc and create a VMEC input file
logger.info('Creating pyQSC configuration')
order = 'r2' if name[1] == '2' else 'r1'
py = Qsc.from_paper(name, nphi=nphi, order=order)
logger.info('Outputing to VMEC')
py.to_vmec(inputFile,r)
# Run VMEC
fcomm = MPI.COMM_WORLD.py2f()
logger.info("Calling runvmec. comm={}".format(fcomm))
ictrl=np.array([15,0,0,0,0], dtype=np.int32)
vmec.runvmec(ictrl, inputFile, True, fcomm, '')
# Check that VMEC converged
assert ictrl[1] == 11
# Open VMEC output file
woutFile="wout_"+str(name).replace(" ","")+".nc"
f = netcdf.netcdf_file(woutFile, 'r')
# Compare the results
logger.info('pyQSC iota on axis = '+str(py.iota))
logger.info('VMEC iota on axis = '+str(-f.variables['iotaf'][()][0]))
logger.info('pyQSC field on axis = '+str(py.B0))
logger.info('VMEC bmnc[1][0] = '+str(f.variables['bmnc'][()][1][0]))
assert np.isclose(py.iota,-f.variables['iotaf'][()][0],rtol=1e-2)
assert np.isclose(py.B0,f.variables['bmnc'][()][1][0],rtol=1e-2)
vmec.cleanup(True)
f.close()
def test_run_read(self):
"""
Try running VMEC, then reading in the output.
"""
ictrl[:] = 0
ictrl[0] = 1 + 2 + 4 + 8
vmec.runvmec(ictrl, filename, verbose, fcomm, reset_file)
self.assertEqual(ictrl[1], 11)
self.assertEqual(vmec.vmec_input.nfp, 3)
self.assertEqual(vmec.vmec_input.mpol, 4)
self.assertEqual(vmec.vmec_input.ntor, 3)
logger.info('rbc.shape: {}'.format(vmec.vmec_input.rbc.shape))
logger.info('rbc: {}'.format(vmec.vmec_input.rbc[101:103, 0:4]))
# n = 0, m = 0:
self.assertAlmostEqual(vmec.vmec_input.rbc[101, 0], 1.3782)
# n = 0, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[101, 1], 4.6465E-01)
# n = 1, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[102, 1], 1.6516E-01)
# Before trying to read the wout file, make sure master is
# done writing it.
MPI.COMM_WORLD.Barrier()
# Now try reading in the output
#wout_file = os.path.join(os.path.dirname(__file__), 'wout_li383_low_res.nc')
wout_file = 'wout_li383_low_res.nc'
logger.info("About to read output file {}".format(wout_file))
#ierr = 0
#vmec.read_wout_mod.read_wout_file(wout_file, ierr)
#self.assertEqual(ierr, 0)
f = netcdf.netcdf_file(wout_file, mmap=False)
#self.assertAlmostEqual(vmec.read_wout_mod.betatot, \
self.assertAlmostEqual(f.variables['betatotal'][()], \
0.0426215030653306, places=4)
#logger.info('iotaf.shape: {}'.format(vmec.read_wout_mod.iotaf.shape))
#logger.info('rmnc.shape: {}'.format(vmec.read_wout_mod.rmnc.shape))
#self.assertAlmostEqual(vmec.read_wout_mod.iotaf[-1], \
iotaf = f.variables['iotaf'][()]
self.assertAlmostEqual(iotaf[-1], \
0.6556508142482989, places=4)
#self.assertAlmostEqual(vmec.read_wout_mod.rmnc[0, 0], \
rmnc = f.variables['rmnc'][()]
self.assertAlmostEqual(rmnc[0, 0], \
1.4760749266902973, places=4)
f.close()
logger.info("About to cleanup(True)")
vmec.cleanup(True)
def tearDown(self):
"""
Tear down the test fixture. This subroutine is automatically
run by python's unittest module before every test. Here we
call runvmec with the cleanup flag to deallocate arrays, such
that runvmec can be called again later.
"""
self.ictrl[0] = 16 # cleanup
vmec.runvmec(self.ictrl, self.filename, self.verbose, \
self.fcomm, reset_file)
def test_run_read(self):
"""
Try running VMEC, then reading in results from the wout file.
"""
self.ictrl[0] = 1 + 2 + 4 + 8
vmec.runvmec(self.ictrl, self.filename, self.verbose, \
self.fcomm, reset_file)
self.assertTrue(self.ictrl[1] in success_codes)
self.assertEqual(vmec.vmec_input.nfp, 3)
self.assertEqual(vmec.vmec_input.mpol, 4)
self.assertEqual(vmec.vmec_input.ntor, 3)
print('rbc.shape:', vmec.vmec_input.rbc.shape)
print('rbc:', vmec.vmec_input.rbc[101:103, 0:4])
# n = 0, m = 0:
self.assertAlmostEqual(vmec.vmec_input.rbc[101, 0], 1.3782)
# n = 0, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[101, 1], 4.6465E-01)
# n = 1, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[102, 1], 1.6516E-01)
# Now try reading in the output
wout_file = os.path.join(os.path.dirname(__file__),
'wout_li383_low_res.nc')
ierr = 0
vmec.read_wout_mod.read_wout_file(wout_file, ierr)
self.assertEqual(ierr, 0)
self.assertAlmostEqual(vmec.read_wout_mod.betatot, \
0.0426215030653306, places=4)
print('iotaf.shape:', vmec.read_wout_mod.iotaf.shape)
print('rmnc.shape:', vmec.read_wout_mod.rmnc.shape)
self.assertAlmostEqual(vmec.read_wout_mod.iotaf[-1], \
0.654868168783638, places=4)
self.assertAlmostEqual(vmec.read_wout_mod.rmnc[0, 0], \
1.4773028173065, places=4)
def test_read_input(self):
"""
Try reading a VMEC input file.
"""
self.ictrl[0] = run_modes['input']
vmec.runvmec(self.ictrl, self.filename, self.verbose, \
self.fcomm, reset_file)
self.assertTrue(self.ictrl[1] in success_codes)
self.assertEqual(vmec.vmec_input.nfp, 3)
self.assertEqual(vmec.vmec_input.mpol, 4)
self.assertEqual(vmec.vmec_input.ntor, 3)
print('rbc.shape:', vmec.vmec_input.rbc.shape)
print('rbc:', vmec.vmec_input.rbc[101:103, 0:4])
# n = 0, m = 0:
self.assertAlmostEqual(vmec.vmec_input.rbc[101, 0], 1.3782)
# n = 0, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[101, 1], 4.6465E-01)
# n = 1, m = 1:
self.assertAlmostEqual(vmec.vmec_input.zbs[102, 1], 1.6516E-01)
def run(self):
"""
Run VMEC, if ``need_to_run_code`` is ``True``.
"""
if not self.need_to_run_code:
logger.info("run() called but no need to re-run VMEC.")
return
logger.info("Preparing to run VMEC.")
# Transfer values from Parameters to VMEC's fortran modules:
vi = vmec.vmec_input # Shorthand
# Convert boundary to RZFourier if needed:
boundary_RZFourier = self.boundary.to_RZFourier()
# VMEC does not allow mpol or ntor above 101:
if vi.mpol > 101:
raise ValueError("VMEC does not allow mpol > 101")
if vi.ntor > 101:
raise ValueError("VMEC does not allow ntor > 101")
vi.rbc[:, :] = 0
vi.zbs[:, :] = 0
mpol_capped = np.min([boundary_RZFourier.mpol, 101])
ntor_capped = np.min([boundary_RZFourier.ntor, 101])
# Transfer boundary shape data from the surface object to VMEC:
for m in range(mpol_capped + 1):
for n in range(-ntor_capped, ntor_capped + 1):
vi.rbc[101 + n, m] = boundary_RZFourier.get_rc(m, n)
vi.zbs[101 + n, m] = boundary_RZFourier.get_zs(m, n)
# Set axis shape to something that is obviously wrong (R=0) to
# trigger vmec's internal guess_axis.f to run. Otherwise the
# initial axis shape for run N will be the final axis shape
# from run N-1, which makes VMEC results depend slightly on
# the history of previous evaluations, confusing the finite
# differencing.
vi.raxis_cc[:] = 0
vi.raxis_cs[:] = 0
vi.zaxis_cc[:] = 0
vi.zaxis_cs[:] = 0
self.iter += 1
input_file = self.input_file + '_{:03d}_{:06d}'.format(
self.mpi.group, self.iter)
self.output_file = os.path.join(
os.getcwd(),
os.path.basename(input_file).replace('input.', 'wout_') + '.nc')
mercier_file = os.path.join(
os.getcwd(),
os.path.basename(input_file).replace('input.', 'mercier.'))
jxbout_file = os.path.join(
os.getcwd(),
os.path.basename(input_file).replace('input.', 'jxbout_') + '.nc')
# I should write an input file here.
logger.info("Calling VMEC reinit().")
vmec.reinit()
logger.info("Calling runvmec().")
self.ictrl[0] = restart_flag + reset_jacdt_flag \
+ timestep_flag + output_flag
self.ictrl[1] = 0 # ierr
self.ictrl[2] = 0 # numsteps
self.ictrl[3] = 0 # ns_index
self.ictrl[4] = 0 # iseq
verbose = True
reset_file = ''
vmec.runvmec(self.ictrl, input_file, verbose, self.fcomm, reset_file)
ierr = self.ictrl[1]
# Deallocate arrays, even if vmec did not converge:
logger.info("Calling VMEC cleanup().")
vmec.cleanup(True)
# See VMEC2000/Sources/General/vmec_params.f for ierr codes.
# 11 = successful_term_flag.
# Error codes that are expected to occur due to lack of
# convergence cause ObjectiveFailure, which the optimizer
# handles gracefully by treating the point as bad. But the
# user/developer should know if an error codes arises that
# should logically never occur, so these codes raise a
# different exception.
if ierr in [0, 5]:
raise RuntimeError(f"runvmec returned an error code that should " \
"never occur: ierr={ierr}")
if ierr != 11:
raise ObjectiveFailure(f"VMEC did not converge. ierr={ierr}")
logger.info("VMEC run complete. Now loading output.")
self.load_wout()
logger.info("Done loading VMEC output.")
# Group leaders handle deletion of files:
if self.mpi.proc0_groups:
# Delete some files produced by VMEC that we never care
# about. For some reason the os.remove statements give a 'file
# not found' error in the CI, hence the try-except blocks.
try:
os.remove(mercier_file)
except FileNotFoundError:
logger.debug(f'Tried to delete the file {mercier_file} but it was not found')
raise
try:
os.remove(jxbout_file)
except FileNotFoundError:
logger.debug(f'Tried to delete the file {jxbout_file} but it was not found')
raise
try:
os.remove("fort.9")
except FileNotFoundError:
logger.debug('Tried to delete the file fort.9 but it was not found')
# If the worker group is not 0, delete all wout files, unless
# keep_all_files is True:
if (not self.keep_all_files) and (self.mpi.group > 0):
os.remove(self.output_file)
# Delete the previous output file, if desired:
for filename in self.files_to_delete:
os.remove(filename)
self.files_to_delete = []
# Record the latest output file to delete if we run again:
if (self.mpi.group == 0) and (self.iter > 0) and (not self.keep_all_files):
self.files_to_delete.append(self.output_file)
self.need_to_run_code = False
def __init__(self,
filename: Union[str, None] = None,
mpi: Union[MpiPartition, None] = None,
keep_all_files: bool = False):
if MPI is None:
raise RuntimeError("mpi4py needs to be installed for running VMEC")
if vmec is None:
raise RuntimeError(
"Running VMEC from simsopt requires VMEC python extension. "
"Install the VMEC python extension from "
"https://https://github.com/hiddenSymmetries/VMEC2000")
if filename is None:
# Read default input file, which should be in the same
# directory as this file:
filename = os.path.join(os.path.dirname(__file__), 'input.default')
logger.info("Initializing a VMEC object from defaults in "
+ filename)
else:
logger.info("Initializing a VMEC object from file: " + filename)
self.input_file = filename
# Get MPI communicator:
if mpi is None:
self.mpi = MpiPartition(ngroups=1)
else:
self.mpi = mpi
comm = self.mpi.comm_groups
self.fcomm = comm.py2f()
self.ictrl = np.zeros(5, dtype=np.int32)
self.iter = -1
self.keep_all_files = keep_all_files
self.files_to_delete = []
self.wout = Struct()
self.ictrl[0] = restart_flag + readin_flag
self.ictrl[1] = 0 # ierr
self.ictrl[2] = 0 # numsteps
self.ictrl[3] = 0 # ns_index
self.ictrl[4] = 0 # iseq
verbose = True
reset_file = ''
logger.info('About to call runvmec to readin')
vmec.runvmec(self.ictrl, filename, verbose, self.fcomm, reset_file)
ierr = self.ictrl[1]
logger.info('Done with runvmec. ierr={}. Calling cleanup next.'.format(ierr))
# Deallocate arrays allocated by VMEC's fixaray():
vmec.cleanup(False)
if ierr != 0:
raise RuntimeError("Failed to initialize VMEC from input file {}. "
"error code {}".format(filename, ierr))
objstr = " for Vmec " + str(hex(id(self)))
# Create an attribute for each VMEC input parameter in VMEC's fortran
# modules,
self.indata = vmec.vmec_input # Shorthand
vi = vmec.vmec_input # Shorthand
# A vmec object has mpol and ntor attributes independent of
# the boundary. The boundary surface object is initialized
# with mpol and ntor values that match those of the vmec
# object, but the mpol/ntor values of either the vmec object
# or the boundary surface object can be changed independently
# by the user.
self.boundary = SurfaceRZFourier(nfp=vi.nfp,
stellsym=not vi.lasym,
mpol=vi.mpol,
ntor=vi.ntor)
self.free_boundary = bool(vi.lfreeb)
# Transfer boundary shape data from fortran to the ParameterArray:
for m in range(vi.mpol + 1):
for n in range(-vi.ntor, vi.ntor + 1):
self.boundary.rc[m, n + vi.ntor] = vi.rbc[101 + n, m]
self.boundary.zs[m, n + vi.ntor] = vi.zbs[101 + n, m]
# Handle a few variables that are not Parameters:
self.depends_on = ["boundary"]
self.need_to_run_code = True
self.fixed = np.full(len(self.get_dofs()), True)
self.names = ['delt', 'tcon0', 'phiedge', 'curtor', 'gamma']
def test_regression(self):
"""
Read in multiple input files, running VMEC various numbers of
times in between. Several output quantities are compared to
known reference values. This test case covers both cases in
which VMEC converges and cases in which it does not, due to an
unachievably small ftol. This test case includes several
geometries, covering both axisymmetry and non-axisymmetry, and
it covers both stellarator-symmetry and
non-stellarator-symmetry.
"""
files = [
'circular_tokamak', 'up_down_asymmetric_tokamak', 'li383_low_res',
'LandremanSenguptaPlunk_section5p3_low_res'
]
#file_indices = list(range(len(files)))
for nrun in range(4):
for jfile in range(len(files)):
filename = os.path.join(os.path.dirname(__file__),
'input.' + files[jfile])
# Read in the input file:
ictrl[:] = 0
ictrl[0] = restart_flag + readin_flag
logger.info("Calling runvmec to read input file {}. ictrl={} comm={}" \
.format(filename, ictrl, fcomm))
vmec.runvmec(ictrl, filename, verbose, fcomm, reset_file)
self.assertEqual(ictrl[1], 0)
ftol_array_orig = np.copy(vmec.vmec_input.ftol_array)
# Deallocate arrays:
logger.info("About to cleanup(False)")
vmec.cleanup(False) # False because we have not time-stepped.
for jrun in range(nrun):
# Set axis to 0 to force VMEC to recompute its
# initial guess for the axis. This way each run of
# the code should produce independent results, no
# matter how many times the code is run.
vmec.vmec_input.raxis_cc = 0
vmec.vmec_input.raxis_cs = 0
vmec.vmec_input.zaxis_cc = 0
vmec.vmec_input.zaxis_cs = 0
# Cover both cases in which VMEC converges and
# cases in which it does not. To do this, for
# every other run we set ftol to a value so small
# that it is unattainable.
should_converge = (np.mod(jrun, 2) == 0)
vmec.vmec_input.ftol_array = ftol_array_orig
if not should_converge:
vmec.vmec_input.ftol_array[:3] = 1.0e-30
logger.info("About to reinit. nrun={} jfile={} jrun={} should_converge={}" \
.format(nrun, jfile, jrun, should_converge))
logger.info("ftol_array_orig: {}".format(
ftol_array_orig[:5]))
logger.info("ftol_array: {}".format(
vmec.vmec_input.ftol_array[:5]))
vmec.reinit()
# Time-step:
ictrl[:] = 0
ictrl[
0] = restart_flag + reset_jacdt_flag + timestep_flag + output_flag
logger.info(
"Calling runvmec to timestep. ictrl={} comm={}".format(
ictrl, fcomm))
vmec.runvmec(ictrl, filename, verbose, fcomm, reset_file)
# Expected return codes for ictrl[1], from vmec_params.f:
# 2 = more_iter_flag
# 11 = successful_term_flag
if should_converge:
self.assertEqual(ictrl[1], 11)
else:
self.assertEqual(ictrl[1], 2)
# Regression tests: compare output to reference files.
if should_converge:
# Before trying to read the wout file, make sure master is
# done writing it.
MPI.COMM_WORLD.Barrier()
# New wout output file is in the current
# working directory, whereas the reference
# wout file is in a fixed directory.
wout_file = 'wout_' + files[jfile] + '.nc'
reference_file = os.path.join(
os.path.dirname(__file__),
'wout_' + files[jfile] + '_reference.nc')
ierr = 0
logger.info(
'About to read output file {}'.format(wout_file))
f1 = netcdf.netcdf_file(wout_file, mmap=False)
#vmec.read_wout_mod.read_wout_file(wout_file, ierr)
#self.assertEqual(ierr, 0)
logger.info(
'About to read reference output file {}'.format(
reference_file))
f2 = netcdf.netcdf_file(reference_file, mmap=False)
compare(f1, f2, 'iotaf')
compare(f1, f2, 'rmnc')
compare(f1, f2, 'zmns')
compare(f1, f2, 'lmns')
compare(f1, f2, 'bmnc')
#np.testing.assert_allclose(vmec.read_wout_mod.iotaf, f2.variables['iotaf'][()])
f1.close()
f2.close()
# Deallocate arrays:
logger.info("About to cleanup(True). nrun={} jfile={} jrun={}" \
.format(nrun, jfile, jrun))
vmec.cleanup(True)
def run(self,
mode='main',
ier=0,
numsteps=-1,
ns_index=-1,
iseq=0,
input_file=None,
verbose=None,
comm=None,
reset_file='',
**kwargs):
"""Interface for Fortran subroutine runvmec in Sources/TimeStep/runvmec.f
Args:
mode (str): The running mode of VMEC. It should be one of the
following options,
('all', 'input', 'output', 'main'). (default: 'main').
ier (int): Flag for error condition. (default: 0).
numsteps (int): Number time steps to evolve the equilibrium.
Iterations will stop EITHER if numsteps > 0 and when the
number of vmec iterations exceeds numsteps; OR if the ftol
condition is satisfied, whichever comes first. The
timestep_flag must be set (in ictrl_flag) for this to be
in effect. If numsteps <= 0, then vmec will choose
consecutive (and increasing) values from the ns_array
until ftol is satisfied on each successive multi-grid.
(default: -1).
ns_index (int): if > 0 on entry, specifies index (in ns_array)
of the radial grid to be used for the present iteration
phase. If ns_index <= 0, vmec will use the previous value
of this index (if the ftol condition was not satisfied
during the last call to runvmec) or the next value of this
index, and it will iterate through each successive
non-zero member of the ns_array until ftol-convergence
occurs on each multigrid.
(default: -1).
iseq (int): specifies a unique sequence label for identifying
output files in a sequential vmec run.
(default: 0).
input_file (str): Filename for VMEC input namelist.
(default: None -> self.input_file).
verbose (bool): If wants scree outputs.
(default: None -> self.verbose).
comm (int): MPI_Communicater, should be converted to Fortran
format using MPI.py2f().
(default: None -> self.comm).
reset_file (str): Filename for reset runs. (default: '').
Returns:
None
"""
# check arguments
mode = mode.lower()
assert mode in run_modes, ("Unsupported running mode. Should "
"be one of [{:}].").format(','.join(
run_modes.keys()))
assert ier == 0, "Error flag should be zero at input."
if input_file is None:
input_file = self.input_file + '_{:06d}'.format(self.iter)
else:
if 'input.' not in input_file:
input_file = 'input.' + input_file
self.output_file = input_file.replace('input.', 'wout_') + '.nc'
if verbose is None:
verbose = self.verbose
if comm is None:
comm = self.comm
# prepare Fortran arguments
self.ictrl[0] = run_modes[mode]
self.ictrl[1] = ier
self.ictrl[2] = numsteps
self.ictrl[3] = ns_index
self.ictrl[4] = iseq
# run VMEC
# print("before: ", self.ictrl)
vmec.runvmec(self.ictrl, input_file, verbose, comm, reset_file)
self.iter += 1
self.success = self.ictrl[1] in self.success_code
return self.success