def step(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'vmec: step')
flags = self.zip_ref.get_state()
if 'state' in flags and flags['state'] == 'needs_update' or 'force_update' in keywords:
task_wait = self.services.launch_task(self.NPROC,
self.services.get_working_dir(),
self.VMEC_EXE,
self.current_vmec_namelist,
logfile = 'vmec.log')
# Update flags.
self.zip_ref.set_state(state='updated')
# Wait for VMEC to finish.
if (self.services.wait_task(task_wait) and not os.path.exists(self.current_wout_file)):
self.services.error('vmec: step failed.')
# Add the wout file to the state.
self.zip_ref.write([self.current_vmec_namelist, self.current_wout_file])
else:
# Update flags.
self.zip_ref.set_state(state='unchanged')
self.zip_ref.close()
self.services.update_state()
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'siesta: init')
self.services.stage_state()
self.current_siesta_namelist = self.services.get_config_param(
'SIESTA_NAMELIST_INPUT')
self.current_siesta_state = self.services.get_config_param(
'CURRENT_SIESTA_STATE')
current_vmec_state = self.services.get_config_param(
'CURRENT_VMEC_STATE')
current_vmec_namelist = self.services.get_config_param(
'VMEC_NAMELIST_INPUT')
current_wout_file = 'wout_{}.nc'.format(
current_vmec_namelist.replace('input.', '', 1))
# Stage state.
self.services.stage_state()
# Unzip files from the state. Use mode a so files can be read and written to.
self.zip_ref = ZipState.ZipState(self.current_siesta_state, 'a')
self.zip_ref.extract(self.current_siesta_namelist)
self.zip_ref.extract(current_vmec_state)
with ZipState.ZipState(current_vmec_state, 'r') as vmec_zip_ref:
vmec_zip_ref.extract(current_wout_file)
flags = vmec_zip_ref.get_state()
if 'state' in flags and flags['state'] == 'updated':
self.zip_ref.set_state(state='needs_update')
# Update parameters in the namelist.
self.set_namelist(wout_file=current_wout_file, **keywords)
def eval_jacobian(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose',
'quasi_newton_driver: eval_jacobian')
# Set the model to a known state.
shutil.copy2(self.current_model_state, 'model_inputs')
for worker in self.model_workers:
worker['wait'] = self.services.call_nonblocking(
worker['init'], 'init', timeStamp)
# Perturb the parameters.
for worker in self.model_workers:
keywords = {worker['name']: worker['value'] + worker['vrnc']}
self.services.wait_call(worker['wait'], True)
worker['wait'] = self.services.call_nonblocking(
worker['driver'], 'init', timeStamp, **keywords)
# Recompute the model.
for worker in self.model_workers:
self.services.wait_call(worker['wait'], True)
worker['wait'] = self.services.call_nonblocking(
worker['driver'],
'step',
timeStamp,
result_file=worker['result'])
# Collect the results.
for worker in self.model_workers:
self.services.wait_call(worker['wait'], True)
worker['wait'] = self.services.call_nonblocking(
worker['init'], 'init', timeStamp)
self.services.stage_subflow_output_files()
# Compute the normalized jacobian A.
#
# A_ij = d e_i/d a_j (1)
#
# Where e is the error vector.
#
# e_i = W_i*((S_i - M_i)/sigma_i)^2 (2)
#
# Note due to the what the memory is laid out the Jacobian is transposed.
for i, worker in enumerate(self.model_workers):
with ZipState.ZipState(worker['output'], 'a') as zip_ref:
zip_ref.extract(worker['result'])
with open(worker['result'], 'r') as result_file:
self.jacobian[i] = self.e - self.get_e(result_file)
with open('jacobian.log', 'a') as jacobian_ref:
jacobian_ref.write('Jacobian step {}\n'.format(timeStamp))
for j in range(len(self.e)):
self.jacobian[:, j].tofile(jacobian_ref,
sep=',',
format='%12.5e')
jacobian_ref.write('\n')
jacobian_ref.write('\n')
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose',
'massive_serial_runner_driver: init')
# Get config filenames.
self.current_state = self.services.get_config_param(
'CURRENT_MSR_STATE')
# Initialize the massive serial runner.
self.massive_serial_runner_port = self.services.get_port('MSR')
self.wait = self.services.call_nonblocking(
self.massive_serial_runner_port, 'init', timeStamp, **keywords)
def finalize(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'siesta_driver: finalize')
self.wait = [
self.services.call_nonblocking(self.vmec_worker['init'],
'finalize', timeStamp),
self.services.call_nonblocking(self.vmec_worker['driver'],
'finalize', timeStamp),
self.services.call_nonblocking(self.siesta_port, 'finalize',
timeStamp)
]
self.services.wait_call_list(self.wait, True)
def finalize(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'cariddi_driver: finalize')
wait_call = [
self.services.call_nonblocking(self.eq_worker['init'], 'finalize',
timeStamp),
self.services.call_nonblocking(self.eq_worker['driver'],
'finalize', timeStamp),
self.services.call_nonblocking(self.cariddi_port, 'finalize',
timeStamp)
]
self.services.wait_call_list(wait_call, True)
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'vmec_driver: init')
# Separate out the vmec keywords.
vmec_keywords = {}
for key, value in keywords.items():
if 'vmec__' in key:
vmec_keywords[key.replace('vmec__', '', 1)] = value
# Initialize vmec.
self.vmec_port = self.services.get_port('VMEC')
self.wait = self.services.call_nonblocking(self.vmec_port, 'init',
timeStamp, **vmec_keywords)
def step(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'massive_serial_runner: step')
flags = self.zip_ref.get_state()
# Run the massive serial workflow.
if 'state' in flags and flags['state'] == 'needs_update':
self.services.call(self.massive_serial_worker['init'], 'init', timeStamp)
self.services.call(self.massive_serial_worker['driver'], 'init', timeStamp)
wait = self.services.call_nonblocking(self.massive_serial_worker['driver'], 'step', timeStamp)
database = 'db_{}.dat'.format(timeStamp)
# Collect results of the workflow into the database file.
if self.services.wait_call(wait, True):
self.services.error('massive_serial_runner: step failed to run massive serial')
task_wait = self.services.launch_task(1, self.services.get_working_dir(),
self.MAKE_DATABASE_EXE,
'--rdir=massive_serial_runner_output_dir',
'--input={}'.format(self.database_config),
'--output={}'.format(database),
'--ndir=0', # FIXME: This command option works around a bug in makedb which shouldn't get called.
logfile='make_db_{}.log'.format(timeStamp))
if self.services.wait_task(task_wait):
self.services.error('massive_serial_runner: step failed to make database')
# Save the new database entries.
self.services.stage_output_files(timeStamp, database)
# Convert the database file to json format.
task_wait = self.services.launch_task(1, self.services.get_working_dir(),
self.TO_JSON_EXE,
'--input_file={}'.format(database),
'--output_file={}'.format(self.current_batch),
'--module_path={}'.format(self.constraint_path),
'--module={}'.format(self.constraint_name),
logfile='to_json_{}.log'.format(timeStamp))
if self.services.wait_task(task_wait):
self.services.error('massive_serial_runner: step failed to make json')
self.zip_ref.write(self.current_batch)
self.zip_ref.set_state(state='updated')
else:
self.zip_ref.write(state='unchanged')
self.zip_ref.close()
self.services.update_state()
def init(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'ml_train_init: init')
# Get config filenames.
if timeStamp == 0.0:
self.current_ml_train_state = self.services.get_config_param('CURRENT_ML_TRAIN_STATE')
self.data_gen_config = self.services.get_config_param('DATA_GEN_CONFIG')
self.data_gen_state = self.services.get_config_param('DATA_GEN_STATE')
self.training_data = self.services.get_config_param('TRAINING_DATA')
self.new_data = self.services.get_config_param('NEW_DATA')
self.prediction_data = self.services.get_config_param('PREDICTION_DATA')
self.nn_model_config = self.services.get_config_param('NN_MODEL_CONFIG')
self.nn_model_matrix = self.services.get_config_param('NN_MODEL_MATRIX')
self.nn_model = self.services.get_config_param('NN_MODEL')
self.ml_train_args = self.services.get_config_param('ML_TRAIN_ARGS')
# Remove old inputs. Stage input files.
for file in os.listdir('.'):
os.remove(file)
# State input files and setup the inital state.
self.services.stage_input_files(self.INPUT_FILES)
# Create plasma state from files. Input files can either be a new plasma state,
# training data file or both. If both file were staged, replace the training
# data input file. If the training data file is present flag the plasma state
# as needing to be updated.
with ZipState.ZipState(self.current_ml_train_state, 'a') as zip_ref:
zip_ref.write_or_check(self.data_gen_config)
zip_ref.write_or_check(self.data_gen_state)
zip_ref.write_optional(self.training_data)
zip_ref.write_optional(self.new_data)
zip_ref.write_optional(self.prediction_data)
zip_ref.write_or_check(self.nn_model_config)
zip_ref.write_optional(self.nn_model_matrix)
if os.path.exists(self.nn_model):
with ZipState.ZipState('{}.zip'.format(self.nn_model), 'w') as nn_ref:
nn_ref.write(self.nn_model)
zip_ref.write_optional('{}.zip'.format(self.nn_model))
zip_ref.write_or_check(self.ml_train_args)
zip_ref.set_state(state='needs_update')
self.services.update_state()
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'vmec: init')
self.current_vmec_namelist = self.services.get_config_param('VMEC_NAMELIST_INPUT')
self.current_wout_file = 'wout_{}.nc'.format(self.current_vmec_namelist.replace('input.','',1))
current_vmec_state = self.services.get_config_param('CURRENT_VMEC_STATE')
# Stage state.
self.services.stage_state()
# Unzip files from the state. Use mode a so files can be read and written to.
self.zip_ref = ZipState.ZipState(current_vmec_state, 'a')
self.zip_ref.extract(self.current_vmec_namelist)
if len(keywords) > 0:
self.zip_ref.set_state(state='needs_update')
# Update parameters in the namelist.
namelist = OMFITnamelist(self.current_vmec_namelist,
collect_arrays={
'ns_array' : {'default' : 0, 'shape' : (100,), 'offset' : (1,), 'sparray' : True},
'niter_array' : {'default' : 0, 'shape' : (100,), 'offset' : (1,), 'sparray' : True},
'rbs' : {'default' : 0, 'shape' : (203,101), 'offset' : (-101,0), 'sparray' : True},
'rbc' : {'default' : 0, 'shape' : (203,101), 'offset' : (-101,0), 'sparray' : True},
'zbs' : {'default' : 0, 'shape' : (203,101), 'offset' : (-101,0), 'sparray' : True},
'zbc' : {'default' : 0, 'shape' : (203,101), 'offset' : (-101,0), 'sparray' : True},
'am' : {'default' : 0, 'shape' : (21,), 'offset' : (0,), 'sparray' : True},
'ai' : {'default' : 0, 'shape' : (21,), 'offset' : (0,), 'sparray' : True},
'ac' : {'default' : 0, 'shape' : (21,), 'offset' : (0,), 'sparray' : True},
'am_aux_s' : {'default' : 0, 'shape' : (10001,), 'offset' : (1,), 'sparray' : True},
'am_aux_f' : {'default' : 0, 'shape' : (10001,), 'offset' : (1,), 'sparray' : True},
'ai_aux_s' : {'default' : 0, 'shape' : (10001,), 'offset' : (1,), 'sparray' : True},
'ai_aux_f' : {'default' : 0, 'shape' : (10001,), 'offset' : (1,), 'sparray' : True},
'ac_aux_s' : {'default' : 0, 'shape' : (10001,), 'offset' : (1,), 'sparray' : True},
'ac_aux_f' : {'default' : 0, 'shape' : (10001,), 'offset' : (1,), 'sparray' : True},
'raxis' : {'default' : 0, 'shape' : (102,), 'offset' : (0,), 'sparray' : True},
'zaxis' : {'default' : 0, 'shape' : (102,), 'offset' : (0,), 'sparray' : True},
'raxis_cc' : {'default' : 0, 'shape' : (102,), 'offset' : (0,), 'sparray' : True},
'raxis_cs' : {'default' : 0, 'shape' : (102,), 'offset' : (0,), 'sparray' : True},
'zaxis_cc' : {'default' : 0, 'shape' : (102,), 'offset' : (0,), 'sparray' : True},
'zaxis_cs' : {'default' : 0, 'shape' : (102,), 'offset' : (0,), 'sparray' : True},
'ftol_array' : {'default' : 0, 'shape' : (100,), 'offset' : (1,), 'sparray' : True},
'extcur' : {'default' : 0, 'shape' : (300,), 'offset' : (1,), 'sparray' : True}
})
for key, value in keywords.items():
NamelistItem.set(namelist['indata'], key, value)
namelist.save()
def step(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'ml_train_driver: step')
with ZipState.ZipState(self.current_ml_train_state, 'r') as zip_ref:
flags = zip_ref.get_state()
# Adaptive training loop. Hard code the number of iterations for now.
for i in range(self.max_iterations):
# Train the NN model.
self.services.call(self.ml_train_port, 'init', timeStamp)
self.services.call(self.ml_train_port, 'step', timeStamp)
# Get the new data batch and add it to the gen data state.
self.services.stage_state()
with ZipState.ZipState(self.current_ml_train_state, 'a') as zip_ref:
zip_ref.extract_or_check(self.new_data)
zip_ref.extract_or_check(self.training_data)
with ZipState.ZipState(self.data_gen_state, 'a') as model_state_ref:
model_state_ref.write(self.new_data)
# Generate new training data. Update the time the first data batch was
# generated in the init method.
timeStamp = timeStamp + 1.0
shutil.copy2(self.data_gen_state, 'data_gen_input_dir')
self.services.call(self.data_gen['init'], 'init', timeStamp)
self.services.call(self.data_gen['driver'], 'init', timeStamp)
self.services.call(self.data_gen['driver'], 'step', timeStamp)
self.services.stage_subflow_output_files()
# Append the new data to he training data.
with ZipState.ZipState(self.data_gen_state, 'r') as model_state_ref:
model_state_ref.extract_or_check(self.new_data)
if os.path.exists(self.training_data):
self.append_data()
else:
os.rename(self.new_data, self.training_data)
zip_ref.write(self.training_data)
zip_ref.set_state(state='needs_update')
flags = zip_ref.get_state()
self.services.update_state()
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'ml_train: init')
if timeStamp == 0.0:
self.current_ml_train_state = self.services.get_config_param(
'CURRENT_ML_TRAIN_STATE')
self.training_data = self.services.get_config_param(
'TRAINING_DATA')
self.new_data = self.services.get_config_param('NEW_DATA')
self.prediction_data = self.services.get_config_param(
'PREDICTION_DATA')
self.nn_model_config = self.services.get_config_param(
'NN_MODEL_CONFIG')
self.nn_model_matrix = self.services.get_config_param(
'NN_MODEL_MATRIX')
self.nn_model = self.services.get_config_param('NN_MODEL')
self.batch_size = self.services.get_config_param('BATCH_SIZE')
self.constraint_path = self.services.get_config_param(
'MODULE_PATH')
self.constraint_name = self.services.get_config_param(
'MODULE_NAME')
# Stage state.
self.services.stage_state()
# Unzip files from the current state. Use mode a so files can be read and
# written to.
self.zip_ref = ZipState.ZipState(self.current_ml_train_state, 'a')
if timeStamp == 0.0:
ml_train_args = self.services.get_config_param('ML_TRAIN_ARGS')
self.zip_ref.extract_or_check(ml_train_args)
self.zip_ref.extract_or_check(self.nn_model_config)
with open(ml_train_args, 'r') as args_ref:
self.task_args = json.load(args_ref)
self.zip_ref.extract_or_check(self.training_data)
self.zip_ref.extract_optional('{}.zip'.format(self.nn_model))
if os.path.exists('{}.zip'.format(self.nn_model)):
with ZipState.ZipState('{}.zip'.format(self.nn_model),
'r') as nn_ref:
nn_ref.extractall()
def init(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'cariddi: init')
# Stage state.
self.services.stage_state()
# Get config filenames.
if timeStamp == '0.0':
self.current_cariddi_matrix = self.services.get_config_param('CARIDDI_MATRIX_FILE')
self.current_cariddi_geometry = self.services.get_config_param('CARIDDI_GEOMETRY_FILE')
self.cariddi_matrix_path = self.services.get_config_param('CARIDDI_MATRIX_PATH')
self.current_cariddi_state = self.services.get_config_param('CURRENT_CARIDDI_STATE')
self.current_v3fit_state = self.services.get_config_param('CURRENT_V3FIT_STATE')
self.current_siesta_state = self.services.get_config_param('CURRENT_SIESTA_STATE')
self.current_vmec_state = self.services.get_config_param('CURRENT_VMEC_STATE')
current_vmec_namelist = self.services.get_config_param('VMEC_NAMELIST_INPUT')
self.current_wout_file = 'wout_{}.nc'.format(current_vmec_namelist.replace('input.','',1))
self.mgrid_file = self.services.get_config_param('MGRID_FILE')
self.current_vmec_profile = self.services.get_config_param('CURRENT_VMEC_PROFILE')
self.zip_ref = ZipState.ZipState(self.current_cariddi_state, 'a')
if 'A1.nc' in self.zip_ref:
self.zip_ref.extract('A1.nc')
else:
self.zip_ref = ZipState.ZipState(self.current_cariddi_state, 'a')
# Extract input files.
self.flags = self.zip_ref.get_state()
if 'state' in self.flags and self.flags['state'] != 'unchanged':
self.zip_ref.extract(self.current_v3fit_state)
with ZipState.ZipState(self.current_v3fit_state, 'r') as v3fit_zip_ref:
if self.current_siesta_state in v3fit_zip_ref:
with ZipState.ZipState(self.current_siesta_state, 'r') as siesta_zip_ref:
siesta_zip_ref.extract(self.current_vmec_state)
else:
v3fit_zip_ref.extract(self.current_vmec_state)
with ZipState.ZipState(self.current_vmec_state, 'r') as vmec_zip_ref:
if self.current_wout_file in vmec_zip_ref:
vmec_zip_ref.extract(self.current_wout_file)
def lm_step(self, step_size):
ScreenWriter.screen_output(self, 'verbose',
'quasi_newton_driver: lm_step')
if step_size > 0.0 and self.delta_a_len[self.k_use] > step_size:
ut_dot_e = numpy.matmul(self.e, self.j_svd_u)
# Find the L-M parameter lambda that corresponds to a step length of step_size.
_lambda = self.lm_get_lambda(step_size, ut_dot_e)
# Find the step.
return numpy.matmul(
ut_dot_e[0:self.k_use] * self.j_svd_w[0:self.k_use] /
(self.j_svd_w[0:self.k_use]**2.0 + _lambda),
self.j_svd_vt[0:self.k_use])
else:
return self.delta_a[self.k_use]
def step(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'v3fit_driver: step')
# Run V3FIT.
self.services.wait_call(self.wait, True)
self.services.call(self.v3fit_port, 'step', timeStamp, **keywords)
# Prepare the output files for a super work flow. Need to remove any old output
# files first before staging the state.
if os.path.exists(self.OUTPUT_FILES):
os.remove(self.OUTPUT_FILES)
self.services.stage_state()
# The super flow may need to rename the output file. Check if the current state
# matches the output file. If it does not rename the state so it can be staged.
if not os.path.exists(self.OUTPUT_FILES):
os.rename(self.current_v3fit_state, self.OUTPUT_FILES)
def get_exp_dg2(self, delta):
ScreenWriter.screen_output(self, 'verbose',
'quasi_newton_driver: get_exp_dg2')
# Linear part.
#
# 2e * A * da (1)
exp_dg2_lin = 2.0 * numpy.dot(
self.e, numpy.matmul(numpy.transpose(self.jacobian), delta))
# Qaudratic part.
#
# da * A^T * A * da = da * alpha * da (2)
#
# Alpha is the hessian matrix. Equation 14 in Hanson et. al.
# doi: 10.1088/0029-5515/49/7/075031
exp_dg2_quad = numpy.dot(delta, numpy.matmul(self.hessian, delta))
return exp_dg2_lin - exp_dg2_quad
def step(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'siesta: step')
flags = self.zip_ref.get_state()
if 'state' in flags and flags['state'] == 'needs_update':
self.set_namelist(ladd_pert=True, lrestart=False)
self.task_wait = self.services.launch_task(
self.NPROC,
self.services.get_working_dir(),
self.SIESTA_EXE,
logfile='siesta1.log')
# Update flags.
self.zip_ref.set_state(state='updated')
# Restart siesta to ensure convergence.
self.services.wait_task(self.task_wait)
self.set_namelist(ladd_pert=False, lrestart=True)
self.task_wait = self.services.launch_task(
self.NPROC,
self.services.get_working_dir(),
self.SIESTA_EXE,
logfile='siesta2.log')
# Wait for SIESTA to finish.
if (self.services.wait_task(self.task_wait)
or not os.path.exists(self.restart_file)):
self.services.error('siesta: step failed.')
# Add the restart file to the state.
self.zip_ref.write(
[self.current_siesta_namelist, self.restart_file])
else:
# Update flags.
self.zip_ref.set_state(state='unchanged')
self.zip_ref.close()
self.services.update_state()
def step(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'v3fit: step')
flags = self.zip_ref.get_state()
if 'state' in flags and flags['state'] == 'needs_update' or 'force_update' in keywords:
task_wait = self.services.launch_task(self.NPROC,
self.services.get_working_dir(),
self.V3FIT_EXE,
self.current_v3fit_namelist,
logfile = 'v3fit_{}.log'.format(timeStamp))
# Update flags.
self.zip_ref.set_state(state='updated')
# Wait for V3FIT to finish.
if (self.services.wait_task(task_wait) and not os.path.exists(self.result_file)):
self.services.error('v3fit: step failed.')
if 'force_update' in keywords:
with ZipState.ZipState(self.current_vmec_state, 'a') as vmec_zip_ref:
vmec_zip_ref.write(self.current_wout_file)
self.zip_ref.write(self.current_vmec_state)
# Add the result file to the state.
self.zip_ref.write([self.current_v3fit_namelist, self.result_file])
else:
# Update flags.
self.zip_ref.set_state(state='unchanged')
if 'result_file' in keywords:
result_nc = OMFITnc(self.result_file)
nsteps = result_nc['nsteps']['data']
result = {'signal_model': result_nc['signal_model_value']['data'][nsteps,:,0].tolist()}
with open(keywords['result_file'], 'w') as result_ref:
json.dump(result, result_ref)
self.zip_ref.write(keywords['result_file'])
self.zip_ref.close()
self.services.update_state()
def init(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'siesta_init: init')
# Get config filenames.
current_vmec_namelist = self.services.get_config_param(
'VMEC_NAMELIST_INPUT')
current_vmec_state = self.services.get_config_param(
'CURRENT_VMEC_STATE')
current_siesta_namelist = self.services.get_config_param(
'SIESTA_NAMELIST_INPUT')
current_siesta_state = self.services.get_config_param(
'CURRENT_SIESTA_STATE')
# Remove old inputs. Stage input files.
for file in os.listdir('.'):
os.remove(file)
self.services.stage_input_files(self.INPUT_FILES)
# Create a vmec state. If the vmec namelist file exists add the namelist input
# file.
with ZipState.ZipState(current_vmec_state, 'a') as zip_ref:
if os.path.exists(current_vmec_namelist):
zip_ref.write(current_vmec_namelist)
zip_ref.set_state(state='needs_update')
# Create state from files. Input files can either be a new state, namelist
# input file or both. If both files were staged, replace the namelist input
# file. If the namelist file is present flag the state as needing to be
# updated. Sub states will automatically merge.
with ZipState.ZipState(current_siesta_state, 'a') as zip_ref:
if os.path.exists(current_siesta_namelist):
zip_ref.write(current_siesta_namelist)
zip_ref.set_state(state='needs_update')
# The vmec state will be merged with any existing vmec state in the siesta
# state.
zip_ref.write(current_vmec_state)
self.services.update_state()
def lm_get_lambda(self, step_size, ut_dot_e):
ScreenWriter.screen_output(self, 'verbose',
'quasi_newton_driver: lm_get_lambda')
# Define a default value incase the root find fails. Note lambda is a python
# builtin add an underscore to avoid this.
_lambda = (self.j_svd_w[0] * self.delta_a_len[self.k_use] /
step_size)**2.0
f_sqrd = ut_dot_e[0:self.k_use]**2.0
step_size_sqrd = step_size**2.0
# Define a lambda function for the root finder.
f = lambda x: numpy.sum(f_sqrd * (self.j_svd_w[0:self.k_use] / (
self.j_svd_w[0:self.k_use]**2.0 + x))**2.0) - step_size_sqrd
# Find the bracketing values of lambda. Look for a small value of lambda, to
# give a positive function value. f(0) should be greater than zero since the
# step size at k_use is larger than step_size.
f_a = f(0)
if f_a < 0.0:
return _lambda
lambda_b = self.j_svd_w[0]**2.0
f_b = f(lambda_b)
for i in range(0, 20):
if f_b <= 0.0:
break
lambda_b = 4.0 * lambda_b
f_b = f(lambda_b)
if f_b > 0.0:
return _lambda
if f_a * f_b > 0.0:
return _lambda
# Found to intervals that bracket the roots. Now find the root.
return optimize.brentq(f, 0.0, lambda_b)
def init(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'quasi_newton_init: init')
# Get config filenames.
current_model_state = self.services.get_config_param('MODEL_INPUT')
quasi_newton_config = self.services.get_config_param(
'QUASI_NEWTON_CONFIG')
current_quasi_newton_state = self.services.get_config_param(
'CURRENT_QUASI_NEWTON_STATE')
# Remove old inputs. Stage input files.
for file in os.listdir('.'):
os.remove(file)
self.services.stage_input_files(self.INPUT_FILES)
# Create state from files.
with ZipState.ZipState(current_quasi_newton_state, 'a') as zip_ref:
zip_ref.write(quasi_newton_config)
zip_ref.write(current_model_state)
self.services.update_state()
def step(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'quasi_newton_init: step')
def finalize(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'cariddi_init: finalize')
def step(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'cariddi_init: step')
def finalize(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'vmec: finalize')
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'v3fit: init')
self.services.stage_state()
# Get config filenames.
self.current_v3fit_namelist = self.services.get_config_param('V3FIT_NAMELIST_INPUT')
self.current_v3fit_state = self.services.get_config_param('CURRENT_V3FIT_STATE')
self.result_file = 'result.{}.nc'.format(self.current_v3fit_namelist)
current_siesta_namelist = self.services.get_config_param('SIESTA_NAMELIST_INPUT')
current_siesta_state = self.services.get_config_param('CURRENT_SIESTA_STATE')
current_vmec_namelist = self.services.get_config_param('VMEC_NAMELIST_INPUT')
self.current_vmec_state = self.services.get_config_param('CURRENT_VMEC_STATE')
self.current_wout_file = 'wout_{}.nc'.format(current_vmec_namelist.replace('input.','',1))
# Stage state.
self.services.stage_state()
# Unzip files from the state. Use mode a so files can be read and written to.
self.zip_ref = ZipState.ZipState(self.current_v3fit_state, 'a')
self.zip_ref.extract(self.current_v3fit_namelist)
if self.result_file in self.zip_ref:
self.zip_ref.extract(self.result_file)
if current_siesta_state in self.zip_ref:
self.zip_ref.extract(current_siesta_state)
with ZipState.ZipState(current_siesta_state, 'r') as siesta_zip_ref:
siesta_zip_ref.extract(current_siesta_namelist)
namelist = OMFITnamelist(current_siesta_namelist)
current_restart_file = 'siesta_{}.nc'.format(namelist['siesta_info']['restart_ext'])
siesta_zip_ref.extract(current_restart_file)
flags = siesta_zip_ref.get_state()
if 'state' in flags and flags['state'] == 'updated':
self.zip_ref.set_state(state='needs_update')
siesta_zip_ref.extract(self.current_vmec_state)
with ZipState.ZipState(self.current_vmec_state, 'r') as vmec_zip_ref:
vmec_zip_ref.extract(self.current_wout_file)
flags = vmec_zip_ref.get_state()
if 'state' in flags and flags['state'] == 'updated':
self.zip_ref.set_state(state='needs_update')
keywords['siesta_nli_filename'] = current_siesta_namelist
keywords['siesta_restart_filename'] = current_restart_file
keywords['vmec_nli_filename'] = current_vmec_namelist
keywords['vmec_wout_input'] = self.current_wout_file
keywords['model_eq_type'] = 'siesta'
else:
self.zip_ref.extract(self.current_vmec_state)
with ZipState.ZipState(self.current_vmec_state, 'r') as vmec_zip_ref:
vmec_zip_ref.extract(self.current_wout_file)
vmec_zip_ref.extract(current_vmec_namelist)
flags = vmec_zip_ref.get_state()
if 'state' in flags and flags['state'] == 'updated':
self.zip_ref.set_state(state='needs_update')
keywords['vmec_nli_filename'] = current_vmec_namelist
keywords['vmec_wout_input'] = self.current_wout_file
keywords['model_eq_type'] = 'vmec'
# Update parameters in the namelist.
self.set_namelist(**keywords)
def finalize(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'massive_serial_runner: finalize')
def init(self, timeStamp=0.0, **keywords):
ScreenWriter.screen_output(self, 'verbose', 'massive_serial_runner: init')
# Stage state.
self.services.stage_state()
# Unzip files from the state. Use mode a so files an be read and written to.
if timeStamp == 0.0:
self.current_state = self.services.get_config_param('CURRENT_MSR_STATE')
self.zip_ref = ZipState.ZipState(self.current_state, 'a')
# Get config filenames.
if timeStamp == 0.0:
self.database_config = self.services.get_config_param('DATABASE_CONFIG')
self.zip_ref.extract(self.database_config)
self.current_batch = self.services.get_config_param('CURRENT_BATCH')
self.constraint_path = self.services.get_config_param('MODULE_PATH')
self.constraint_name = self.services.get_config_param('MODULE_NAME')
# Keys for the massiver serial subworkflow.
keys = {
'PWD' : self.services.get_config_param('PWD'),
'SIM_NAME' : 'massive_serial_runner_sub',
'LOG_FILE' : 'log.massive_serial_runner',
'NNODES' : self.services.get_config_param('MSR_NNODES'),
'INPUT_DIR_SIM' : 'massive_serial_runner_input_dir',
'OUTPUT_DIR_SIM' : '{}/massive_serial_runner_output_dir'.format(os.getcwd())
}
if os.path.exists('massive_serial_runner_input_dir'):
shutil.rmtree('massive_serial_runner_input_dir')
os.mkdir('massive_serial_runner_input_dir')
self.massive_serial_worker = {
'sim_name' : None,
'init' : None,
'driver' : None
}
msr_global = self.services.get_config_param('MSR_GLOBAL_CONFIG')
self.zip_ref.extract(msr_global)
(self.massive_serial_worker['sim_name'],
self.massive_serial_worker['init'],
self.massive_serial_worker['driver']) = self.services.create_sub_workflow('massive_serial',
msr_global,
keys,
'massive_serial_runner_input_dir')
self.zip_ref.extract('inscan')
shutil.copy2('inscan', 'massive_serial_runner_input_dir')
# These files should never change so only extract them once.
if timeStamp == 0.0:
ms_state = self.services.get_config_param('MSR_SERIAL_STATE')
self.zip_ref.extract(ms_state)
shutil.copy2(ms_state, 'massive_serial_runner_input_dir')
os.chdir('massive_serial_runner_input_dir')
# We need the input directory to exist in a directory called input. So we must
# make that directory first than extract the files. Remember to change back to
# the orginal working directory after extraction.
with ZipState.ZipState(ms_state, 'r') as zip_ref:
zip_ref.extractall()
with ZipState.ZipState('input.zip', 'r') as input_ref:
input_ref.extractall()
override = ConfigObj(infile=self.services.get_config_param('MSR_SERIAL_NODE_CONFIG'), interpolation='template', file_error=True)
override['INPUT_DIR_SIM'] = os.getcwd()
override.write()
override2 = ConfigObj(infile=self.services.get_config_param('MSR_MODEL_CONFIG'), interpolation='template', file_error=True)
override2['INPUT_DIR_SIM'] = os.getcwd()
override2.write()
os.chdir('../')
def finalize(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose',
'quasi_newton_init: finalize')
def init(self, timeStamp=0.0):
ScreenWriter.screen_output(self, 'verbose', 'cariddi_init: init')
# Get config filenames.
current_cariddi_state = self.services.get_config_param('CURRENT_CARIDDI_STATE')
current_vmec_namelist = self.services.get_config_param('VMEC_NAMELIST_INPUT')
current_vmec_state = self.services.get_config_param('CURRENT_VMEC_STATE')
current_siesta_namelist = self.services.get_config_param('SIESTA_NAMELIST_INPUT')
current_siesta_state = self.services.get_config_param('CURRENT_SIESTA_STATE')
current_v3fit_namelist = self.services.get_config_param('V3FIT_NAMELIST_INPUT')
current_v3fit_state = self.services.get_config_param('CURRENT_V3FIT_STATE')
# Remove old inputs. Stage input files.
for file in os.listdir('.'):
os.remove(file)
self.services.stage_input_files(self.INPUT_FILES)
# All v3fit runs require a vmec state at the minimum. Create a vmec state. If
# the vmec namelist file exists add the namelist input file.
with ZipState.ZipState(current_vmec_state, 'a') as zip_ref:
if os.path.exists(current_vmec_namelist):
zip_ref.write(current_vmec_namelist)
zip_ref.set_state(state='needs_update')
# A siesta state is optional. If a siesta state or namelist exist, create a
# siesta state. If the siesta namelist or vmec state files exists add
# them to the siesta state.
if os.path.exists(current_siesta_state) or os.path.exists(current_siesta_namelist):
with ZipState.ZipState(current_siesta_state, 'a') as zip_ref:
if os.path.exists(current_siesta_namelist):
zip_ref.write(current_siesta_namelist)
zip_ref.set_state(state='needs_update')
# The vmec state will be merged with any existing vmec state in the siesta
# state.
zip_ref.write(current_vmec_state)
# Create state from files. Input files can either be a new state, namelist
# input file or both. If both files were staged, replace the namelist input
# file. If the namelist file is present flag the state as needing to be
# updated.
with ZipState.ZipState(current_v3fit_state, 'a') as zip_ref:
if os.path.exists(current_v3fit_namelist):
zip_ref.write(current_v3fit_namelist)
zip_ref.set_state(state='needs_update')
# If a siesta state exists at this point add it to the archive. Otherwise add
# the vmec state.
if os.path.exists(current_siesta_state):
zip_ref.write(current_siesta_state)
else:
zip_ref.write(current_vmec_state)
# Create state from files. Input files can either be a new state, input file or
# both. If both files were staged, replace the input file. If the input file is
# present flag the state as needing to be updated.
with ZipState.ZipState(current_cariddi_state, 'a') as zip_ref:
zip_ref.write(current_v3fit_state)
self.services.update_state()
