def _find_ip_positions(self):
""" Returns IP positions from Dataframe
Load model into twiss_df first!
"""
tw = self.twiss_df
return tw.loc[regex_in(r"\AIP\d$", tw.index), 'S']
def _remove_nonnecessaries(self):
""" Removies unnecessary entries from model
Load model into twiss_df first!
"""
LOG.debug("Removing unnecessary entries:")
LOG.debug(" Entries total: {:d}".format(self.twiss_df.shape[0]))
with timeit(lambda t: LOG.debug(" Removed in {:f}s".format(t))):
cleaned = self.twiss_df.loc[
regex_in(r"\A(M|BPM)", self.twiss_df.index), :]
LOG.debug(" Entries left: {:d}".format(cleaned.shape[0]))
return cleaned
def get_rdts(self, rdt_names=None):
""" Return Resonance Driving Terms. """
if rdt_names:
if not isinstance(rdt_names, list):
rdt_names = [rdt_names]
rdt_names = [rdt.upper() for rdt in rdt_names]
new_rdts = [
rdt for rdt in rdt_names if rdt not in self._results_df
]
if len(new_rdts) > 0:
self.calc_rdts(new_rdts)
return self._results_df.loc[:, ["S"] + rdt_names]
else:
return self._results_df.loc[:,
regex_in(r'\A(S|F\d{4})$', self.
_results_df.columns)]
def _get_model_twiss(model_or_path):
""" Load model, but keep only BPMs and Magnets """
try:
model = tfs.read_tfs(model_or_path, index="NAME")
except TypeError:
LOG.debug("Received model as DataFrame")
model = model_or_path
else:
LOG.debug("Loaded Model from file '{:s}'".format(model_or_path))
# Remove not needed entries
LOG.debug("Removing non-necessary entries:")
LOG.debug(" Entries total: {:d}".format(model.shape[0]))
model = model.loc[regex_in(r"\A(M|BPM)", model.index), :]
LOG.debug(" Entries left: {:d}".format(model.shape[0]))
# make a copy to suppress "SettingWithCopyWarning"
model = model.copy()
# Add Dummy for Phase Calculations
model.loc[DUMMY_ID, ["S", "MUX", "MUY"]] = 0.0
return model
def plot_rdts(self, rdt_names=None, apply_fun=np.abs, combined=True):
""" Plot Resonance Driving Terms """
LOG.debug("Plotting Resonance Driving Terms")
rdts = self.get_rdts(rdt_names)
is_s = regex_in(r'\AS$', rdts.columns)
rdts = rdts.dropna()
rdts.loc[:, ~is_s] = rdts.loc[:, ~is_s].applymap(apply_fun)
pstyle.set_style(self._plot_options.style, self._plot_options.manual)
if combined:
ax = rdts.plot(x='S')
ax.set_title('Resonance Driving Terms')
pstyle.small_title(ax)
pstyle.set_name('Resonance Driving Terms', ax)
pstyle.set_yaxis_label(apply_fun.__name__, 'F_{{jklm}}', ax)
self._nice_axes(ax)
else:
for rdt in rdts.loc[:, ~is_s]:
ax = rdts.plot(x='S', y=rdt)
ax.set_title('Resonance Driving Term ' + rdt)
pstyle.small_title(ax)
pstyle.set_name('Resonance Driving Term ' + rdt, ax)
pstyle.set_yaxis_label(apply_fun.__name__, rdt, ax)
self._nice_axes(ax)
def _sort_element_types(self):
""" Sorts Elements by types
Load model into twiss_df first!
"""
LOG.debug("Sorting Elements into types:")
tw = self.twiss_df
with timeit(lambda t: LOG.debug(" Sorted in {:f}s".format(t))):
return {
'BPM': regex_in(r'\ABPM', tw.index),
'MB': regex_in(r'\AMB[^S]', tw.index),
'MBS': regex_in(r'\AMBS', tw.index),
'MQ': regex_in(r'\AMQ[^ST]', tw.index),
'MQS': regex_in(r'\AMQS', tw.index),
'MQT': regex_in(r'\AMQT', tw.index),
'MS': regex_in(r'\AMS[^S]', tw.index),
'MSS': regex_in(r'\AMSS', tw.index),
'MO': regex_in(r'\AMO[^S]', tw.index),
'MOS': regex_in(r'\AMOS', tw.index),
'MCB': regex_in(r'\AMCB', tw.index),
'MCS': regex_in(r'\AMCS', tw.index),
'MCO': regex_in(r'\AMCO', tw.index),
'MCD': regex_in(r'\AMCD', tw.index),
}