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 add_moving_average(kickac_df, bbq_df, filter_args):
""" Adds the moving average of the bbq data to kickac_df and bbq_df. """
LOG.debug("Calculating moving average.")
for idx, plane in enumerate(PLANES):
if filter_args.bbq_filtering_method == 'outliers':
bbq_mav, bbq_std, mask = bbq_tools.clean_outliers_moving_average(
bbq_df[get_bbq_col(plane)],
length=filter_args.window_length,
limit=filter_args.outlier_limit)
else:
bbq_mav, bbq_std, mask = bbq_tools.get_moving_average(
bbq_df[get_bbq_col(plane)],
length=filter_args.window_length,
min_val=filter_args.tunes_minmax[2 * idx],
max_val=filter_args.tunes_minmax[2 * idx + 1],
fine_length=filter_args.fine_window,
fine_cut=filter_args.fine_cut,
)
bbq_df[get_mav_col(plane)] = bbq_mav
bbq_df[get_mav_err_col(plane)] = bbq_std
bbq_df[get_used_in_mav_col(plane)] = mask
kickac_df = add_bbq_data(kickac_df, bbq_mav, get_mav_col(plane))
kickac_df = add_bbq_data(kickac_df, bbq_std, get_mav_err_col(plane))
return kickac_df, bbq_df
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
def add_corrected_natural_tunes(kickac_df):
""" Adds the corrected natural tunes to kickac
Args:
kickac_df: Dataframe containing the data
Returns:
Modified kick_ac
"""
for plane in PLANES:
kickac_df[get_natq_corr_col(plane)] = (kickac_df[get_natq_col(plane)] -
kickac_df[get_mav_col(plane)])
return kickac_df
def _plot_bbq_data(bbq_df,
interval=None,
x_lim=None,
y_lim=None,
two_plots=False):
"""
Plot BBQ data.
Args:
bbq_df: BBQ Dataframe with moving average columns.
interval: start and end time of used interval, will be marked with red bars.
x_lim: x limits (time).
y_lim: y limits (tune).
output: Path to the output file.
show: Shows plot if ``True``.
two_plots: Plots each tune in it's own axes if ``True``.
Returns:
Plotted figure.
"""
LOG.debug("Plotting BBQ data.")
pstyle.set_style(
"standard", {
u'figure.figsize': [12.24, 7.68],
u"lines.marker": u"",
u"lines.linestyle": u""
})
fig, axs = plt.subplots(1 + two_plots, 1)
if not two_plots:
axs = [axs, axs]
handles = [None] * (3 * len(PLANES))
for idx, plane in enumerate(PLANES):
color = pcolors.get_mpl_color(idx)
mask = np.array(bbq_df[get_used_in_mav_col(plane)], dtype=bool)
# plot and save handles for nicer legend
handles[idx] = axs[idx].plot(
[i.datetime for i in bbq_df.index],
bbq_df[get_bbq_col(plane)],
color=pcolors.change_color_brightness(color, .4),
marker="o",
markerfacecolor="None",
label="$Q_{:s}$".format(plane.lower(), ))[0]
filtered_data = bbq_df.loc[mask, get_bbq_col(plane)].dropna()
handles[len(PLANES) + idx] = axs[idx].plot(
filtered_data.index,
filtered_data.values,
color=pcolors.change_color_brightness(color, .7),
marker=".",
label="filtered".format(plane.lower()))[0]
handles[2 * len(PLANES) + idx] = axs[idx].plot(
bbq_df.index,
bbq_df[get_mav_col(plane)],
color=color,
linestyle="-",
label="moving av.".format(plane.lower()))[0]
if (y_lim is None or y_lim[0] is None) and two_plots:
axs[idx].set_ylim(bottom=min(bbq_df.loc[mask, get_bbq_col(plane)]))
if (y_lim is None or y_lim[1] is None) and two_plots:
axs[idx].set_ylim(top=max(bbq_df.loc[mask, get_bbq_col(plane)]))
# things to add/do only once if there is only one plot
for idx in range(1 + two_plots):
if interval:
axs[idx].axvline(x=interval[0], color="red")
axs[idx].axvline(x=interval[1], color="red")
if two_plots:
axs[idx].set_ylabel("$Q_{:s}$".format(PLANES[idx]))
else:
axs[idx].set_ylabel('Tune')
if y_lim is not None:
axs[idx].set_ylim(y_lim)
axs[idx].yaxis.set_major_formatter(FormatStrFormatter('%.5f'))
if x_lim is not None:
axs[idx].set_xlim(x_lim)
axs[idx].set_xlabel('Time')
axs[idx].xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))
if idx:
# don't show labels on upper plot (if two plots)
# use the visibility to allow cursor x-position to be shown
axs[idx].tick_params(labelbottom=False)
axs[idx].xaxis.get_label().set_visible(False)
if not two_plots or idx:
# reorder legend
axs[idx].legend(
handles,
[h.get_label() for h in handles],
loc='lower right',
bbox_to_anchor=(1.0, 1.01),
ncol=3,
)
fig.tight_layout()
fig.tight_layout()
return fig