def _get_bbq_data(beam, input_, kick_df):
"""Return BBQ data from input, either file or timber fill."""
try:
fill_number = int(input_)
except ValueError:
# input is a string
if input_ == "kick":
LOG.debug("Getting timber data from kick-times.")
timber_keys, bbq_cols = _get_timber_keys_and_bbq_columns(beam)
t_start = min(kick_df.index.values)
t_end = max(kick_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:
LOG.debug(f"Getting bbq data from file '{input_:s}'")
data = read_timed_dataframe(input_)
data.drop([get_mav_col(p) for p in PLANES
if get_mav_col(p) in data.columns], axis='columns')
else:
# input is a number
LOG.debug(f"Getting timber data from fill '{input_:d}'")
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)))
return data
def main(opt):
LOG.info("Plotting Amplitude Detuning Results.")
_save_options(opt)
_check_opt(opt)
kick_plane = opt.plane
figs = {}
_set_plotstyle(opt.manual_style)
limits = opt.get_subdict(['x_lim', 'y_lim'])
for tune_plane in PLANES:
for corrected in [False, True]:
corr_label = "_corrected" if corrected else ""
acd_corr = 1
if opt.correct_acd and (kick_plane == tune_plane):
acd_corr = 0.5
fig = plt.figure()
ax = fig.add_subplot(111)
for idx, (kick, label) in enumerate(zip(opt.kicks, opt.labels)):
kick_df = kick_mod.read_timed_dataframe(kick) if isinstance(kick, str) else kick
try:
data_df = kick_mod.get_ampdet_data(kick_df, kick_plane, tune_plane, corrected=corrected)
except KeyError:
continue # should only happen when there is no 'corrected' columns
odr_fit = kick_mod.get_odr_data(kick_df, kick_plane, tune_plane,
order=opt.detuning_order, corrected=corrected)
data_df, odr_fit = _correct_and_scale(data_df, odr_fit,
opt.action_unit, opt.action_plot_unit,
10**opt.tune_scale, acd_corr)
_plot_detuning(ax, data=data_df, label=label, limits=limits,
odr_fit=odr_fit,
color=pcolors.get_mpl_color(idx),
action_unit=opt.action_plot_unit, tune_scale=10**opt.tune_scale)
ax_labels = const.get_paired_lables(tune_plane, kick_plane, opt.tune_scale)
id_str = f"dQ{tune_plane.upper():s}d2J{kick_plane.upper():s}{corr_label:s}"
pannot.set_name(id_str, fig)
_format_axes(ax, labels=ax_labels, limits=limits)
if opt.show:
plt.show()
if opt.output:
output = Path(opt.output)
fig.savefig(f'{output.with_suffix("")}_{id_str}{output.suffix}')
figs[id_str] = fig
return figs
def main(opt):
""" Plot BBQ wrapper.
"""
LOG.info("Plotting BBQ.")
_save_options(opt)
bbq_df = kick_mod.read_timed_dataframe(opt.input) if isinstance(
opt.input, str) else opt.input
opt.pop("input")
if opt.kick is not None:
if opt.interval is not None:
raise ValueError(
"interval and kick-file given. Both are used for the same purpose. Please only use one."
)
window = 0 # not too important, bars will then indicate first and last kick directly
for p in PLANES:
with suppress(KeyError):
window = bbq_df.headers[get_mav_window_header(p)]
kick_df = kick_mod.read_timed_dataframe(opt.kick) if isinstance(
opt.kick, str) else opt.kick
opt.interval = ad_ana.get_approx_bbq_interval(bbq_df, kick_df.index,
window)
bbq_df = bbq_df.loc[opt.interval[0]:opt.interval[1]]
opt.pop("kick")
show = opt.pop("show")
out = opt.pop("output")
fig = _plot_bbq_data(bbq_df, **opt)
if show:
plt.show()
if out:
fig.savefig(out)
return fig
def _get_bbq_data(beam: int, input_: str,
kick_df: TfsDataFrame) -> TfsDataFrame:
"""
Return BBQ data from input, either file or timber fill, as a ``TfsDataFrame``.
Note: the ``input_`` parameter is always parsed from the commandline as a string, but could be 'kick'
or a kickfile name or an integer. All these options will be tried until one works.
"""
try:
fill_number = int(input_)
except ValueError: # input_ is a file name or the string 'kick'
if input_ == "kick":
LOG.debug("Getting timber data from kick times")
timber_keys, bbq_cols = _get_timber_keys_and_bbq_columns(beam)
t_start = min(kick_df.index.to_numpy())
t_end = max(kick_df.index.to_numpy())
t_delta = timedelta(seconds=DTIME)
data = timber_extract.extract_between_times(
t_start - t_delta,
t_end + t_delta,
keys=timber_keys,
names=dict(zip(timber_keys, bbq_cols)))
else: # input_ is a file name
LOG.debug(f"Getting bbq data from file '{input_:s}'")
data = read_timed_dataframe(input_)
# Drop old moving average columns as these will be computed again later
data.drop([
get_mav_col(p)
for p in PLANES if get_mav_col(p) in data.columns
],
axis="columns")
else: # input_ is a number, assumed to be a fill number
LOG.debug(f"Getting timber data from fill '{input_:d}'")
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)))
return data