def _take_next(phases: tfs.TfsDataFrame, shift: int = 1): """ Takes the following BPM for momentum reconstruction by a given shift. Args: phases (tfs.TfsDataFrame): Dataframe matrix of phase advances, as calculated in phase.py. shift (int): Value to determine the BPM pairing. If ``0`` is given, tries to find the best candidate. If a value ``n>=1`` is given, then takes the n-th following BPM downstream for the pairing. """ indices = np.roll(np.arange(phases.to_numpy().shape[0]), shift) return indices, phases.to_numpy()[np.arange(phases.to_numpy().shape[0]), indices] - 0.25
def _find_candidate(phases: tfs.TfsDataFrame) -> Tuple[np.ndarray, np.ndarray]: """ Finds the best candidate for momentum reconstruction. Args: phases (tfs.TfsDataFrame): Dataframe matrix of phase advances, as calculated in phase.py. bd (int): beam direction, will be negative for beam 2. Returns: The indices of best candidates, and the corresponding phase advances between these indices. """ slice_ = _tilt_slice_matrix( phases.to_numpy(), 0, 2 * CUTOFF) - 0.25 # do not overwrite built-in 'slice' indices = np.argmin(abs(slice_), axis=0) deltas = slice_[indices, range(len(indices))] indices = (indices + np.arange(len(indices))) % len(indices) return np.array(indices), deltas