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), }