def _calculate_dispersion(model, input_files, plane, header, unit, cut, output, order=1):
df_orbit = pd.DataFrame(model).loc[:, ['S', 'MU' + plane, 'DP' + plane, 'D' + plane, plane]]
df_orbit.rename(columns={'MU' + plane: 'MU' + plane + 'MDL', 'DP' + plane: 'DP' + plane + 'MDL',
'D' + plane: 'D' + plane + 'MDL', plane: plane + 'MDL'}, inplace=True)
df_orbit = pd.merge(df_orbit, input_files.joined_frame(plane, ['CO', 'CORMS']), how='inner',
left_index=True, right_index=True)
df_orbit['COUNT'] = len(input_files.get_columns(df_orbit, 'CO'))
dpps = input_files.dpps(plane)
if np.max(dpps) - np.min(dpps) == 0.0:
return # temporary solution
# raise ValueError('Cannot calculate dispersion, only a single momentum data')
fit = np.polyfit(dpps, SCALES[unit] * input_files.get_data(df_orbit, 'CO').T, order, cov=True)
# in the fit results the coefficients are sorted by power in decreasing order
df_orbit['D' + plane] = fit[0][-2, :].T
df_orbit['STDD' + plane] = np.sqrt(fit[1][-2, -2, :].T)
df_orbit[plane] = fit[0][-1, :].T
df_orbit['STD' + plane] = np.sqrt(fit[1][-1, -1, :].T)
# since we get variances from the fit, maybe we can include the variances of fitted points
df_orbit = df_orbit.loc[np.abs(df_orbit.loc[:, plane]) < cut*SCALES[unit], :]
df_orbit['DP' + plane] = _calculate_dp(model,
df_orbit.loc[:, ['D' + plane, 'STDD' + plane]], plane)
df_orbit['DELTAD' + plane] = df_orbit.loc[:, 'D'+ plane] - df_orbit.loc[:, 'D' + plane + 'MDL']
output_df = df_orbit.loc[
:, ['S', 'COUNT', 'D' + plane, 'STDD' + plane, plane, 'STD' + plane, 'DP' + plane,
'D' + plane + 'MDL', 'DP' + plane + 'MDL', 'MU' + plane + 'MDL', 'DELTAD' + plane]]
tfs_pandas.write_tfs(join(output, header['FILENAME']), output_df, header, save_index='NAME')
return output_df
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 extract_knob_value_and_definition(knob_names, time, cwd="./", server=None):
""" Extract knob values from online model using the knob name and the time.
Example Call::
extract_knob_values("LHCBEAM/2018_global_ats_flat_b1_for_ip5_waist",
"2018-10-30 15:00:00.0",
cwd="/afs/cern.ch/user/j/jdilly/extractor/",
server="cs-ccr-dev3")
Args:
knob_names (list): List of knob names to extract
time: UTC time in ISOformat to extract
cwd: output directory for results and log (default: current directory)
server: server to run on (default: runs local)
"""
knob_file = os.path.join(cwd, const.get_extractor_knobs_filename())
# first get the trim value and madx_changefile
_run_knob_extraction(knob_names, time, cwd, server, "trim")
if not os.path.exists(knob_file):
raise IOError(
"Something went wrong while extracting data form online model. See log file."
)
trim, _ = dc.post_trim_extract(knob_file)
# run again and get the values and deltas
_run_knob_extraction(knob_names, time, cwd, server, "k")
df_list = dc.knobs_k_to_tfs(knob_file, trim)
for knob, df in zip(knob_names, df_list):
filename = const.get_knob_tfs_filename(knob)
tfs.write_tfs(os.path.join(cwd, filename), df)
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 _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 _write_coupling_diff_file(meas_path, meas, model):
LOG.debug("Calculating coupling diff.")
tw = pd.merge(meas.coupling,
model,
how='inner',
left_index=True,
right_index=True)
out_columns = ['NAME', 'S']
for idx, rdt in enumerate(['F1001', 'F1010']):
tw[rdt + 're'] = tw.loc[:, rdt + 'R']
tw[rdt + 'im'] = tw.loc[:, rdt + 'I']
tw[rdt + 'e'] = tw.loc[:, 'FWSTD{:d}'.format(idx + 1)]
tw[rdt + 're_m'] = np.real(tw.loc[:, rdt + '_c'])
tw[rdt + 'im_m'] = np.imag(tw.loc[:, rdt + '_c'])
tw[rdt +
're_prediction'] = tw.loc[:, rdt + 're'] - tw.loc[:, rdt + 're_m']
tw[rdt +
'im_prediction'] = tw.loc[:, rdt + 'im'] - tw.loc[:, rdt + 'im_m']
tw[rdt + 'W_prediction'] = np.sqrt(
np.square(tw[rdt + 're_prediction']) +
np.square(tw[rdt + 'im_prediction']))
out_columns += [
rdt + 're', rdt + 'im', rdt + 'e', rdt + 're_m', rdt + 'im_m',
rdt + 'W', rdt + 'W_prediction', rdt + 're_prediction',
rdt + 'im_prediction'
]
tw['in_use'] = 1
out_columns += ['in_use']
write_tfs(join(meas_path, get_diff_filename('couple')),
tw.loc[:, out_columns])
def extract_knob_value_and_definition(knob_names, time, cwd="./", server=None):
""" Extract knob values from online model using the knob name and the time.
Example Call::
extract_knob_values("LHCBEAM/2018_global_ats_flat_b1_for_ip5_waist",
"2018-10-30 15:00:00.0",
cwd="/afs/cern.ch/user/j/jdilly/extractor/",
server="cs-ccr-dev3")
Args:
knob_names (list): List of knob names to extract
time: UTC time in ISOformat to extract
cwd: output directory for results and log (default: current directory)
server: server to run on (default: runs local)
"""
knob_file = os.path.join(cwd, const.get_extractor_knobs_filename())
# first get the trim value and madx_changefile
_run_knob_extraction(knob_names, time, cwd, server, "trim")
if not os.path.exists(knob_file):
raise IOError("Something went wrong while extracting data form online model. See log file.")
trim, _ = dc.post_trim_extract(knob_file)
# run again and get the values and deltas
_run_knob_extraction(knob_names, time, cwd, server, "k")
df_list = dc.knobs_k_to_tfs(knob_file, trim)
for knob, df in zip(knob_names, df_list):
filename = const.get_knob_tfs_filename(knob)
tfs.write_tfs(os.path.join(cwd, filename), df)
def _write_betabeat_diff_file(meas_path, meas, model, plane, betafile):
LOG.debug("Calculating beta diff.")
if betafile == "getbeta":
meas_beta = meas.beta[plane]
elif betafile == "getampbeta":
meas_beta = meas.amp_beta[plane]
elif betafile == "getkmodbeta":
meas_beta = meas.kmod_beta[plane]
else:
raise KeyError("Unknown beta file name '{}'.".format(betafile))
up = plane.upper()
tw = pd.merge(meas_beta,
model,
how='inner',
left_index=True,
right_index=True)
tw['MEA'] = ((tw.loc[:, 'BET' + up] - tw.loc[:, 'BET' + up + 'MDL']) /
tw.loc[:, 'BET' + up + 'MDL'])
tw['ERROR'] = tw.loc[:, 'ERRBET' + up] / tw.loc[:, 'BET' + up + 'MDL']
tw['MODEL'] = (
(tw.loc[:, 'BET' + up + '_c'] - tw.loc[:, 'BET' + up + '_n']) /
tw.loc[:, 'BET' + up + '_n'])
tw['EXPECT'] = tw['MEA'] - tw['MODEL']
write_tfs(join(meas_path, get_diff_filename('bb' + plane)),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
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 _calculate_normalised_dispersion(model, input_files, beta, header, unit, cut, output, accelerator):
#TODO there are no errors from orbit
df_orbit = pd.DataFrame(model).loc[:, ['S', 'MUX', 'DPX', 'DX', 'X', 'BETX']]
df_orbit['NDXMDL'] = df_orbit.loc[:, 'DX'] / np.sqrt(df_orbit.loc[:, 'BETX'])
df_orbit.rename(columns={'MUX': 'MUXMDL', 'DPX': 'DPXMDL', 'DX': 'DXMDL', 'X': 'XMDL'}, inplace=True)
df_orbit['COUNT'] = len(input_files.get_columns(df_orbit, 'CO'))
dpps = input_files.dpps("X")
df_orbit = pd.merge(df_orbit, input_files.joined_frame("X", ['CO', 'CORMS', 'AMPX']),
how='inner', left_index=True, right_index=True)
df_orbit = pd.merge(df_orbit, beta.loc[:, ['BETX', 'ERRBETX']], how='inner', left_index=True,
right_index=True, suffixes=('', '_phase'))
if np.max(dpps) - np.min(dpps) == 0.0:
return # temporary solution
# raise ValueError('Cannot calculate dispersion, only a single dpoverp')
fit = np.polyfit(dpps, SCALES[unit] * input_files.get_data(df_orbit, 'CO').T, 1, cov=True)
df_orbit['NDX_unscaled'] = fit[0][-2, :].T / stats.weighted_mean(input_files.get_data(df_orbit, 'AMPX'), axis=1) # TODO there is no error from AMPX
df_orbit['STDNDX_unscaled'] = np.sqrt(fit[1][-2, -2, :].T) / stats.weighted_mean(input_files.get_data(df_orbit, 'AMPX'), axis=1)
df_orbit = df_orbit.loc[np.abs(fit[0][-1, :].T) < cut * SCALES[unit], :]
mask = accelerator.get_element_types_mask(df_orbit.index, ["arc_bpm"])
global_factor = np.sum(df_orbit.loc[mask, 'NDXMDL'].values) / np.sum(df_orbit.loc[mask, 'NDX_unscaled'].values)
df_orbit['NDX'] = global_factor * df_orbit.loc[:, 'NDX_unscaled']
df_orbit['STDNDX'] = global_factor * df_orbit.loc[:, 'STDNDX_unscaled']
df_orbit['DX'] = df_orbit.loc[:, 'NDX'] * np.sqrt(df_orbit.loc[:, 'BETX_phase'])
df_orbit['STDDX'] = df_orbit.loc[:, 'STDNDX'] * np.sqrt(df_orbit.loc[:, 'BETX_phase'])
df_orbit['DPX'] = _calculate_dp(model, df_orbit.loc[:, ['DX', 'STDDX']], "X")
df_orbit['DELTANDX'] = df_orbit.loc[:, 'NDX'] - df_orbit.loc[:, 'NDXMDL']
output_df = df_orbit.loc[:, ['S', 'COUNT', 'NDX', 'STDNDX', 'DX', 'DPX',
'NDXMDL', 'DXMDL', 'DPXMDL', 'MUXMDL', 'DELTANDX']]
tfs_pandas.write_tfs(join(output, header['FILENAME']), output_df, header, save_index='NAME')
return output_df
def write_harpy_output(main_input, harpy_data_frame, headers, spectrum, plane):
output_file = get_outpath_with_suffix(
main_input.file, main_input.outputdir, ".lin" + plane
)
tfs_pandas.write_tfs(output_file, harpy_data_frame, headers)
if not main_input.skip_files:
_write_full_spectrum(main_input, spectrum, plane)
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 write_bad_bpms(first_file, bad_bpms_to_write):
meas_dir = os.path.abspath(os.path.join(first_file, os.pardir))
for plane in PLANE:
bad_bpms_summary_path = os.path.join(
meas_dir, "bad_bpms_iforest_{}.tfs".format(plane))
tfs_pandas.write_tfs(bad_bpms_summary_path, bad_bpms_to_write[plane])
LOGGER.info("Bad BPMs summary from Isolation Forest written to: %s",
meas_dir)
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_table(results, output):
rows = [[BEAM_STR[beam], SIDE_STR[side],
results[(beam, side, HOR)][0], results[(beam, side, HOR)][1],
results[(beam, side, VER)][0], results[(beam, side, VER)][1]]
for beam in BEAMS for side in SIDES]
data = pd.DataFrame(data=rows,
columns=("BEAM", "SIDE",
"OFFSETX", "ERROFFSETX",
"OFFSETY", "ERROFFSETY"))
tfs_pandas.write_tfs(output, data)
def _write_table(results, output):
rows = [[
BEAM_STR[beam], SIDE_STR[side], results[(beam, side, HOR)][0],
results[(beam, side, HOR)][1], results[(beam, side, VER)][0],
results[(beam, side, VER)][1]
] for beam in BEAMS for side in SIDES]
data = pd.DataFrame(data=rows,
columns=("BEAM", "SIDE", "OFFSETX", "ERROFFSETX",
"OFFSETY", "ERROFFSETY"))
tfs_pandas.write_tfs(output, data)
def extract_overview(knob_names,
time=None,
cwd="./",
server=None,
show_plot=False):
""" Extract overview-data consisting of
- beamprocess
- optics used
- values of given Knobs
- orbit plot
Args:
knob_names (list): List of knob names to extract
time: UTC time in ISO format (default: now)
cwd: output directory for results and log (default: current directory)
server: server to run on (default: runs local)
"""
if time is None:
time = dc.get_utc_now()
if knob_names is None or len(knob_names) == 0:
raise NotImplementedError(
"Knob names need to be provided, due to bug in extractor.")
# knob_names = extract_all_knob_names(time=time, only_active=False, cwd=cwd, server=server)
# extraction
_run_overview_extraction(knob_names, time, cwd, server)
# knobs
knobs_file = os.path.join(cwd, const.get_extractor_knobs_filename())
if not os.path.exists(knobs_file):
raise RuntimeError("Knobs output file '{:s}' ".format(knobs_file) +
"not found!")
df = dc.knobs_kvalues_to_tfs(knobs_file)
# optics
optics_file = os.path.join(cwd, const.get_extractor_output_filename())
if not os.path.exists(optics_file):
raise RuntimeError("Optics output file '{:s}' ".format(optics_file) +
"not found!")
df.headers[const.get_optics_header()] = dc.get_optics(optics_file)
# orbit
filenames = [
os.path.join(cwd, const.get_default_orbit_filename(b)) for b in [1, 2]
]
if not all([os.path.exists(fn) for fn in filenames]):
raise RuntimeError("Orbit files not found in '{:s}'.".format(cwd))
df.headers[const.get_fill_header()] = dc.get_fill_from_orbitfile(
filenames[0])
tfs.write_tfs(os.path.join(cwd, const.get_overview_filename()), df)
_log_df(df)
_plot_orbits(cwd, show_plot)
def _write_full_spectrum(main_input, spectrum, plane):
spectr_amps_files = get_outpath_with_suffix(
main_input.file, main_input.outputdir, ".amps" + plane
)
amps_df = spectrum["COEFS"].abs().T
tfs_pandas.write_tfs(spectr_amps_files, amps_df)
spectr_freqs_files = get_outpath_with_suffix(
main_input.file, main_input.outputdir, ".freqs" + plane
)
freqs_df = spectrum["FREQS"].T
tfs_pandas.write_tfs(spectr_freqs_files, freqs_df)
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 _write_full_spectrum(main_input, spectrum, plane):
spectr_amps_files = get_outpath_with_suffix(main_input.file,
main_input.outputdir,
".amps" + plane)
amps_df = spectrum["COEFS"].abs().T
tfs_pandas.write_tfs(spectr_amps_files, amps_df)
spectr_freqs_files = get_outpath_with_suffix(main_input.file,
main_input.outputdir,
".freqs" + plane)
freqs_df = spectrum["FREQS"].T
tfs_pandas.write_tfs(spectr_freqs_files, freqs_df)
def filter_by_ip(ip, file_in, file_out):
""" Filter the Error table by IP
Args:
ip (str): IPs to keep, e.g. "1" or "15" (filters IP1 and IP5)
file_in: path to input error def file
file_out: path to output file
"""
df = tfs.read_tfs(file_in, index="NAME")
mask = df.index.str.match(r".*[LR][{:s}](\.V[12])?".format(ip), case=False)
tfs.write_tfs(file_out, df.loc[mask, :], save_index="NAME")
def _calculate_orbit(model, input_files, plane, header, output):
df_orbit = pd.DataFrame(model).loc[:, ['S', 'MU' + plane, plane]]
df_orbit.rename(columns={'MU' + plane: 'MU' + plane + 'MDL', plane: plane + 'MDL'}, inplace=True)
df_orbit = pd.merge(df_orbit, input_files.joined_frame(plane, ['CO', 'CORMS']), how='inner',
left_index=True, right_index=True)
df_orbit['COUNT'] = len(input_files.get_columns(df_orbit, 'CO'))
df_orbit[plane] = stats.weighted_mean(input_files.get_data(df_orbit, 'CO'), axis=1)
df_orbit['STD' + plane] = stats.weighted_error(input_files.get_data(df_orbit, 'CO'), axis=1)
df_orbit['DELTA' + plane] = df_orbit.loc[:, plane] - df_orbit.loc[:, plane + 'MDL']
output_df = df_orbit.loc[:, ['S', 'COUNT', plane, 'STD' + plane, plane + 'MDL', 'MU' + plane + 'MDL', 'DELTA' + plane]]
tfs_pandas.write_tfs(join(output, header['FILENAME']), output_df, header, save_index='NAME')
return output_df
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 _write_phase_diff_file(meas_path, meas, model, plane):
LOG.debug("Calculating phase diff.")
up = plane.upper()
tw = pd.merge(meas.phase[plane], model, how='inner', on='NAME')
tw['MEA'] = tw.loc[:, 'PHASE' + up]
tw['ERROR'] = tw.loc[:, 'STDPH' + up]
tw['MODEL'] = np.concatenate((np.diff(tw.loc[:, 'MU' + up + '_c']), np.array([0.0])))
tw['DIFF'] = tw.loc[:, 'PHASE' + up] - tw.loc[:, 'PH' + up + 'MDL']
tw['DIFF_MDL'] = tw.loc[:, 'MODEL'] - tw.loc[:, 'PH' + up + 'MDL']
tw['EXPECT'] = tw['DIFF'] - tw['DIFF_MDL']
write_tfs(join(meas_path, get_diff_filename('phase' + plane)),
tw.loc[tw.index[:-1],
['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'DIFF', 'DIFF_MDL', 'EXPECT']])
def _write_closed_orbit_diff_file(meas_path, meas, model, plane):
LOG.debug("Calculating orbit diff.")
up = plane.upper()
try:
tw = pd.merge(meas.orbit[plane], model, how='inner', on='NAME')
except IOError:
LOG.debug("Orbit measurements not found. Skipped.")
else:
tw['MEA'] = tw.loc[:, up] - tw.loc[:, up + 'MDL']
tw['ERROR'] = tw.loc[:, 'STD' + up]
tw['MODEL'] = (tw.loc[:, up + '_c'] - tw.loc[:, up + '_n']) * 1000
tw['EXPECT'] = tw['MEA'] - tw['MODEL'] * 1000
write_tfs(join(meas_path, get_diff_filename('co' + plane)),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
def _write_norm_disp_diff_file(meas_path, meas, model):
LOG.debug("Calculating normalized dispersion diff.")
try:
tw = pd.merge(meas.norm_disp, model, how='inner', on='NAME')
except IOError:
LOG.debug("Normalized dispersion measurements not found. Skipped.")
else:
tw['MEA'] = tw.loc[:, 'NDX'] - tw.loc[:, 'NDXMDL']
tw['ERROR'] = tw.loc[:, 'STDNDX']
tw['MODEL'] = (tw.loc[:, 'DX_c'] / np.sqrt(tw.loc[:, 'BETX_c']) -
tw.loc[:, 'DX_n'] / np.sqrt(tw.loc[:, 'BETX_n']))
tw['EXPECT'] = tw['MEA'] - tw['MODEL']
write_tfs(join(meas_path, get_diff_filename('ndx')),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
def _write_norm_disp_diff_file(meas_path, meas, model):
LOG.debug("Calculating normalized dispersion diff.")
try:
tw = pd.merge(meas.norm_disp, model, how='inner', on='NAME')
except IOError:
LOG.debug("Normalized dispersion measurements not found. Skipped.")
else:
tw['MEA'] = tw.loc[:, 'NDX'] - tw.loc[:, 'NDXMDL']
tw['ERROR'] = tw.loc[:, 'STDNDX']
tw['MODEL'] = (tw.loc[:, 'DX_c'] / np.sqrt(tw.loc[:, 'BETX_c'])
- tw.loc[:, 'DX_n'] / np.sqrt(tw.loc[:, 'BETX_n']))
tw['EXPECT'] = tw['MEA'] - tw['MODEL']
write_tfs(join(meas_path, get_diff_filename('ndx')),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
def _write_disp_diff_file(meas_path, meas, model, plane):
LOG.debug("Calculating dispersion diff.")
up = plane.upper()
try:
tw = pd.merge(meas.disp[plane], model, how='inner', on='NAME')
except IOError:
LOG.debug("Dispersion measurements not found. Skipped.")
else:
tw['MEA'] = tw.loc[:, 'D' + up] - tw.loc[:, 'D' + up + 'MDL']
tw['ERROR'] = tw.loc[:, 'STDD' + up]
tw['MODEL'] = tw.loc[:, 'D' + up + '_c'] - tw.loc[:, 'D' + up + '_n']
tw['EXPECT'] = tw['MEA'] - tw['MODEL']
write_tfs(join(meas_path, get_diff_filename('d' + plane)),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
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 _write_disp_diff_file(meas_path, meas, model, plane):
LOG.debug("Calculating dispersion diff.")
up = plane.upper()
try:
tw = pd.merge(meas.disp[plane], model, how='inner', on='NAME')
except IOError:
LOG.debug("Dispersion measurements not found. Skipped.")
else:
tw['MEA'] = tw.loc[:, 'D' + up] - tw.loc[:, 'D' + up + 'MDL']
tw['ERROR'] = tw.loc[:, 'STDD' + up]
tw['MODEL'] = tw.loc[:, 'D' + up + '_c'] - tw.loc[:, 'D' + up + '_n']
tw['EXPECT'] = tw['MEA'] - tw['MODEL']
write_tfs(join(meas_path, get_diff_filename('d' + plane)),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
def _write_closed_orbit_diff_file(meas_path, meas, model, plane):
LOG.debug("Calculating orbit diff.")
up = plane.upper()
try:
tw = pd.merge(meas.orbit[plane], model, how='inner', on='NAME')
except IOError:
LOG.debug("Orbit measurements not found. Skipped.")
else:
tw['MEA'] = tw.loc[:, up] - tw.loc[:, up + 'MDL']
tw['ERROR'] = tw.loc[:, 'STD' + up]
tw['MODEL'] = (tw.loc[:, up + '_c'] - tw.loc[:, up + '_n']) * 1000
tw['EXPECT'] = tw['MEA'] - tw['MODEL'] * 1000
write_tfs(join(meas_path, get_diff_filename('co' + plane)),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
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 _write_phase_diff_file(meas_path, meas, model, plane):
LOG.debug("Calculating phase diff.")
up = plane.upper()
tw = pd.merge(meas.phase[plane], model, how='inner', on='NAME')
tw['MEA'] = tw.loc[:, 'PHASE' + up]
tw['ERROR'] = tw.loc[:, 'STDPH' + up]
tw['MODEL'] = np.concatenate(
(np.diff(tw.loc[:, 'MU' + up + '_c']), np.array([0.0])))
tw['DIFF'] = tw.loc[:, 'PHASE' + up] - tw.loc[:, 'PH' + up + 'MDL']
tw['DIFF_MDL'] = tw.loc[:, 'MODEL'] - tw.loc[:, 'PH' + up + 'MDL']
tw['EXPECT'] = tw['DIFF'] - tw['DIFF_MDL']
write_tfs(
join(meas_path, get_diff_filename('phase' + plane)),
tw.loc[tw.index[:-1], [
'NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'DIFF', 'DIFF_MDL', 'EXPECT'
]])
def _gather_data(cwd, machine, beam, xing, error_types, error_loc, optic_types,
seeds, unused_stages):
""" Gather the data and write it to a file """
for optic_type in optic_types:
# ampdet data
for output_id in _get_all_output_ids(machine, beam, unused_stages):
seed_data = tfs.TfsDataFrame(
index=seeds,
columns=AMPDET_NAMES,
)
for seed in seeds:
output_dir = tripcor.get_output_dir(
tripcor.get_seed_dir(cwd, seed), xing, error_types,
error_loc, optic_type)
# define seed folder and error definition paths
seed_data.loc[seed, :] = get_values_from_tfs(
get_tfs_name(output_dir, machine, beam, optic_type,
output_id))
title, filename = get_seed_data_title_and_filename(
beam, xing, error_types, error_loc, optic_type, output_id)
LOG.info("Gathered data for '{:s}'".format(title))
seed_data.headers["Title"] = title
seed_data = seed_data.astype(np.float64)
tfs.write_tfs(os.path.join(get_data_output_folder(cwd), filename),
seed_data,
save_index="SEED")
# cta data
seed_data_cta = tfs.TfsDataFrame(index=seeds, columns=["QX", "QY"])
for seed in seeds:
output_dir = tripcor.get_output_dir(
tripcor.get_seed_dir(cwd, seed), xing, error_types, error_loc,
optic_type)
seed_data_cta.loc[seed, :] = get_cta_values(output_dir, beam)
title, filename = get_cta_seed_data_title_and_filename(
beam, xing, error_types, error_loc, optic_type)
LOG.info("Gathered cta data for '{:s}'".format(title))
seed_data_cta.headers["Title"] = title
seed_data_cta = seed_data_cta.astype(np.float64)
tfs.write_tfs(os.path.join(get_data_output_folder(cwd), filename),
seed_data_cta,
save_index="SEED")
def extract_overview(knob_names, time=None, cwd="./", server=None, show_plot=False):
""" Extract overview-data consisting of
- beamprocess
- optics used
- values of given Knobs
- orbit plot
Args:
knob_names (list): List of knob names to extract
time: UTC time in ISO format (default: now)
cwd: output directory for results and log (default: current directory)
server: server to run on (default: runs local)
"""
if time is None:
time = dc.get_utc_now()
if knob_names is None or len(knob_names) == 0:
raise NotImplementedError("Knob names need to be provided, due to bug in extractor.")
# knob_names = extract_all_knob_names(time=time, only_active=False, cwd=cwd, server=server)
# extraction
_run_overview_extraction(knob_names, time, cwd, server)
# knobs
knobs_file = os.path.join(cwd, const.get_extractor_knobs_filename())
if not os.path.exists(knobs_file):
raise RuntimeError("Knobs output file '{:s}' ".format(knobs_file) +
"not found!")
df = dc.knobs_kvalues_to_tfs(knobs_file)
# optics
optics_file = os.path.join(cwd, const.get_extractor_output_filename())
if not os.path.exists(optics_file):
raise RuntimeError("Optics output file '{:s}' ".format(optics_file) +
"not found!")
df.headers[const.get_optics_header()] = dc.get_optics(optics_file)
# orbit
filenames = [os.path.join(cwd, const.get_default_orbit_filename(b)) for b in [1, 2]]
if not all([os.path.exists(fn) for fn in filenames]):
raise RuntimeError("Orbit files not found in '{:s}'.".format(cwd))
df.headers[const.get_fill_header()] = dc.get_fill_from_orbitfile(filenames[0])
tfs.write_tfs(os.path.join(cwd, const.get_overview_filename()), df)
_log_df(df)
_plot_orbits(cwd, show_plot)
def revert_forest_cleaning(files):
"""
Reverts the cleaning. The backup files are renamed back to the original names (e.g .linx.notcleaned --> .linx)
:param paths: list of files, where bad bpms identified by iForest are removed
"""
files_list = files.split(',')
for path in files_list:
src_dir = os.path.abspath(os.path.join(path, os.pardir))
filename = os.path.basename(path)
notcleaned_file = os.path.join(src_dir, filename + ".notcleaned")
original_file_tfs = tfs_pandas.read_tfs(notcleaned_file).set_index(
"NAME", drop=False)
os.remove(path)
lin_file = os.path.join(src_dir,
notcleaned_file.replace(".notcleaned", ""))
os.rename(notcleaned_file, lin_file)
tfs_pandas.write_tfs(lin_file, original_file_tfs)
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 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 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 get_ndx(folder1, folder2):
file_name = "/getNDx.out"
result = get_merged_measurement(folder1, folder2, file_name)
meas = 'NDX'
model = 'NDXMDL'
err = "STDNDX"
result.loc[:, meas] = (result.loc[:, meas] + result.loc[:, model] -
result.loc[:, meas + '_t'])
model_diff = np.std(result.loc[:, model].values -
result.loc[:, model + '_t'].values)
if model_diff:
print("Models are different")
print(model_diff)
result.loc[:, err] = np.sqrt(
np.square(result.loc[:, err].values) +
np.square(result.loc[:, err + '_t'].values))
write_tfs(rfold + file_name, result, {})
return result
def write_betastar_from_phase(ips_d, headers, output_dir):
""" TODO
"""
assert not any(
tfs_pandas.write_tfs(
join(output_dir, MODE_TO_SUFFIX[mode].format(plane.lower())),
ips_d[plane][mode],
headers_dict=headers,
) for plane in PLANES for mode in ips_d[plane])
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 get_beta_from_phase(folder1, folder2, plane='x', free=True):
file_name = "/getbeta" + plane + free * "_free" + ".out"
result = get_merged_measurement(folder1, folder2, file_name)
meas = 'BET' + plane.upper()
model = 'BET' + plane.upper() + 'MDL'
err = "ERRBET" + plane.upper()
result.loc[:, meas] = (result.loc[:, meas] + result.loc[:, model] -
result.loc[:, meas + '_t'])
model_diff = np.std(result.loc[:, model].values -
result.loc[:, model + '_t'].values)
if model_diff:
print("Models are different")
print(model_diff)
result.loc[:, err] = np.sqrt(
np.square(result.loc[:, err].values) +
np.square(result.loc[:, err + '_t'].values))
file_name = "/getbeta" + plane + ".out"
write_tfs(rfold + file_name, result, {})
return result
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 write_betastar_from_phase(ips_d, headers, output_dir):
""" TODO
"""
assert not any(
tfs_pandas.write_tfs(
join(output_dir, MODE_TO_SUFFIX[mode].format(plane.lower())),
ips_d[plane][mode],
headers_dict=headers,
)
for plane in PLANES for mode in ips_d[plane]
)
def _write_betabeat_diff_file(meas_path, meas, model, plane, betafile):
LOG.debug("Calculating beta diff.")
if betafile == "getbeta":
meas_beta = meas.beta[plane]
elif betafile == "getampbeta":
meas_beta = meas.amp_beta[plane]
elif betafile == "getkmodbeta":
meas_beta = meas.kmod_beta[plane]
else:
raise KeyError("Unknown beta file name '{}'.".format(betafile))
up = plane.upper()
tw = pd.merge(meas_beta, model, how='inner', on='NAME')
tw['MEA'] = ((tw.loc[:, 'BET' + up] - tw.loc[:, 'BET' + up + 'MDL'])
/ tw.loc[:, 'BET' + up + 'MDL'])
tw['ERROR'] = tw.loc[:, 'ERRBET' + up] / tw.loc[:, 'BET' + up + 'MDL']
tw['MODEL'] = ((tw.loc[:, 'BET' + up + '_c'] - tw.loc[:, 'BET' + up + '_n'])
/ tw.loc[:, 'BET' + up + '_n'])
tw['EXPECT'] = tw['MEA'] - tw['MODEL']
write_tfs(join(meas_path, get_diff_filename('bb' + plane)),
tw.loc[:, ['NAME', 'S', 'MEA', 'ERROR', 'MODEL', 'EXPECT']])
def get_phase(folder1, folder2, plane='x', free=True):
file_name = "/getphase" + plane + free * "_free" + ".out"
res, header = get_merged_measurement_with_header(folder1, folder2,
file_name)
mask = res.loc[:, 'NAME2'] == res.loc[:, 'NAME2_t']
result = res.loc[res.index[mask], :]
meas = 'PHASE' + plane.upper()
model = 'PH' + plane.upper() + 'MDL'
err = "STDPH" + plane.upper()
result.loc[:, meas] = (result.loc[:, meas] + result.loc[:, model] -
result.loc[:, meas + '_t'])
model_diff = np.std(result.loc[:, model].values -
result.loc[:, model + '_t'].values)
if model_diff:
print("Models are different")
print(model_diff)
result.loc[:, err] = np.sqrt(
np.square(result.loc[:, err].values) +
np.square(result.loc[:, err + '_t'].values))
file_name = "/getphase" + plane + ".out"
write_tfs(rfold + file_name, result, header)
return result
def _build_changeparameters_file(input_data):
original_file =\
os.path.join(input_data.match_path, "changeparameters.madx")
results_dirname = "results_" + datetime.datetime.now().strftime(
"%Y-%m-%d_%H:%M:%S")
output_dir = os.path.join(input_data.match_path, results_dirname)
os.mkdir(output_dir)
vars_dict = OrderedDict()
with open(original_file, "r") as original_file_data:
for original_line in original_file_data:
parts = original_line.split("=")
variable_name = parts[0].replace("d", "", 1).strip()
variable_value = float(parts[1].replace(";", "").strip())
vars_dict[variable_name] = variable_value
tfs_pandas.write_tfs(
os.path.join(output_dir, "changeparameters.tfs"),
pd.DataFrame(data={
"NAME": vars_dict.keys(),
"DELTA": vars_dict.values()
}).loc[:, ["NAME", "DELTA"]],
headers_dict={
"DATE": datetime.datetime.now().strftime("%a %b %d %H:%M:%S %Y")
},
)
changeparameters_correct =\
os.path.join(output_dir, "changeparameters_correct.madx")
changeparameters_match =\
os.path.join(output_dir, "changeparameters.madx")
with open(changeparameters_correct, "w") as correct_data,\
open(changeparameters_match, "w") as match_data:
for varname in vars_dict:
value = vars_dict[varname]
sign = "+" if value >= 0 else "-"
sign_correct = "-" if value >= 0 else "+" # Flip sign to correct
correct_data.write("{name} = {name} {sign} {value};\n".format(
name=varname, sign=sign_correct, value=abs(value)))
match_data.write("{name} = {name} {sign} {value};\n".format(
name=varname, sign=sign, value=abs(value)))
def calculate_kick(measure_input, input_files, model, mad_ac, beta_d, header_dict):
"""
Fills the following TfsFiles:
- getkick.out getkickac.out
Args:
measure_input: Optics_input object
input_files: Stores the input files Tfs_pandas
model: Model tfs panda
mad_ac: Model tfs panda with AC-dipole in
beta_d: measured beta functions
header_dict: OrderedDict containing information about the analysis
Returns:
"""
try:
tunes_actions = _getkick(input_files, _get_model_arc_betas(measure_input, model), ac=False)
except IndexError: # occurs if either no x or no y files exist
return pd.DataFrame, pd.DataFrame
column_names = ["DPP", "QX", "QXRMS", "QY", "QYRMS", "NATQX", "NATQXRMS", "NATQY", "NATQYRMS",
"sqrt2JX", "sqrt2JXSTD", "sqrt2JY", "sqrt2JYSTD", "2JX", "2JXSTD", "2JY",
"2JYSTD"]
kick_frame = pd.DataFrame(data=tunes_actions, columns=column_names)
header = _get_header(header_dict, beta_d)
tfs_pandas.write_tfs(join(measure_input.outputdir, header['FILENAME']), kick_frame, header)
actions_x, actions_y = tunes_actions[:, 9:11], tunes_actions[:, 11:13] # sqrt2jx, sqrt2Jy
if measure_input.accelerator.excitation != AccExcitationMode.FREE:
column_names_ac = column_names + ["sqrt2JXRES", "sqrt2JXSTDRES", "sqrt2JYRES", "sqrt2JYSTDRES", "2JXRES",
"2JXSTDRES", "2JYRES", "2JYSTDRES"]
tunes_actions_ac = _getkick(input_files, _get_model_arc_betas(measure_input, mad_ac), ac=True)
x, y = beta_d["X"], beta_d["Y"]
tunes_actions_ac_res = tunes_actions_ac[:, 9:] / np.array([np.sqrt(x), np.sqrt(x), np.sqrt(y), np.sqrt(y), x, x, y, y])
kick_frame_ac = pd.DataFrame(data=np.hstack((tunes_actions_ac, tunes_actions_ac_res)), columns=column_names_ac)
header_ac = _get_header(header_dict, beta_d, ac=True)
tfs_pandas.write_tfs(join(measure_input.outputdir, header_ac['FILENAME']), kick_frame_ac, header_ac)
actions_x, actions_y = tunes_actions_ac[:, 9:11], tunes_actions_ac[:, 11:13]
return actions_x, actions_y
def _write_coupling_diff_file(meas_path, meas, model):
LOG.debug("Calculating coupling diff.")
tw = pd.merge(meas.coupling, model, how='inner', on='NAME')
out_columns = ['NAME', 'S']
for idx, rdt in enumerate(['F1001', 'F1010']):
tw[rdt+'re'] = tw.loc[:, rdt+'R']
tw[rdt+'im'] = tw.loc[:, rdt+'I']
tw[rdt+'e'] = tw.loc[:, 'FWSTD{:d}'.format(idx+1)]
tw[rdt+'re_m'] = np.real(tw.loc[:, rdt+'_c'])
tw[rdt+'im_m'] = np.imag(tw.loc[:, rdt+'_c'])
tw[rdt+'re_prediction'] = tw.loc[:, rdt+'re'] - tw.loc[:, rdt+'re_m']
tw[rdt+'im_prediction'] = tw.loc[:, rdt+'im'] - tw.loc[:, rdt+'im_m']
tw[rdt+'W_prediction'] = np.sqrt(np.square(tw[rdt+'re_prediction'])
+ np.square(tw[rdt+'im_prediction']))
out_columns += [rdt+'re', rdt+'im', rdt+'e',
rdt+'re_m', rdt+'im_m',
rdt+'W', rdt+'W_prediction',
rdt+'re_prediction', rdt+'im_prediction']
tw['in_use'] = 1
out_columns += ['in_use']
write_tfs(join(meas_path, get_diff_filename('couple')), tw.loc[:, out_columns])
def remove_bpms_from_file(paths, bad_bpm_names, plane):
"""
Writes a backup of the original .lin files (e.g .linx --> .linx.notcleaned)
and removes the BPNs identified by iForest as bad.
:param paths: original lin files
:param bad_bpm_names: list of the names of bad BPMs identified by iForest
"""
for path in paths:
src_dir = os.path.abspath(os.path.join(path, os.pardir))
filename = os.path.basename(path)
new_filename = os.path.join(src_dir, filename + ".notcleaned")
os.rename(path, new_filename)
original_file_tfs = tfs_pandas.read_tfs(new_filename).set_index(
"NAME", drop=False)
original_file_tfs = original_file_tfs.loc[~original_file_tfs.index.
isin(bad_bpm_names)]
pln_num = "1" if plane == "x" else "2"
original_file_tfs.headers["Q{}".format(pln_num)] =\
original_file_tfs["TUNE" + plane.upper()].mean()
original_file_tfs.headers["Q{}RMS".format(pln_num)] =\
np.std(original_file_tfs["TUNE" + plane.upper()])
tfs_pandas.write_tfs(path, original_file_tfs,
original_file_tfs.headers)
def _build_changeparameters_file(input_data):
original_file =\
os.path.join(input_data.match_path, "changeparameters.madx")
output_dir = os.path.join(input_data.match_path, "results")
os.mkdir(output_dir)
vars_dict = OrderedDict()
with open(original_file, "r") as original_file_data:
for original_line in original_file_data:
parts = original_line.split("=")
variable_name = parts[0].replace("d", "", 1).strip()
variable_value = float(parts[1].replace(";", "").strip())
vars_dict[variable_name] = variable_value
tfs_pandas.write_tfs(
os.path.join(output_dir, "changeparameters.tfs"),
pd.DataFrame(data={"NAME": vars_dict.keys(),
"DELTA": vars_dict.values()}).loc[:, ["NAME", "DELTA"]],
headers_dict={"DATE": datetime.datetime.now().strftime("%a %b %d %H:%M:%S %Y")},
)
changeparameters_correct =\
os.path.join(output_dir, "changeparameters_correct.madx")
changeparameters_match =\
os.path.join(output_dir, "changeparameters.madx")
with open(changeparameters_correct, "w") as correct_data,\
open(changeparameters_match, "w") as match_data:
for varname in vars_dict:
value = vars_dict[varname]
sign = "+" if value >= 0 else "-"
sign_correct = "-" if value >= 0 else "+" # Flip sign to correct
correct_data.write(
"{name} = {name} {sign} {value};\n"
.format(name=varname, sign=sign_correct, value=abs(value))
)
match_data.write(
"{name} = {name} {sign} {value};\n"
.format(name=varname, sign=sign, value=abs(value))
)
def get_diff_two_dataframes(df1, df2, diff_columns, is_error=None, prefix="", index=None, keep_colums=(), out_file=None):
""" Get the difference of common elements of specific columns between two dataframes.
Merges on index.
Args:
df1 (DataFrame or Path): First dataframe, Minuend
df2 (DataFrame or Path): Second dataframe, Subtrahend
diff_columns (list of stings): List of columns to get the difference of
is_error (list of booleans): defines if the column in question is an error column
prefix (str): Prefix for difference columns (default: "")
index (str): index column - most likely needed when reading/writing files
keep_colums (list of strings): additional columns to keep in the returned dataframe
out_file (Path): if given, writes the result into this file
Returns:
DataFrame containing difference columns and kept columns.
"""
# convert from files to dataframes
df1 = _get_dataframe(df1, index)
df2 = _get_dataframe(df2, index)
# check input
_check_for_missing_columns(df1, df2, diff_columns)
_check_for_missing_columns(df1, df2, keep_colums)
# merge dataframes
merged = pd.merge(df1, df2, how='inner',
left_index=True, right_index=True,
suffixes=('_df1', '_df2'))
# calculate difference
if is_error is None:
is_error = [False] * len(diff_columns)
else:
if len(is_error) != len(diff_columns):
raise ValueError(
"The length of the is_error switch needs to correspond to the columns."
)
for idx, col in enumerate(diff_columns):
if is_error[idx]:
merged[prefix + col] = .5 * np.sqrt(np.square(merged[col + '_df1'])
+ np.square(merged[col + '_df2']))
else:
merged[prefix + col] = merged[col + '_df1'] - merged[col + '_df2']
# copy columns to be kept
for col in keep_colums:
for suffix in ["", "_df1", "_df2"]:
try:
merged[col] = merged[col + suffix]
except KeyError:
pass
else:
break
# prepare output
merged = merged.loc[:, keep_colums + [prefix + c for c in diff_columns]]
if out_file:
tfs_pandas.write_tfs(out_file, merged, save_index=index)
return merged
def _write_tfs(self, filename, data_frame):
if self.allow_write:
tfs_pandas.write_tfs(os.path.join(self.directory, filename), data_frame)
self._buffer[filename] = data_frame
def writting_calibration_factors_to_file(name,position,beta_calibration,beta_calibration_error,file_output):
tfs_calibration_pandas = file_output
all_summary = pd.DataFrame(OrderedDict([('NAME', name),('S', position),('CALIBRATION',beta_calibration),('ERROR_CALIBRATION',beta_calibration_error)]))
tfs_pandas.write_tfs(tfs_calibration_pandas,all_summary,save_index=False)
def _write_betastar_diff_file(meas_path, meas, twiss_cor, twiss_no):
LOG.debug("Calculating betastar diff at the IPs.")
try:
meas = meas.kmod_betastar.set_index(bsft.RES_COLUMNS[0])
except IOError:
LOG.debug("Beta* measurements not found. Skipped.")
else:
# get all IPs
ip_map = {}
beam = ''
for label in meas.index.values:
ip, beam = re.findall(r'\d', label)[-2:] # beam should be the same for all
if ip not in "1258":
raise NotImplementedError(
"Beta-Star comparison is not yet implemented for measurements in IP" + ip)
ip_label = "IP" + ip
ip_map[label] = ip_label
beam = int(beam)
all_ips = set(ip_map.values())
try:
# calculate waist and so on
model = bsft.get_beta_star_and_waist_from_ip(twiss_cor, beam, all_ips)
design = bsft.get_beta_star_and_waist_from_ip(twiss_no, beam, all_ips)
except KeyError:
LOG.warn("Can't find all IPs in twiss files. Skipped beta* calculations.")
else:
# extract data
tw = pd.DataFrame()
for label in meas.index:
plane = label[-1]
ip_name = bsft.get_full_label(ip_map[label], beam, plane)
tw.loc[label, "S"] = model.loc[ip_name, "S"]
# calculate alpha* but with s-oriented waist definition
meas["ALPHASTAR"] = meas["WAIST"] / meas["BETAWAIST"]
meas["ALPHASTAR_ERR"] = ((meas["WAIST_ERR"] / meas["WAIST"] +
meas["BETAWAIST_ERR"] / meas["BETAWAIST"]) *
meas["ALPHASTAR"]
)
for attr in bsft.RES_COLUMNS[2:]:
# default diff parameter
tw.loc[label, attr + "_MEA"] = (meas.loc[label, attr]
- design.loc[ip_name, attr])
tw.loc[label, attr + "_ERROR"] = meas.loc[label, attr + "_ERR"]
tw.loc[label, attr + "_MODEL"] = (model.loc[ip_name, attr]
- design.loc[ip_name, attr])
# additional for checks (e.g. for betastar* panel)
tw.loc[label, attr + "_MEAVAL"] = meas.loc[label, attr]
tw.loc[label, attr + "_DESIGNVAL"] = design.loc[ip_name, attr]
tw.loc[label, attr + "_MODELVAL"] = model.loc[ip_name, attr]
# and the beatings
tw.loc[label, "B{}_MEA".format(attr)] = (tw.loc[label, attr + "_MEA"]
/ design.loc[ip_name, attr])
tw.loc[label, "B{}_MODEL".format(attr)] = (tw.loc[label, attr + "_MODEL"]
/ design.loc[ip_name, attr])
# special handling for the expectation values, as waist and betawaist
# should be derived directly from alpha* and beta*
if attr in bsft.RES_COLUMNS[2:4]:
# beta* and alpha*: as usual
tw.loc[label, attr + "_EXPECT"] = (tw.loc[label, attr + "_MEA"]
- tw.loc[label, attr + "_MODEL"])
tw.loc[label, attr + "_EXPECTVAL"] = (design.loc[ip_name, attr]
+ tw.loc[label, attr + "_EXPECT"])
tw.loc[label, "B{}_EXPECT".format(attr)] = (
tw.loc[label, "B{}_MEA".format(attr)]
- tw.loc[label, "B{}_MODEL".format(attr)])
else:
# waist and betawaist: calculate expected value directly and go from there
tw.loc[label, attr + "_EXPECTVAL"] = (
bsft.get_waist_wrapper(attr,
tw.loc[label, "BETASTAR_EXPECTVAL"],
tw.loc[label, "ALPHASTAR_EXPECTVAL"],
)
)
tw.loc[label, attr + "_EXPECT"] = (
tw.loc[label, attr + "_EXPECTVAL"] - design.loc[ip_name, attr])
tw.loc[label, "B{}_EXPECT".format(attr)] = (
tw.loc[label, attr + "_EXPECTVAL"] / design.loc[ip_name, attr])
write_tfs(join(meas_path, get_diff_filename('betastar')), tw,
save_index=bsft.RES_COLUMNS[0])
