def _write_chromatic_coupling_files(meas_path, cor_path):
LOG.debug("Calculating chromatic coupling diff.")
# TODO: Add Cf1010
try:
twiss_plus = read_tfs(join(split(cor_path)[0], TWISS_CORRECTED_PLUS), index='NAME')
twiss_min = read_tfs(join(split(cor_path)[0], TWISS_CORRECTED_MINUS), index='NAME')
except IOError:
LOG.debug("Chromatic coupling measurements not found. Skipped.")
else:
deltap = np.abs(twiss_plus.DELTAP - twiss_min.DELTAP)
plus = TwissOptics(twiss_plus, quick_init=True).get_coupling(method='cmatrix')
minus = TwissOptics(twiss_min, quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(plus, minus, how='inner', left_index=True, right_index=True,
suffixes=('_p', '_m'))
model['NAME'] = model.index.values
if exists(join(meas_path, "chromcoupling_free.out")):
meas = read_tfs(join(meas_path, "chromcoupling_free.out"))
else:
meas = read_tfs(join(meas_path, "chromcoupling.out"))
tw = pd.merge(meas, model, how='inner', on='NAME')
cf1001 = (tw.loc[:, 'F1001_p'] - tw.loc[:, 'F1001_m']) / deltap
tw['Cf1001r_model'] = np.real(cf1001)
tw['Cf1001i_model'] = np.imag(cf1001)
tw['Cf1001r_prediction'] = tw.loc[:, 'Cf1001r'] - tw.loc[:, 'Cf1001r_model']
tw['Cf1001i_prediction'] = tw.loc[:, 'Cf1001i'] - tw.loc[:, 'Cf1001i_model']
write_tfs(join(meas_path, get_diff_filename('chromatic_coupling')),
tw.loc[:, ['NAME', 'S',
'Cf1001r', 'Cf1001rERR',
'Cf1001i', 'Cf1001iERR',
'Cf1001r_model', 'Cf1001i_model',
'Cf1001r_prediction', 'Cf1001i_prediction']])
def init_from_model_dir(cls, model_dir): # prints only for debugging
print("Creating accelerator instance from model dir")
instance = cls()
instance.modelpath = model_dir
if os.path.isfile(model_dir):
model_dir = os.path.dirname(model_dir)
instance.model_tfs = tfs_pandas.read_tfs(os.path.join(model_dir, "twiss.dat")).set_index("NAME")
print("model path =", os.path.join(model_dir, "twiss.dat"))
try:
model_best_knowledge_path = os.path.join(model_dir, "twiss_best_knowledge.dat")
if os.path.isfile(model_best_knowledge_path):
instance.model_best_knowledge = tfs_pandas.read_tfs(model_best_knowledge_path).set_index("NAME")
except IOError:
instance.model_best_knowledge = None
elements_path = os.path.join(model_dir, "twiss_elements.dat")
if os.path.isfile(elements_path):
instance.elements = tfs_pandas.read_tfs(elements_path).set_index("NAME")
else:
raise AcceleratorDefinitionError("Elements twiss not found")
elements_path = os.path.join(model_dir, "twiss_elements_centre.dat")
if os.path.isfile(elements_path):
instance.elements_centre = tfs_pandas.read_tfs(elements_path).set_index("NAME")
else:
instance.elements_centre = instance.elements
instance.nat_tune_x = float(instance.model_tfs.headers["Q1"])
instance.nat_tune_y = float(instance.model_tfs.headers["Q2"])
return instance
def init_from_model_dir(self, model_dir):
LOGGER.debug("Creating accelerator instance from model dir")
self.model_dir = model_dir
LOGGER.debug(" model path = " + os.path.join(model_dir, "twiss.dat"))
self._model = tfs_pandas.read_tfs(os.path.join(model_dir, "twiss.dat"), index="NAME")
self.nat_tune_x = float(self._model.headers["Q1"])
self.nat_tune_y = float(self._model.headers["Q2"])
# Elements #####################################
elements_path = os.path.join(model_dir, "twiss_elements.dat")
if os.path.isfile(elements_path):
self._elements = tfs_pandas.read_tfs(elements_path, index="NAME")
else:
raise AcceleratorDefinitionError("Elements twiss not found")
center_path = os.path.join(model_dir, "twiss_elements_centre.dat")
if os.path.isfile(center_path):
self._elements_centre = tfs_pandas.read_tfs(center_path, index="NAME")
else:
self._elements_centre = self._elements
# Error Def #####################################
self._errordefspath = None
errordefspath = os.path.join(self.model_dir, "errordefs")
if os.path.exists(errordefspath):
self._errordefspath = errordefspath
else: # until we have a proper file name convention
errordefspath = os.path.join(self.model_dir, "error_deffs.txt")
if os.path.exists(errordefspath):
self._errordefspath = errordefspath
def _get_dfs(df_a, df_b):
if isinstance(df_a, basestring):
df_a = tfs.read_tfs(df_a)
if isinstance(df_b, basestring):
df_b = tfs.read_tfs(df_b)
return df_a, df_b
def main(input_path, input_path_model, output_path):
utils.iotools.create_dirs(output_path)
(beta_from_phase_x,
beta_from_amp_x) = _get_twiss_for_one_of(input_path, "getbetax_free.out",
"getbetax.out")
(beta_from_phase_y,
beta_from_amp_y) = _get_twiss_for_one_of(input_path, "getbetay_free.out",
"getbetay.out")
nominal_model = tfs_pandas.read_tfs(input_path_model)
_configure_plots()
files_phase = (beta_from_phase_x, beta_from_phase_y)
files_amplitude = (beta_from_amp_x, beta_from_amp_y)
beam = _get_beam_from_model(nominal_model)
for i in range(len(PLANES)):
for ip in IPS:
(summary_amplitude,
summary_beta) = _compute_calibration_for_ip_and_plane(
ip, PLANES[i], files_phase[i], files_amplitude[i],
nominal_model, beam, output_path)
name_file = OUTPUT_FILE_PREFIX_CALIBRATION_FILE + "_" + str(
PLANES[i].lower()) + ".out"
calibration_file_path = os.path.join(output_path, name_file)
try:
summary_amplitude_tot = tfs_pandas.read_tfs(
calibration_file_path)
summary_amplitude_tot = summary_amplitude_tot.append(
summary_amplitude)
tfs_pandas.write_tfs(calibration_file_path,
summary_amplitude_tot,
save_index=False)
except:
tfs_pandas.write_tfs(calibration_file_path,
summary_amplitude,
save_index=False)
def _get_twiss_for_one_of(directory, *file_names):
for file_name in file_names:
file_path = os.path.join(directory, file_name)
if os.path.isfile(file_path):
file_name_amplitude = "getampbeta" + file_name[7:]
file_path_amplitude = os.path.join(directory, file_name_amplitude)
return (tfs_pandas.read_tfs(file_path),
tfs_pandas.read_tfs(file_path_amplitude))
raise IOError("None of the files exist:\n\t" + "\n\t".join(file_names))
def getNDX(files,model,output):
#getting the list of BPMs and arcBPMs
file_list = [(file_name.strip() + ".linx") for file_name in files.strip("\"").split(",")]
file_dict = {}
model_tfs=tfs.read_tfs(model)
bpms=model_tfs.loc[:,"NAME"].values
for file_name in file_list:
filetfs = tfs.read_tfs(file_name)
bpms = intersect(bpms,filetfs.loc[:,"NAME"].values)
bpms=np.array(bpms)
arc_bpms=get_arc_bpms(model_tfs, bpms)
#merging the Dataframes
model_panda=load_panda(model)
for i,file_name in enumerate(file_list):
file_panda = load_panda(file_name)
model_panda = pd.merge(model_panda, file_panda, how='inner', on='NAME', suffixes=('', str(i+1)))
model_panda['NDXMDL'] = (model_panda.loc[:, 'DX'].values / np.sqrt(model_panda.loc[:, 'BETX'].values))
columns=['NAME','S', 'NDXMDL']
for c in model_panda.columns.values:
if c.startswith('AMPZ') or c.startswith('MUZ'):
columns.append(c)
results = model_panda.loc[:,columns]
results.set_index("NAME", inplace=True, drop=False)
columns=['S', 'NDXMDL']
cols=[]
# scaling to the model, and getting the synchrotron phase in the arcs
for c in results.columns.values:
if c.startswith('MUZ'):
results['sc'+c.replace('MU','AMP')]=results.loc[:,c.replace('MU','AMP')].values * np.sum(results.loc[arc_bpms,'NDXMDL']) / np.sum(results.loc[arc_bpms,c.replace('MU','AMP')])
results['s'+c]=np.angle(np.exp(PI2I*results.loc[:,c].values))/(2*np.pi)
field = results.loc[:,'s'+c].values- np.angle(np.sum(np.exp(PI2I * results.loc[arc_bpms,'s'+c])))/(2*np.pi)
results['sc'+c]=np.abs(np.where(np.abs(field)>0.5,field-np.sign(field),field))
d='sc'+c
cols.append(d)
#resolving the sign of dispersion
for c in cols:
results[c.replace('scMUZ','fNDX')]=results.loc[:,c.replace('MU','AMP')] * np.sign(0.25 - np.abs(np.angle(np.sum(np.exp(PI2I * results.loc[:,cols]),axis=1)))/(2*np.pi))
columns.append(c.replace('scMUZ','fNDX'))
forfile=results.loc[:,columns]
f=[]
#averaging over files and error calculation
for c in forfile.columns.values:
if c.startswith('fNDX'):
f.append(c)
if len(f)>1:
forfile['STDNDX']=np.std(forfile.loc[:,f],axis=1)*t_value_correction(len(f))
else:
forfile['STDNDX']=0.0
forfile['NDX']=np.mean(forfile.loc[:,f],axis=1)
forfile['DNDX']=forfile.loc[:,'NDX']-forfile.loc[:,'NDXMDL']
forfile['NAME']=results.index
tfs.write_tfs(os.path.join(output, "getNDx.out"), forfile)
return
def _get_bpm_names(orbit_path_left, orbit_path_right):
return {
(BEAM1, LEFT): tfs_pandas.read_tfs(
os.path.join(orbit_path_left, "LHCB1.orbit1.tfs"))["NAME"],
(BEAM1, RIGHT): tfs_pandas.read_tfs(
os.path.join(orbit_path_right, "LHCB1.orbit1.tfs"))["NAME"],
(BEAM2, LEFT): tfs_pandas.read_tfs(
os.path.join(orbit_path_left, "LHCB2.orbit1.tfs"))["NAME"],
(BEAM2, RIGHT): tfs_pandas.read_tfs(
os.path.join(orbit_path_right, "LHCB2.orbit1.tfs"))["NAME"],
}
def getdiff(meas_path=None, beta_file_name="getbeta"):
""" Calculates the differences between measurement, corrected and uncorrected model.
After running madx and creating the model with (twiss_cor) and without (twiss_no)
corrections, run this functions to get tfs-files with the differences between measurements
and models.
Args:
meas_path (str): Path to the measurement folder.
Needs to contain twiss_cor.dat and twiss_no.dat.
"""
if meas_path is None:
meas_path = sys.argv[1]
LOG.debug("Started 'getdiff' for measurment dir '{:s}'".format(meas_path))
if not isdir(meas_path):
raise IOError("No valid measurement directory:" + meas_path)
corrected_model_path = join(meas_path, TWISS_CORRECTED)
uncorrected_model_path = join(meas_path, TWISS_NOT_CORRECTED)
meas = OpticsMeasurement(meas_path)
twiss_cor = read_tfs(corrected_model_path).set_index('NAME')
twiss_no = read_tfs(uncorrected_model_path).set_index('NAME')
coup_cor = TwissOptics(twiss_cor,
quick_init=True).get_coupling(method='cmatrix')
coup_no = TwissOptics(twiss_no,
quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(twiss_cor,
twiss_no,
how='inner',
left_index=True,
right_index=True,
suffixes=('_c', '_n'))
coupling_model = pd.merge(coup_cor,
coup_no,
how='inner',
left_index=True,
right_index=True,
suffixes=('_c', '_n'))
#coupling_model['NAME'] = coupling_model.index.values
for plane in ['x', 'y']:
_write_betabeat_diff_file(meas_path, meas, model, plane,
beta_file_name)
_write_phase_diff_file(meas_path, meas, model, plane)
_write_disp_diff_file(meas_path, meas, model, plane)
_write_closed_orbit_diff_file(meas_path, meas, model, plane)
_write_coupling_diff_file(meas_path, meas, coupling_model)
_write_norm_disp_diff_file(meas_path, meas, model)
_write_chromatic_coupling_files(meas_path, corrected_model_path)
_write_betastar_diff_file(meas_path, meas, twiss_cor, twiss_no)
LOG.debug("Finished 'getdiff'.")
def generate_lin_files(model_dir, outfile='test', dpp=0.0):
free = read_tfs(join(model_dir, 'twiss.dat'))
driven = read_tfs(join(model_dir, 'twiss_ac.dat'))
nattune = {"X": np.remainder(free.headers['Q1'], 1), "Y": np.remainder(free.headers['Q2'], 1)}
tune = {"X": np.remainder(driven.headers['Q1'], 1), "Y": np.remainder(driven.headers['Q2'], 1)}
model = pd.merge(free, driven, how='inner', on='NAME', suffixes=('_f', '_d'))
model['S'] = model.loc[:, 'S_f']
nbpms = len(model.index.values)
coup = 0.01
for plane in ['X', 'Y']:
lin = model.loc[:, ['NAME', 'S']]
lin['NOISE'] = np.abs(np.random.randn(nbpms) * 0.0002 + 0.0002)
lin['AVG_NOISE'] = lin.loc[:, 'NOISE']
lin['CO'] = dpp * 1000 * model.loc[:, 'D' + plane + '_d'] # meters to millimeters
lin['CORMS'] = np.abs(np.random.randn(nbpms) * 0.003 + 0.003)
lin['PK2PK'] = 0.07 * np.sqrt(model.loc[:, 'BET' + plane + '_d'])
lin['TUNE' + plane] = tune[plane] + 1e-7 * np.random.randn(nbpms)
lin['MU' + plane] = np.remainder(model.loc[:, 'MU' + plane + '_d']
+ dpp * model.loc[:, 'DMU' + plane + '_d']
+ 0.0001 *np.random.randn(nbpms) + np.random.rand(), 1)
lin['AMP' + plane] = 0.03 * np.sqrt(model.loc[:, 'BET' + plane + '_d'] *
(1 + dpp * np.sin(2 * np.pi * model.loc[:, 'PHI' + plane + '_d'])
* model.loc[:, 'W' + plane + '_d']))\
+ 0.001 *np.random.randn(nbpms)
lin['NATTUNE' + plane] = nattune[plane] + 1e-6 * np.random.randn(nbpms)
lin['NATMU' + plane] = np.remainder(model.loc[:, 'MU' + plane + '_f']
+ 0.001 * np.random.randn(nbpms) + np.random.rand(), 1)
lin['NATAMP' + plane] = 0.0003 * np.sqrt(model.loc[:, 'BET' + plane + '_f'])\
+ 0.00001 *np.random.randn(nbpms)
plane_number = {"X": "1", "Y": "2"}[plane]
header = OrderedDict()
header["Q" + plane_number] = tune[plane]
header["Q" + plane_number + "RMS"] = 1e-7
header["DPP"] = dpp
header["NATQ" + plane_number] = nattune[plane]
header["NATQ" + plane_number + "RMS"] = 1e-6
if plane == 'X':
lin['PHASE01'] = np.remainder(model.loc[:, 'MUY_d'] + dpp * model.loc[:, 'DMUY_d']
+ 0.0001 *np.random.randn(nbpms) + np.random.rand(), 1)
lin['AMP01'] = coup * np.sqrt(model.loc[:, 'BETY_d']
* (1 + dpp * np.sin(model.loc[:, 'PHIY_d']) * model.loc[:, 'WY_d'])) \
+ 0.0001 * np.random.randn(nbpms)
write_tfs(outfile + '.linx', lin, header)
else:
lin['PHASE10'] = np.remainder(model.loc[:, 'MUX_d']
+ 0.0001 *np.random.randn(nbpms) + np.random.rand(), 1)
lin['AMP10'] = coup * np.sqrt(model.loc[:, 'BETX_d'] \
* (1 + dpp * np.sin(model.loc[:, 'PHIX_d']) * model.loc[:, 'WX_d'])) \
+ 0.0001 * np.random.randn(nbpms)
write_tfs(outfile + '.liny', lin, header)
def _get_model_tunes(opt):
twiss_kick_path = _get_twiss_kicker_path(opt)
if not twiss_kick_path:
raise IOError(
"Tune Model not found. Hint: Works only with adt or acd kicks!")
df_kick = tfs.read_tfs(twiss_kick_path)
df = tfs.read_tfs(_get_twiss_path(opt))
data = {
'tunex': df_kick.headers['Q1'] % 1,
'tuney': df_kick.headers['Q2'] % 1,
'nattunex': df.headers['Q1'] % 1,
'nattuney': df.headers['Q2'] % 1,
}
return data
def get_dispersion_normalized_dispersion_ir(dispersion_file,
normalized_dispersion_file,
bpm_set):
dispersion_measurement = tfs_pandas.read_tfs(dispersion_file)
normalized_dispersion_measurement = tfs_pandas.read_tfs(
normalized_dispersion_file)
dispersion = dispersion_measurement.set_index("NAME").loc[bpm_set, "DX"]
dispersion_error = dispersion_measurement.set_index("NAME").loc[bpm_set,
"STDDX"]
normalized_dispersion = normalized_dispersion_measurement.set_index(
"NAME").loc[bpm_set, "NDX"]
normalized_dispersion_error = normalized_dispersion_measurement.set_index(
"NAME").loc[bpm_set, "STDNDX"]
return dispersion, dispersion_error, normalized_dispersion, normalized_dispersion_error
def _get_bpm_names(orbit_path_left, orbit_path_right):
return {
(BEAM1, LEFT):
tfs_pandas.read_tfs(os.path.join(orbit_path_left,
"LHCB1.orbit1.tfs"))["NAME"],
(BEAM1, RIGHT):
tfs_pandas.read_tfs(os.path.join(orbit_path_right,
"LHCB1.orbit1.tfs"))["NAME"],
(BEAM2, LEFT):
tfs_pandas.read_tfs(os.path.join(orbit_path_left,
"LHCB2.orbit1.tfs"))["NAME"],
(BEAM2, RIGHT):
tfs_pandas.read_tfs(os.path.join(orbit_path_right,
"LHCB2.orbit1.tfs"))["NAME"],
}
def _log_rms(files, legends, column_name, mask):
""" Calculate and print rms value into log """
file_name = os.path.splitext(os.path.basename(files[0]))[0]
collected = {}
LOG.info("Results for '{:s}':".format(file_name))
if mask is not None:
LOG.info(" {:<20s} {:^15s} {:^15s} {:^15s} {:^15s}".format(
"", "RMS", "Mean", "RMS w/ cut", "Mean w/ cut"))
else:
LOG.info(" {:<20s} {:^15s} {:^15s}".format("", "RMS", "Mean"))
for f, l in zip(files, legends):
collected[l] = {}
data = tfs_pandas.read_tfs(f, index="NAME")[column_name]
collected[l]["rms"] = _rms(data)
collected[l]["mean"] = np.mean(data)
if mask is not None:
data = data.loc[mask]
rms_val = _rms(data)
mean_val = np.mean(data)
collected[l]["rms_f"] = rms_val
collected[l]["mean_f"] = mean_val
LOG.info(" {:<20s} {:+15.4e} {:+15.4e} {:+15.4e} {:+15.4e}".format(
l, collected[l]["rms"], collected[l]["mean"],
collected[l]["rms_f"], collected[l]["mean_f"]
))
else:
LOG.info(" {:<20s} {:+15.4e} {:+15.4e}".format(
l, collected[l]["rms"], collected[l]["mean"]))
return collected
def get_variables(cls, frm=None, to=None, classes=None):
correctors_dir = os.path.join(LHC_DIR, "2012", "correctors")
all_corrs = _merge_jsons(
os.path.join(correctors_dir, "correctors_b" + str(cls.get_beam()),
"beta_correctors.json"),
os.path.join(correctors_dir, "correctors_b" + str(cls.get_beam()),
"coupling_correctors.json"),
cls._get_triplet_correctors_file(),
)
my_classes = classes
if my_classes is None:
my_classes = all_corrs.keys()
vars_by_class = set(_flatten_list(
[all_corrs[corr_cls] for corr_cls in my_classes if corr_cls in all_corrs])
)
if frm is None and to is None:
return list(vars_by_class)
elems_matrix = tfs_pandas.read_tfs(
cls._get_corrector_elems()
).sort_values("S")
if frm is not None and to is not None:
if frm > to:
elems_matrix = elems_matrix[(elems_matrix.S >= frm) | (elems_matrix.S <= to)]
else:
elems_matrix = elems_matrix[(elems_matrix.S >= frm) & (elems_matrix.S <= to)]
elif frm is not None:
elems_matrix = elems_matrix[elems_matrix.S >= frm]
elif to is not None:
elems_matrix = elems_matrix[elems_matrix.S <= to]
vars_by_position = _remove_dups_keep_order(_flatten_list(
[raw_vars.split(",") for raw_vars in elems_matrix.loc[:, "VARS"]]
))
return _list_intersect_keep_order(vars_by_position, vars_by_class)
def get_segment(cls, label, first_elem, last_elem, optics_file, twiss_file):
# this creates a new class called PsSegment
segment_cls = type(cls.__name__ + "Segment",
(_PsSegmentMixin,cls),
{})
segment_inst = segment_cls()
bpms_file = os.path.join(PS_DIR,str(CURRENT_YEAR),"sequence/bpms.tfs")
bpms_file_data = tfs_pandas.read_tfs(bpms_file).set_index("NAME")
first_elem_s = bpms_file_data.loc[first_elem, "S"]
last_elem_s = bpms_file_data.loc[last_elem, "S"]
segment_inst.label = label
segment_inst.start = Element(first_elem, first_elem_s)
segment_inst.end = Element(last_elem, last_elem_s)
segment_inst.optics_file = optics_file
segment_inst.fullresponse = None
LOGGER.debug('twiss_file is <%s>',twiss_file)
tw = twiss(twiss_file)
LOGGER.debug('twiss_file has tunes %f %f ',tw.Q1,tw.Q2)
segment_inst.nat_tune_x = tw.Q1
segment_inst.nat_tune_y = tw.Q2
segment_inst.energy = tw.ENERGY
segment_inst.kind = '' # '' means beta from phase, can be 'betaamp', in the future 'betakmod'
return segment_inst
def run_all_for_file(tbt_file, main_input, clean_input, harpy_input):
tbt_file = _cut_tbt_file(tbt_file, main_input.startturn,
main_input.endturn)
file_date = tbt_file.date
bpm_datas = {
"x": tbt_file.samples_matrix_x,
"y": tbt_file.samples_matrix_y
}
if main_input.write_raw:
output_handler.write_raw_file(tbt_file, main_input)
elif clean_input is None and harpy_input is None:
raise ValueError("No clean or harpy or --write_raw...")
# Defaults in case we dont run clean
usvs = {"x": None, "y": None}
all_bad_bpms = {"x": [], "y": []}
bpm_ress = {"x": None, "y": None}
dpp = 0.0
model_tfs = tfs.read_tfs(main_input.model).loc[:, ('NAME', 'S', 'DX')]
if clean_input is not None:
usvs, all_bad_bpms, bpm_ress, dpp = _do_clean(main_input, clean_input,
bpm_datas, file_date,
model_tfs)
lin = {"x": None, "y": None}
if harpy_input is not None:
all_bad_bpms, lin = _do_harpy(main_input, harpy_input, bpm_datas, usvs,
model_tfs, bpm_ress, dpp, all_bad_bpms)
for plane in ("x", "y"):
output_handler.write_bad_bpms(main_input.file, all_bad_bpms[plane],
main_input.outputdir, plane)
return lin
def clean_tunes(files, limit=DEF_LIMIT):
for file in files:
file_df = tfs_pandas.read_tfs(file)
mask = _get_mask(file_df, limit)
file_df = file_df.loc[mask, :]
_recompute_tune_stats(file_df)
tfs_pandas.write_tfs(file, file_df)
def get_segment(cls, label, first_elem, last_elem, optics_file,
twiss_file):
segment_cls = type(cls.__name__ + "Segment",
(_PsboosterSegmentMixin, cls), {})
LOGGER.debug('twiss_file is <%s>', twiss_file)
tw = twiss(twiss_file)
LOGGER.debug('twiss_file has tunes %f %f ', tw.Q1, tw.Q2)
ring = _get_ring_from_seqname(tw.SEQUENCE)
#ring = cls.get_ring()
segment_inst = segment_cls()
bpms_file = _get_file_for_ring(ring)
bpms_file_data = tfs_pandas.read_tfs(bpms_file).set_index("NAME")
first_elem_s = bpms_file_data.loc[first_elem, "S"]
last_elem_s = bpms_file_data.loc[last_elem, "S"]
segment_inst.label = label
segment_inst.start = Element(first_elem, first_elem_s)
segment_inst.end = Element(last_elem, last_elem_s)
segment_inst.optics_file = optics_file
segment_inst.fullresponse = None
segment_inst.nat_tune_x = tw.Q1
segment_inst.nat_tune_y = tw.Q2
segment_inst.energy = tw.ENERGY
segment_inst.sequence = tw.SEQUENCE
segment_inst.ring = ring
segment_inst.kind = '' # '' means beta from phase, can be 'betaamp', in the future 'betakmod'
return segment_inst
def get_fill_from_orbitfile(orbit_file):
""" Extracts the fill number from orbit file (-oe switch in online model extractor).
Args:
orbit_file: Path to orbit file.
"""
return int(tfs.read_tfs(orbit_file).headers[FILLNUMBER])
def _log_rms(files, legends, column_name, mask):
""" Calculate and print rms value into log """
file_name = os.path.splitext(os.path.basename(files[0]))[0]
collected = {}
LOG.info("Results for '{:s}':".format(file_name))
if mask is not None:
LOG.info(" {:<20s} {:^15s} {:^15s} {:^15s} {:^15s}".format(
"", "RMS", "Mean", "RMS w/ cut", "Mean w/ cut"))
else:
LOG.info(" {:<20s} {:^15s} {:^15s}".format("", "RMS", "Mean"))
for f, l in zip(files, legends):
collected[l] = {}
data = tfs_pandas.read_tfs(f, index="NAME")[column_name]
collected[l]["rms"] = _rms(data)
collected[l]["mean"] = np.mean(data)
if mask is not None:
data = data.loc[mask]
rms_val = _rms(data)
mean_val = np.mean(data)
collected[l]["rms_f"] = rms_val
collected[l]["mean_f"] = mean_val
LOG.info(" {:<20s} {:+15.4e} {:+15.4e} {:+15.4e} {:+15.4e}".format(
l, collected[l]["rms"], collected[l]["mean"],
collected[l]["rms_f"], collected[l]["mean_f"]
))
else:
LOG.info(" {:<20s} {:+15.4e} {:+15.4e}".format(
l, collected[l]["rms"], collected[l]["mean"]))
return collected
def get_beta_from_phase_ir(beta_phase_file, bpm_set):
beta_phase_measurement = tfs_pandas.read_tfs(beta_phase_file)
position_phase = beta_phase_measurement.set_index("NAME").loc[bpm_set, "S"]
beta_phase = beta_phase_measurement.set_index("NAME").loc[bpm_set, "BETX"]
beta_phase_error = beta_phase_measurement.set_index("NAME").loc[bpm_set,
"ERRBETX"]
return position_phase, beta_phase, beta_phase_error
def get_fill_from_orbitfile(orbit_file):
""" Extracts the fill number from orbit file (-oe switch in online model extractor).
Args:
orbit_file: Path to orbit file.
"""
return int(tfs.read_tfs(orbit_file).headers[FILLNUMBER])
def clean_tunes(files, limit=DEF_LIMIT):
for file in files:
file_df = tfs_pandas.read_tfs(file)
mask = _get_mask(file_df, limit)
file_df = file_df.loc[mask, :]
_recompute_tune_stats(file_df)
tfs_pandas.write_tfs(file, file_df)
def get_segment(cls, label, first_elem, last_elem, optics_file,
twiss_file):
# this creates a new class called PsSegment
segment_cls = type(cls.__name__ + "Segment", (_PsSegmentMixin, cls),
{})
segment_inst = segment_cls()
bpms_file = os.path.join(PS_DIR, str(CURRENT_YEAR),
"sequence/bpms.tfs")
bpms_file_data = tfs_pandas.read_tfs(bpms_file).set_index("NAME")
first_elem_s = bpms_file_data.loc[first_elem, "S"]
last_elem_s = bpms_file_data.loc[last_elem, "S"]
segment_inst.label = label
segment_inst.start = Element(first_elem, first_elem_s)
segment_inst.end = Element(last_elem, last_elem_s)
segment_inst.optics_file = optics_file
segment_inst.fullresponse = None
LOGGER.debug('twiss_file is <%s>', twiss_file)
tw = twiss(twiss_file)
LOGGER.debug('twiss_file has tunes %f %f ', tw.Q1, tw.Q2)
segment_inst.nat_tune_x = tw.Q1
segment_inst.nat_tune_y = tw.Q2
segment_inst.energy = tw.ENERGY
segment_inst.kind = '' # '' means beta from phase, can be 'betaamp', in the future 'betakmod'
return segment_inst
def _get_beta_beat_file(self, label, plane):
outpath = self._matcher_model.get_output_path()
file_data = tfs_pandas.read_tfs(os.path.join(
outpath, "sbs",
"sbsbetabeat" + plane.lower() + "_" + label + ".out")
)
return file_data
def _get_amp_beta_beat_file(self, label, plane):
outpath = self._matcher_model.get_output_path()
file_data = tfs_pandas.read_tfs(
os.path.join(
outpath, "sbs",
"sbsampbetabeat" + plane.lower() + "_" + label + ".out"))
return file_data
def get_beta_from_amplitude_ir(beta_amp_file, bpm_set):
beta_amp_measurement = tfs_pandas.read_tfs(beta_amp_file)
position_amp = beta_amp_measurement.set_index("NAME").loc[bpm_set, "S"]
beta_amp = beta_amp_measurement.set_index("NAME").loc[bpm_set, "BETX"]
beta_amp_error = beta_amp_measurement.set_index("NAME").loc[bpm_set,
"BETXSTD"]
return position_amp, beta_amp, beta_amp_error
def _get_timber_data(beam, input, output, kickac_df):
""" Return Timber data from input """
try:
fill_number = int(input)
except ValueError:
# input is a string
LOG.debug("Getting timber data from file '{:s}'".format(input))
data = tfs.read_tfs(input, index=COL_TIME())
data.drop([COL_MAV(p) for p in PLANES if COL_MAV(p) in data.columns],
axis='columns')
except TypeError:
# input is None
LOG.debug("Getting timber data from kickac-times.")
timber_keys, bbq_cols = _get_timber_keys_and_bbq_columns(beam)
t_start = min(kickac_df.index.values)
t_end = max(kickac_df.index.values)
data = timber_extract.extract_between_times(t_start-DTIME, t_end+DTIME,
keys=timber_keys,
names=dict(zip(timber_keys, bbq_cols)))
else:
# input is a number
LOG.debug("Getting timber data from fill '{:d}'".format(input))
timber_keys, bbq_cols = _get_timber_keys_and_bbq_columns(beam)
data = timber_extract.lhc_fill_to_tfs(fill_number,
keys=timber_keys,
names=dict(zip(timber_keys, bbq_cols)))
if output:
tfs.write_tfs(output, data, save_index=COL_TIME())
return data
def _get_param_data(param, method, opt):
"""
Get all getLLM data as DataFrame from subfolder in method and parameter param
"""
directory = _get_getllm_fullpath(method, opt)
df = tfs.read_tfs(os.path.join(directory, 'get' + param + '_free.out'))
df.set_index('NAME', inplace=True)
return df
def _create_tfs_data(filepaths, plane):
bpm_data_rows = []
for filepath in filepaths:
bpm_tfs_file = tfs_pandas.read_tfs(filepath)
bpms_tfs_data = bpm_tfs_file[FEATURES_WITH_NAME.format(
plane.upper()).split(",")]
bpm_data_rows.append(bpms_tfs_data)
return pandas.concat(bpm_data_rows)
def get_values_from_tfs(tfs_path):
df = tfs.read_tfs(tfs_path)
anhx1000 = df.query('NAME == "ANHX" and '
'ORDER1 == 1 and ORDER2 == 0')["VALUE"].values[0]
anhy0100 = df.query('NAME == "ANHY" and ORDER1 == 0 and '
'ORDER2 == 1')["VALUE"].values[0]
anhx0100 = df.query('NAME == "ANHX" and ORDER1 == 0 and '
'ORDER2 == 1')["VALUE"].values[0]
return anhx1000, anhy0100, anhx0100
def _copy_calibration_files(outputdir, calibrationdir):
if calibrationdir is None:
return None
calibs = {}
for plane in PLANES:
cal_file = "calibration_{}.out".format(plane.lower())
iotools.copy_item(join(calibrationdir, cal_file), join(outputdir, cal_file))
calibs[plane] = tfs_pandas.read_tfs(join(outputdir, cal_file)).set_index("NAME")
return calibs
def _load_and_remove_twiss(var_and_path):
""" Function for pool to retrieve results """
(var, path) = var_and_path
twissfile = os.path.join(path, "twiss." + var)
tfs_data = tfs.read_tfs(twissfile, index="NAME")
tfs_data['Q1'] = tfs_data.Q1
tfs_data['Q2'] = tfs_data.Q2
os.remove(twissfile)
return var, tfs_data
def _load_and_remove_twiss(var_and_path):
""" Function for pool to retrieve results """
(var, path) = var_and_path
twissfile = os.path.join(path, "twiss." + var)
tfs_data = tfs.read_tfs(twissfile, index="NAME")
tfs_data['Q1'] = tfs_data.Q1
tfs_data['Q2'] = tfs_data.Q2
os.remove(twissfile)
return var, tfs_data
def compute_offset(model1_path, model2_path, orbit_path_left, orbit_path_right,
kleft_path, kright_path, ip):
ks, orbits = _collect_orbit_data(orbit_path_left, orbit_path_right,
kleft_path, kright_path)
model1 = tfs_pandas.read_tfs(model1_path)
model1.set_index("NAME", inplace=True)
model2 = tfs_pandas.read_tfs(model2_path)
model2.set_index("NAME", inplace=True)
bpm_names = _get_bpm_names(orbit_path_left, orbit_path_right)
models = {BEAM1: model1, BEAM2: model2}
results = _apply_to_beam_side_plane(
lambda beam, side, plane: _compute_and_clean(
ip, beam, side, plane, models, bpm_names, ks, orbits))
return results
def _load_source_data(src, index=None):
data = None
try:
data = tfs_pandas.read_tfs(src, index=index)
except IOError:
LOG.warn("Cannot find kmod data in '{:s}', won't copy those files.".format(src))
else:
LOG.warn("Loaded kmod data from '{:s}'.".format(src))
return data
def get_phases(files, model, output):
for plane in ["X","Y"]:
file_list = [(file_name.strip() + ".lin" + plane.lower()) for file_name in files.strip("\"").split(",")]
#merging the Dataframes
model_panda = load_panda(model)
tune=0.0
for i,file_name in enumerate(file_list):
file_panda = tfs.read_tfs(file_name)
if plane =="X":
tune = tune + file_panda.headers['Q1']
else:
tune = tune + file_panda.headers['Q2']
file_panda = pd.DataFrame(file_panda)
model_panda = pd.merge(model_panda, file_panda, how='inner', on='NAME', suffixes=('', str(i+1)))
tune = tune / len(file_list)
model_panda.rename(columns={'MU'+plane:'MU'+plane+'MDL'}, inplace=True)
columns=['NAME','S']
for c in model_panda.columns.values:
if c.startswith('MU'+plane):
columns.append(c)
all_data = model_panda.loc[:,columns]
all_data.set_index("NAME", inplace=True, drop=False)
#Here is what we need from the model and all the measured phases for the intersected BPMs
bpms = all_data.loc[:,'NAME'].values
columns = ['NAME', 'S', 'MU' + plane + 'MDL']
results = all_data.loc[bpms[:-1], columns]
results['NAME2'] = bpms[1:]
results['S2'] = all_data.loc[bpms[1:],'S'].values
cols=[]
for c in all_data.columns.values:
if c.startswith('MU'):
field = all_data.loc[bpms[1:],c].values - all_data.loc[bpms[:-1],c].values
results[c.replace('MU','PHASE')] = np.where(np.abs(field)>0.5,field-np.sign(field),field)
d = c.replace('MU','PHASE')
cols.append(d)
cols.remove('PHASE' + plane + 'MDL')
results.rename(columns={'PHASE' + plane + 'MDL':'PH' + plane + 'MDL'}, inplace=True)
results['PHASE'+plane]= np.angle(np.sum(np.exp(PI2I * results.loc[:,cols]),axis=1))/(2*np.pi)
f =[]
for c in cols:
field = results.loc[:,c] - results.loc[:,'PHASE'+plane]
results['d'+ c] = np.where(np.abs(field)>0.5,field-np.sign(field),field)
f.append('d'+ c)
if len(f)>1:
results['STDPH' + plane]=np.std(results.loc[:,f],axis=1)*t_value_correction(len(f))
else:
results['STDPH' + plane]=0.0
if plane =="X":
header = {'Q1': tune}
else:
header = {'Q2': tune}
heads = ['NAME', 'NAME2', 'S', 'S2', 'PHASE'+plane, 'STDPH' + plane, 'PH' + plane+'MDL', 'MU' + plane+'MDL']
tfs.write_tfs(os.path.join(output, "getphase" + plane.lower() + ".out"), results.loc[:, heads], header)
return
def _write_chromatic_coupling_files(meas_path, cor_path):
LOG.debug("Calculating chromatic coupling diff.")
# TODO: Add Cf1010
try:
twiss_plus = read_tfs(join(split(cor_path)[0], TWISS_CORRECTED_PLUS),
index='NAME')
twiss_min = read_tfs(join(split(cor_path)[0], TWISS_CORRECTED_MINUS),
index='NAME')
except IOError:
LOG.debug("Chromatic coupling measurements not found. Skipped.")
else:
deltap = np.abs(twiss_plus.DELTAP - twiss_min.DELTAP)
plus = TwissOptics(twiss_plus,
quick_init=True).get_coupling(method='cmatrix')
minus = TwissOptics(twiss_min,
quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(plus,
minus,
how='inner',
left_index=True,
right_index=True,
suffixes=('_p', '_m'))
model['NAME'] = model.index.values
if exists(join(meas_path, "chromcoupling_free.out")):
meas = read_tfs(join(meas_path, "chromcoupling_free.out"))
else:
meas = read_tfs(join(meas_path, "chromcoupling.out"))
tw = pd.merge(meas, model, how='inner', on='NAME')
cf1001 = (tw.loc[:, 'F1001_p'] - tw.loc[:, 'F1001_m']) / deltap
tw['Cf1001r_model'] = np.real(cf1001)
tw['Cf1001i_model'] = np.imag(cf1001)
tw['Cf1001r_prediction'] = tw.loc[:,
'Cf1001r'] - tw.loc[:,
'Cf1001r_model']
tw['Cf1001i_prediction'] = tw.loc[:,
'Cf1001i'] - tw.loc[:,
'Cf1001i_model']
write_tfs(
join(meas_path, get_diff_filename('chromatic_coupling')),
tw.loc[:, [
'NAME', 'S', 'Cf1001r', 'Cf1001rERR', 'Cf1001i', 'Cf1001iERR',
'Cf1001r_model', 'Cf1001i_model', 'Cf1001r_prediction',
'Cf1001i_prediction'
]])
def make_histogram_plots(opt):
alpha_mean = .2
alpha_hist = .6
title = ""
ps.set_style('standard')
opt = tripcor.check_opt(opt)
cwd = get_data_output_folder(opt.cwd)
output_folder = get_plot_output_folder(opt.cwd)
for beam, xing, error_types, error_loc, optic_type in hist_loop(opt):
fig, axs = plt.subplots(len(AMPDET_NAMES), 1)
for idx_data, output_id in _ordered_output_ids(opt.machine, beam,
opt.unused_stages):
title, filename = get_seed_data_title_and_filename(
beam, xing, error_types, error_loc, optic_type, output_id)
seed_df = tfs.read_tfs(os.path.join(cwd, filename), index="SEED")
y_max = len(seed_df.index)
for idx_ax, term in enumerate(AMPDET_NAMES):
ax = axs[idx_ax]
data = seed_df[term]
x_pos = data.mean()
# plot mean
stem_cont = ax.stem([x_pos], [y_max],
markerfmt="",
basefmt="",
label="_nolegend_")
plt.setp(stem_cont[1],
color=ps.get_mpl_color(idx_data),
alpha=alpha_mean)
# plot std
# error = data.std()
# ebar_cont = ax.errorbar(x_pos, y_max, xerr=error, color=ps.get_mpl_color(idx_data),
# label="_nolegend_", marker="")
# ps.change_ebar_alpha_for_line(ebar_cont, alpha_mean)
# plot histogram
data.hist(ax=ax,
alpha=alpha_hist,
color=ps.get_mpl_color(idx_data),
label=output_id)
for idx_ax, term in enumerate(AMPDET_NAMES):
axs[idx_ax].set_xlabel(term)
axs[idx_ax].set_ylabel("Seeds")
legend = axs[0].legend()
_reorder_legend(
legend, _get_all_output_ids(opt.machine, beam, opt.unused_stages))
ps.sync2d(axs)
title = " ".join(title.split(" ")[:-1])
figfilename = "{:s}.{:s}".format(title.replace(' ', '.'), "histogram")
fig.canvas.set_window_title(title)
make_bottom_text(axs[-1], title)
fig.savefig(os.path.join(output_folder, figfilename + ".pdf"))
fig.savefig(os.path.join(output_folder, figfilename + ".png"))
def _assign_uncertainties(twiss_full, errordefspath):
'''
Adds uncertainty information to twiss_full.
:Sources of Errors:
dK1: quadrupolar field errors
dS: quadrupole longitudinal misalignments
dX: sextupole transverse misalignments
BPMdS: BPM longitudinal misalignments
'''
LOGGER.debug("Start creating uncertainty information")
errdefs = tfs_pandas.read_tfs(errordefspath)
# create new columns
twiss_full = twiss_full.assign(UNC=False,
dK1=0,
KdS=0,
mKdS=0,
dX=0,
BPMdS=0)
# loop over uncertainty definitions, fill the respective columns, set UNC to true
for indx in errdefs.index:
patt = errdefs.loc[indx, "PATTERN"]
if patt.startswith("key:"):
LOGGER.debug(
"creating uncertainty information for {:s}".format(patt))
mask = patt.split(":")[1]
else:
reg = re.compile(patt)
LOGGER.debug(
"creating uncertainty information for RegEx {:s}".format(patt))
mask = twiss_full.index.str.contains(
reg) # pandas chose to call it 'contains'.
twiss_full.loc[mask, "dK1"] = (errdefs.loc[indx, "dK1"] *
twiss_full.loc[mask, "K1L"])**2
twiss_full.loc[mask, "dX"] = errdefs.loc[indx, "dX"] * 2
if errdefs.loc[indx, "MAINFIELD"] == "BPM":
twiss_full.loc[mask, "BPMdS"] = errdefs.loc[indx, "dS"]**2
else:
twiss_full.loc[mask, "KdS"] = (errdefs.loc[indx, "dS"] *
twiss_full.loc[mask, "K1L"])**2
twiss_full.loc[mask, "UNC"] = True
# in case of quadrupole longitudinal misalignments, the element (DRIFT) in front of the
# misaligned quadrupole will be used for the thin lens approximation of the misalignment
twiss_full["mKdS"] = np.roll(twiss_full.loc[:]["KdS"], 1)
twiss_full.loc[:, "UNC"] = np.logical_or(
abs(np.roll(twiss_full.loc[:, "dK1"], -1)) > 1.0e-12,
twiss_full.loc[:, "UNC"])
LOGGER.debug("DONE creating uncertainty information")
return twiss_full[twiss_full["UNC"]]
def get_ip_positions(path):
""" Returns a dict of IP positions from tfs-file of path.
Args:
path (str): Path to the tfs-file containing IP-positions
"""
df = tfs.read_tfs(path).set_index('NAME')
ip_names = ["IP" + str(i) for i in range(1, 9)]
ip_pos = df.loc[ip_names, 'S'].values
return dict(zip(ip_names, ip_pos))
def _collect_orbit_data(orbit_path_left, orbit_path_right,
kleft_path, kright_path):
k_file_left = tfs_pandas.read_tfs(kleft_path)
k_file_left["TIME"] = k_file_left["TIME"].astype(long)
k_file_left.set_index("TIME", inplace=True)
k_file_right = tfs_pandas.read_tfs(kright_path)
k_file_right["TIME"] = k_file_right["TIME"].astype(long)
k_file_right.set_index("TIME", inplace=True)
orbit_paths = {LEFT: orbit_path_left, RIGHT: orbit_path_right}
k_files = {LEFT: k_file_left, RIGHT: k_file_right}
k = {}
_apply_to_beam_side_plane(
lambda beam, side, plane: k.__setitem__((beam, side, plane), [])
)
bpm = {}
_apply_to_beam_side_plane(
lambda beam, side, plane: bpm.__setitem__((beam, side, plane), [])
)
for side in SIDES:
orbit_path = orbit_paths[side]
k_file = k_files[side]
for filename in os.listdir(orbit_path):
orbit_data = tfs_pandas.read_tfs(os.path.join(
orbit_path,
filename
))
timestamp = _date_str_to_timestamp(orbit_data.headers["OrbitDate"])
if filename.startswith('LHCB1'):
beam = BEAM1
elif filename.startswith('LHCB2'):
beam = BEAM2
for plane in (HOR, VER):
closest = np.argmin(np.abs(timestamp - k_file.index.values))
val = _get_sign(beam, plane) * k_file.K.iloc[closest]
k[(beam, side, plane)].append(val)
orbit = orbit_data.loc[:, PLANE_STR[plane]]
bpm[(beam, side, plane)].append(orbit - np.mean(orbit))
return k, bpm
def do_plot(data_tuple, labels, colorplot):
xdata, ydata, xdata_err, ydata_err = data_tuple
fig = plt.figure(figsize=(7, 4))
fig.patch.set_facecolor('white')
axt = fig.add_subplot(111)
odr_output = do_odr(data_tuple)
xi = np.linspace(0, 0.018, 3)
line = odr_output.beta[0] * xi
axt.plot(xi,
line,
linestyle='--',
color='k',
label='${:.4f}\, \pm\, {:.4f}$'.format(odr_output.beta[0],
odr_output.sd_beta[0]))
axt.errorbar(xdata,
ydata - odr_output.beta[1],
xerr=xdata_err,
yerr=ydata_err,
fmt='o',
color=colorplot,
label='IP5 Xing')
com = tfs_pandas.read_tfs(
'/media/jdilly/Storage/Projects/NOTE.18.MD1_Amplitude_Detuning/overleaf/results/2018_commissioning/B1_crossing_H_amp_det.dat'
)
xdata_com = com['2JXRES']
ydata_com = com['NATQX'] - com['BBQb1qx']
xdata_err_com = com['2JXSTDRES']
ydata_err_com = com['NATQXRMS']
data_tuple_com = xdata_com, ydata_com, xdata_err_com, ydata_err_com
odr_output_com = do_odr(data_tuple_com)
line_com = odr_output_com.beta[0] * xi
axt.errorbar(xdata_com,
ydata_com - odr_output_com.beta[1],
xerr=xdata_err_com,
yerr=ydata_err_com,
fmt='o',
color='c',
label='Crossing angles')
axt.plot(xi,
line_com,
linestyle='--',
color='m',
label='${:.4f}\, \pm\, {:.4f}$'.format(odr_output_com.beta[0],
odr_output_com.sd_beta[0]))
axt.set_xlim([0., 0.017])
axt.set_ylim([-0.0008, 0.0013])
axt.set_xlabel(labels[0], fontsize=14)
axt.set_ylabel(labels[1], fontsize=14)
fig.tight_layout()
plt.legend(fontsize=14, ncol=2, loc='lower right')
def getdiff(meas_path=None, beta_file_name="getbeta"):
""" Calculates the differences between measurement, corrected and uncorrected model.
After running madx and creating the model with (twiss_cor) and without (twiss_no)
corrections, run this functions to get tfs-files with the differences between measurements
and models.
Args:
meas_path (str): Path to the measurement folder.
Needs to contain twiss_cor.dat and twiss_no.dat.
"""
if meas_path is None:
meas_path = sys.argv[1]
LOG.debug("Started 'getdiff' for measurment dir '{:s}'".format(meas_path))
if not isdir(meas_path):
raise IOError("No valid measurement directory:" + meas_path)
corrected_model_path = join(meas_path, TWISS_CORRECTED)
uncorrected_model_path = join(meas_path, TWISS_NOT_CORRECTED)
meas = OpticsMeasurement(meas_path)
twiss_cor = read_tfs(corrected_model_path).set_index('NAME', drop=False)
twiss_no = read_tfs(uncorrected_model_path).set_index('NAME', drop=False)
coup_cor = TwissOptics(twiss_cor, quick_init=True).get_coupling(method='cmatrix')
coup_no = TwissOptics(twiss_no, quick_init=True).get_coupling(method='cmatrix')
model = pd.merge(twiss_cor, twiss_no, how='inner', on='NAME', suffixes=('_c', '_n'))
coupling_model = pd.merge(coup_cor, coup_no, how='inner', left_index=True, right_index=True,
suffixes=('_c', '_n'))
coupling_model['NAME'] = coupling_model.index.values
for plane in ['x', 'y']:
_write_betabeat_diff_file(meas_path, meas, model, plane, beta_file_name)
_write_phase_diff_file(meas_path, meas, model, plane)
_write_disp_diff_file(meas_path, meas, model, plane)
_write_closed_orbit_diff_file(meas_path, meas, model, plane)
_write_coupling_diff_file(meas_path, meas, coupling_model)
_write_norm_disp_diff_file(meas_path, meas, model)
_write_chromatic_coupling_files(meas_path, corrected_model_path)
_write_betastar_diff_file(meas_path, meas, twiss_cor, twiss_no)
LOG.debug("Finished 'getdiff'.")
def plot(self):
self._figure.clear()
self._axes_to_data = {}
axes_x = self._figure.add_subplot(2, 1, 1)
axes_x.set_title(self._matcher_model.get_name())
file_horizontal = tfs_pandas.read_tfs(os.path.join(
self._matcher_model.get_output_path(), "sbs",
"sbsphasext_" + str(self._matcher_model.get_label()) + ".out")
)
self._axes_to_data[axes_x] = file_horizontal
self._plot_match(axes_x, file_horizontal, "X")
file_vertical = tfs_pandas.read_tfs(os.path.join(
self._matcher_model.get_output_path(), "sbs",
"sbsphaseyt_" + str(self._matcher_model.get_label()) + ".out")
)
axes_y = self._figure.add_subplot(2, 1, 2)
self._axes_to_data[axes_y] = file_vertical
self._plot_match(axes_y, file_vertical, "Y")
self._figure.tight_layout()
self._figure.patch.set_visible(False)
self._figure.canvas.draw()
def compute_offset(model1_path, model2_path,
orbit_path_left, orbit_path_right,
kleft_path, kright_path,
ip):
ks, orbits = _collect_orbit_data(orbit_path_left, orbit_path_right,
kleft_path, kright_path)
model1 = tfs_pandas.read_tfs(model1_path)
model1.set_index("NAME", inplace=True)
model2 = tfs_pandas.read_tfs(model2_path)
model2.set_index("NAME", inplace=True)
bpm_names = _get_bpm_names(orbit_path_left, orbit_path_right)
models = {BEAM1: model1, BEAM2: model2}
results = _apply_to_beam_side_plane(
lambda beam, side, plane: _compute_and_clean(
ip, beam, side, plane, models, bpm_names, ks, orbits
)
)
return results
def _create_base_file(source_dir, source_file, meas, error, expect, outname):
""" Copy Measurement into a base-file. """
data = tfs_pandas.read_tfs(source_file)
if error == "":
new_data = data.loc[:, ["S", "NAME", meas]]
new_data.columns = ["S", "NAME", expect]
else:
new_data = data.loc[:, ["S", "NAME", meas, error]]
new_data.columns = ["S", "NAME", expect, error]
path_out = os.path.join(source_dir, outname + BASE_ID)
tfs_pandas.write_tfs(path_out, new_data)
return path_out
def read_tfs(self, filename):
"""Actually reads the TFS file from self.directory with filename.
This function can be ovewriten to use something instead of tfs_pandas
to load the files.
Arguments:
filename: The name of the file to load.
Returns:
A tfs_pandas instance of the requested file.
"""
tfs_data = tfs_pandas.read_tfs(os.path.join(self.directory, filename))
if "NAME" in tfs_data:
tfs_data = tfs_data.set_index("NAME", drop=False)
return tfs_data
def test_measure_optics(_create_input):
input_files, measure_input = _create_input
measure_optics.measure_optics(input_files, measure_input)
for f in _find_measurement_files(measure_input.outputdir):
if not f.startswith("get"):
continue
a = tfs_pandas.read_tfs(join(measure_input.outputdir, f))
cols = [column for column in a.columns.values if column.startswith('DELTA')]
for col in cols:
mask = outliers.get_filter_mask(a.loc[:, col].values, limit=0.1)
rms = stats.weighted_rms(a.loc[mask, col].values)
if col[5] in LIMITS.keys():
assert rms < LIMITS[col[5]], "\nFile: {:25} Column: {:15} RMS: {:.6f}".format(f, col, rms)
else:
assert rms < DEFAULT_LIMIT, "\nFile: {:25} Column: {:15} RMS: {:.6f}".format(f, col, rms)
print("\nFile: {:25} Column: {:15} RMS: {:.6f}".format(f, col, rms))
def remove_nan_from_files(list_of_files, replace=False):
""" Remove NAN-Entries from files in list_of_files.
If replace=False a new file with .dropna in it's name is created, otherwise the file is
overwritten.
"""
for filepath in list_of_files:
try:
df = tfs.read_tfs(filepath)
LOG.info("Read file {:s}".format(filepath))
except (IOError, tfs.TfsFormatError):
LOG.info("Skipped file {:s}".format(filepath))
else:
df = df.dropna(axis='index')
if not replace:
filepath += ".dropna"
tfs.write_tfs(filepath, df)
def segment_by_segment(accel_cls, options):
"""
TODO
"""
segments = _parse_segments(options.segments)
elements = _parse_elements(options.elements)
if not segments and not elements:
raise SbsDefinitionError("No segments or elements provided in the input.")
if _there_are_duplicated_names(segments, elements):
raise SbsDefinitionError("Duplicated names in segments and elements.")
model = tfs_pandas.read_tfs(options.model).set_index("NAME", drop=False)
meas = OpticsMeasurement(options.measurement)
elem_segments = [Segment.init_from_element(name) for name in elements]
for segment in elem_segments + segments:
propagable = run_for_segment(accel_cls, segment, model, meas,
options.optics, options.output)
write_beatings(segment, propagable, options.output)
def _assign_uncertainties(twiss_full, errordefspath):
'''
Adds uncertainty information to twiss_full.
:Sources of Errors:
dK1: quadrupolar field errors
dS: quadrupole longitudinal misalignments
dX: sextupole transverse misalignments
BPMdS: BPM longitudinal misalignments
'''
LOGGER.debug("Start creating uncertainty information")
errdefs = tfs_pandas.read_tfs(errordefspath)
# create new columns
twiss_full = twiss_full.assign(UNC=False, dK1=0, KdS=0, mKdS=0, dX=0, BPMdS=0)
# loop over uncertainty definitions, fill the respective columns, set UNC to true
for indx in errdefs.index:
patt = errdefs.loc[indx, "PATTERN"]
if patt.startswith("key:"):
LOGGER.debug("creating uncertainty information for {:s}".format(patt))
mask = patt.split(":")[1]
else:
reg = re.compile(patt)
LOGGER.debug("creating uncertainty information for RegEx {:s}".format(patt))
mask = twiss_full.index.str.contains(reg) # pandas chose to call it 'contains'.
twiss_full.loc[mask, "dK1"] = (errdefs.loc[indx, "dK1"] * twiss_full.loc[mask, "K1L"]) ** 2
twiss_full.loc[mask, "dX"] = errdefs.loc[indx, "dX"]*2
if errdefs.loc[indx, "MAINFIELD"] == "BPM":
twiss_full.loc[mask, "BPMdS"] = errdefs.loc[indx, "dS"]**2
else:
twiss_full.loc[mask, "KdS"] = (errdefs.loc[indx, "dS"] * twiss_full.loc[mask, "K1L"]) ** 2
twiss_full.loc[mask, "UNC"] = True
# in case of quadrupole longitudinal misalignments, the element (DRIFT) in front of the
# misaligned quadrupole will be used for the thin lens approximation of the misalignment
twiss_full["mKdS"] = np.roll(twiss_full.loc[:]["KdS"], 1)
twiss_full.loc[:, "UNC"] = np.logical_or(abs(np.roll(twiss_full.loc[:, "dK1"], -1)) > 1.0e-12,
twiss_full.loc[:, "UNC"])
LOGGER.debug("DONE creating uncertainty information")
return twiss_full[twiss_full["UNC"]]
def _calc_dp_over_p(main_input, bpm_data):
model_twiss = tfs.read_tfs(main_input.model)
model_twiss.set_index("NAME", inplace=True)
sequence = model_twiss.headers["SEQUENCE"].lower().replace("b1", "").replace("b2", "")
if sequence != "lhc":
return 0.0 # TODO: What do we do with other accels.
accel_cls = manager.get_accel_class(accel=sequence)
arc_bpms_mask = accel_cls.get_element_types_mask(bpm_data.index, types=["arc_bpm"])
arc_bpms_mask = np.array([bool(x) for x in arc_bpms_mask])
arc_bpm_data = bpm_data.loc[arc_bpms_mask]
# We need it in mm:
dispersions = model_twiss.loc[arc_bpm_data.index, "DX"] * 1e3
closed_orbits = np.mean(arc_bpm_data, axis=1)
numer = np.sum(dispersions * closed_orbits)
denom = np.sum(dispersions ** 2)
if denom == 0.:
raise ValueError("Cannot compute dpp probably no arc BPMs.")
dp_over_p = numer / denom
return dp_over_p
def merge_and_copy_kmod_output(beam, kmod_dir, res_dir, mod_path):
""" Merges the needed files into dataframes and saves them.
Args:
beam (str): Currently used beam.
kmod_dir (str): Parent folder of the kmod results
res_dir (str): Destination folder for the merged output
mod_path (str): Path to the model.
"""
LOG.debug("Merging and copying of kmod data started.")
pattern = re.compile(".*R[0-9]\." + beam.upper())
model_twiss = tfs_pandas.read_tfs(mod_path, index="NAME")
ip_dir_names = [d for _, dirs, _ in os.walk(kmod_dir) for d in dirs if pattern.match(d)]
# copy beta data
for plane in "xy":
up = plane.upper()
new_data = tfs_pandas.TfsDataFrame()
for ip_dir_name in ip_dir_names:
src = join(kmod_dir, ip_dir_name, get_beta_filename(plane))
data = _load_source_data(src, "NAME")
if data is not None:
in_model = data.index.isin(model_twiss.index)
new_data = new_data.append(data.loc[in_model, :])
# Todo: Let Kmod fill these columns in the future
new_data["S"] = model_twiss.loc[new_data.index, "S"]
new_data["BET{}MDL".format(up)] = model_twiss.loc[new_data.index, "BET{}".format(up)]
new_data = new_data.rename(columns={"BET{}STD".format(up): "ERRBET{}".format(up)})
dst = join(res_dir, get_beta_merged_filename(plane))
tfs_pandas.write_tfs(dst, new_data, save_index="NAME")
# copy beta* data
new_data = tfs_pandas.TfsDataFrame()
for ip_dir_name in ip_dir_names:
src = join(kmod_dir, ip_dir_name, get_beta_star_filename())
data = _load_source_data(src)
new_data = new_data.append(data)
dst = join(res_dir, get_beta_star_merged_filename())
tfs_pandas.write_tfs(dst, new_data)
def plot(self):
self._figure.clear()
self._axes_to_data = {}
label = self._matcher_model.get_label()
axes_f1001 = self._figure.add_subplot(2, 1, 1)
axes_f1010 = self._figure.add_subplot(2, 1, 2)
outpath = self._matcher_model.get_output_path()
file_coup = tfs_pandas.read_tfs(os.path.join(
outpath, "sbs",
"sbscouple_" + label + ".out")
)
self._axes_to_data[axes_f1001] = file_coup
self._axes_to_data[axes_f1010] = file_coup
self._plot_match(axes_f1001, file_coup, "f1001")
self._plot_match(axes_f1010, file_coup, "f1010")
self._figure.tight_layout()
self._figure.patch.set_visible(False)
self._figure.canvas.draw()
