def _run_tracking_model(directory, options):
print("Creating model and tracking...")
madx_script = _get_madx_script(BEAM, directory, options)
with silence():
madx_wrapper.resolve_and_run_string(
madx_script,
madx_path=MADX_PATH,
output_file=os.path.join(directory, 'job.test.madx'),
log_file=os.path.join(directory, 'madx_log.txt')
)
track_path = _get_track_path(directory, one=True)
tbt_path = _get_tbt_path(directory)
with silence():
ADDbpmerror.convert_files(infile=track_path, outfile=tbt_path)
headers = []
lines = []
with open(tbt_path, "r") as tbt_data:
for line in tbt_data:
if line.startswith("#"):
headers.append(line)
else:
lines.append(line)
with open(tbt_path, "w") as tbt_data:
for header in headers:
tbt_data.write(header)
for line in reversed(lines):
tbt_data.write(line)
def run_madx(madx_script, logfile=None, writeto=None):
madx_wrapper.resolve_and_run_string(madx_script,
output_file=writeto,
log_file=logfile)
def create_samples_with_validation(templates, models, optics, output_path,
madx_bin, index):
sample = None
print("Start creating dataset")
np.random.seed(seed=None)
print("Doing index: ", str(index))
seed = random.randint(0, 999999999)
with open(templates["B1"], 'r') as template:
template_str = template.read()
madx_wrapper.resolve_and_run_string(template_str % {
"DIR": output_path,
"INDEX": str(index),
"OPTICS": optics,
"SEED": seed
},
madx_path=madx_bin)
with open(templates["B2"], 'r') as template:
template_str = template.read()
madx_wrapper.resolve_and_run_string(template_str % {
"DIR": output_path,
"INDEX": str(index),
"OPTICS": optics,
"SEED": seed
},
madx_path=madx_bin)
# save data in sample
b1_errors_file_path = os.path.join(output_path,
"b1_errors_{}.tfs".format(index))
b2_errors_file_path = os.path.join(output_path,
"b2_errors_{}.tfs".format(index))
b1_tw_perturbed_path = os.path.join(output_path,
"b1_twiss_{}.tfs".format(index))
b2_tw_perturbed_path = os.path.join(output_path,
"b2_twiss_{}.tfs".format(index))
b1_tw_before_match = os.path.join(
output_path, "b1_twiss_before_match_{}.tfs".format(index))
b1_tw_after_match = os.path.join(
output_path, "b1_twiss_after_match_{}.tfs".format(index))
b2_tw_before_match = os.path.join(
output_path, "b2_twiss_before_match_{}.tfs".format(index))
b2_tw_after_match = os.path.join(
output_path, "b2_twiss_after_match_{}.tfs".format(index))
common_errors_path = os.path.join(output_path,
"common_errors_{}.tfs".format(index))
if os.path.isfile(b1_tw_perturbed_path) and os.path.isfile(
b2_tw_perturbed_path):
beta_star_b1, delta_mux_b1, delta_muy_b1, delta_dx_b1 = get_input_for_beam(
b1_tw_perturbed_path, models["B1"], 1)
beta_star_b2, delta_mux_b2, delta_muy_b2, delta_dx_b2 = get_input_for_beam(
b2_tw_perturbed_path, models["B2"], 2)
errors = get_errors_from_file(common_errors_path, b1_errors_file_path,
b2_errors_file_path, b1_tw_before_match,
b1_tw_after_match, b2_tw_before_match,
b2_tw_after_match)
# TODO: when creating validation set: activate it
# add_mqts_to_errtab(b1_errors_file_path, b1_tw_before_match, b1_tw_after_match)
# add_mqts_to_errtab(b2_errors_file_path, b2_tw_before_match, b2_tw_after_match)
sample = beta_star_b1, beta_star_b2, delta_mux_b1, delta_muy_b1, delta_dx_b1, delta_mux_b2, delta_muy_b2, delta_dx_b2, np.array(
errors, dtype=float)
os.remove(b1_tw_perturbed_path)
os.remove(b2_tw_perturbed_path)
os.remove(b1_errors_file_path)
os.remove(b2_errors_file_path)
os.remove(b1_tw_before_match)
os.remove(b1_tw_after_match)
os.remove(b2_tw_before_match)
os.remove(b2_tw_after_match)
os.remove(common_errors_path)
return sample
def create_nominal_twiss(optics, nominal_twiss_templ, madx_bin):
with open(nominal_twiss_templ, 'r') as template:
template_str = template.read()
madx_wrapper.resolve_and_run_string(template_str % {"OPTICS": optics},
madx_path=madx_bin)
def _callMadx(madx_script):
return madx_wrapper.resolve_and_run_string(
madx_script,
log_file=DEV_NULL, # TODO: Redirect to logger
)