diff --git a/siriuspy/siriuspy/devices/sofb.py b/siriuspy/siriuspy/devices/sofb.py index ffec4e613..1bdb7b38c 100644 --- a/siriuspy/siriuspy/devices/sofb.py +++ b/siriuspy/siriuspy/devices/sofb.py @@ -6,6 +6,8 @@ from ..sofb.csdev import SOFBFactory from ..sofb.utils import si_calculate_bump as _si_calculate_bump +from ..sofb.utils import filter_local_distortions as _filter_local_distortions +from ..sofb.utils import find_closest_corrs_ss as _find_closest_corrs_ss from .device import Device as _Device @@ -1155,6 +1157,18 @@ def wait_drive(self, timeout=None): 'DriveState-Sts', self._data.DriveState.Open, timeout=timeout ) + def filter_local_distortions(self, corr_idcs, dorb=None, respmat=None): + """Filters local distortions. + + Returns orbit distortion due to local perturbation. + """ + respmat = respmat or self.respmat + if dorb is None: + orb = _np.concatenate((self.orbx, self.orby)) + dorb = orb - _np.concatenate((self.refx, self.refy)) + dorb_filtered = _filter_local_distortions(corr_idcs, dorb, respmat) + return dorb_filtered + @staticmethod def si_calculate_bumps( orbx, @@ -1180,6 +1194,11 @@ def si_calculate_bumps( minsingval=minsingval, ) + def find_closest_corrs_ssection(self, ss, nr_ch, nr_cv, inc_rf=False): + """Find the closest correctors indices to a given sector indices in SOFB matrix.""" + idcs_corrs = _find_closest_corrs_ss(ss, nr_ch, nr_cv, inc_rf=inc_rf) + return idcs_corrs + @property def trajx(self): """.""" diff --git a/siriuspy/siriuspy/sofb/utils.py b/siriuspy/siriuspy/sofb/utils.py index f70cb7d62..764e9fe83 100644 --- a/siriuspy/siriuspy/sofb/utils.py +++ b/siriuspy/siriuspy/sofb/utils.py @@ -74,3 +74,51 @@ def si_calculate_bump( orby[bpm1] += pos_bpm[2] orby[bpm2] += pos_bpm[3] return orbx, orby + + +def find_closest_corrs_ss(ss, nr_ch, nr_cv, inc_rf=False): + """Find the closest correctors indices to a given sector indices in SOFB matrix.""" + nch = 6 + ncv = 8 + idx_upstream = (ss-1)*nch + idx_downstream = (ss-1)*nch + 1 + idcs_ch_up = idx_upstream - _np.arange(int(nr_ch/2)) + idcs_ch_dw = idx_downstream + _np.arange(int(nr_ch/2)) + idcs_ch = _np.sort(_np.concatenate((idcs_ch_up, idcs_ch_dw))) + idcs_ch = _np.mod(idcs_ch, 120) + + idx_upstream = (ss-1)*ncv + idx_downstream = (ss-1)*ncv + 1 + idcs_cv_up = idx_upstream - _np.arange(int(nr_cv/2)) + idcs_cv_dw = idx_downstream + _np.arange(int(nr_cv/2)) + idcs_cv = _np.sort(_np.concatenate((idcs_cv_up, idcs_cv_dw))) + idcs_cv = _np.mod(idcs_cv, 160) + 120 + + idcs_corrs = _np.concatenate((idcs_ch, idcs_cv)) + if inc_rf: + idcs_corrs = _np.r_[idcs_corrs, 280] + return idcs_corrs + + +def filter_local_distortions(corr_idcs, dorb=None, respmat=None): + """Filter local orbit distortions. + + Args: + corr_idcs (numpy 1D array or list): Indices of correctors to use. + dorb (2D or 1D numpy array, optional): array containing orbit + distortions x and y, in the case of a 2D array, + the second dimension must have length 320, + ie. (npts, 320). Defaults to None. + respmat (2D numpy array, optional): Orbit response matrix. + Defaults to None. + + Returns: + (2D or 1D numpy array): Returns the orbit distortion filtered + in an array with the same chape as the dorb array, ie (npts, 320). + """ + kicks = _np.linalg.lstsq(respmat, dorb.T, rcond=None)[0] + sel = _np.zeros(respmat.shape[1], dtype=bool) + sel[corr_idcs] = True + kicks[~sel] = 0 + dorb_filt = respmat @ kicks + return dorb_filt.T