ENH:add local orbit distortion filter function to sofb device

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):
         """."""

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