Merge pull request #283 from lnls-fac/add-qn-selection-tune-corr

add idcs_out argument
lnls-fac · Oct 31, 2024 · 313dbdb · 313dbdb
2 parents d049268 + 39877b2
commit 313dbdb
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git a/apsuite/optics_analysis/base_correction.py b/apsuite/optics_analysis/base_correction.py
@@ -27,7 +27,7 @@ class BaseCorr():
     GROUPING = _get_namedtuple('Grouping', ['Individual', 'TwoKnobs'])
     CORR_STATUS = _get_namedtuple('CorrStatus', ['Fail', 'Sucess'])
 
-    def __init__(self, model, acc=None, method=None, grouping=None):
+    def __init__(self, model, acc=None, method=None, grouping=None, idcs_out=None):
         """."""
         self.model = model
         self.acc = acc.upper()
@@ -37,6 +37,7 @@ def __init__(self, model, acc=None, method=None, grouping=None):
         self.fam = None
         self.method = method
         self.grouping = grouping
+        self.idcs_out = idcs_out
 
     @property
     def grouping(self):
@@ -96,8 +97,9 @@ def _get_strength(self, model=None, knobs=None):
             for mag in self.fam[knb]['index']:
                 stren_seg = []
                 for seg in mag:
-                    stren_seg.append(getattr(
-                        model[seg], self._STRENGTH_TYPE))
+                    if self.idcs_out is None or seg not in self.idcs_out:
+                        stren_seg.append(getattr(
+                            model[seg], self._STRENGTH_TYPE))
                 stren_mag.append(sum(stren_seg))
             stren.append(_np.mean(stren_mag))
         return _np.array(stren)
@@ -168,12 +170,13 @@ def _add_delta_stren(self, delta_stren, model=None):
                 delta = delta_stren[idx_knb]
             for mag in self.fam[knb]['index']:
                 for seg in mag:
-                    stren = getattr(model[seg], self._STRENGTH_TYPE)
-                    if self._method == BaseCorr.METHODS.Proportional:
-                        stren *= (1 + delta/len(mag))
-                    else:
-                        stren += delta/len(mag)
-                    setattr(model[seg], self._STRENGTH_TYPE, stren)
+                    if self.idcs_out is None or seg not in self.idcs_out:
+                        stren = getattr(model[seg], self._STRENGTH_TYPE)
+                        if self._method == BaseCorr.METHODS.Proportional:
+                            stren *= (1 + delta/len(mag))
+                        else:
+                            stren += delta/len(mag)
+                        setattr(model[seg], self._STRENGTH_TYPE, stren)
 
     def _group_2knobs_matrix(self, jacobian_matrix=None):
         """."""

diff --git a/apsuite/optics_analysis/tune_correction.py b/apsuite/optics_analysis/tune_correction.py
@@ -20,12 +20,14 @@ class TuneCorr(BaseCorr):
     OPTICS = _get_namedtuple('Optics', ['EdwardsTeng', 'Twiss'])
 
     def __init__(self, model, acc, qf_knobs=None, qd_knobs=None,
-                 method=None, grouping=None, type_optics=None):
+                 method=None, grouping=None, type_optics=None, idcs_out=None):
         """."""
         super().__init__(
-            model=model, acc=acc, method=method, grouping=grouping)
+            model=model, acc=acc, method=method, grouping=grouping,
+            idcs_out=idcs_out)
         self._type_optics = TuneCorr.OPTICS.EdwardsTeng
         self.type_optics = type_optics
+        self.idcs_out = idcs_out
         if self.type_optics == self.OPTICS.EdwardsTeng:
             self._optics_func = pyaccel.optics.calc_edwards_teng
         elif self._type_optics == self.OPTICS.Twiss:
@@ -100,7 +102,8 @@ def calc_jacobian_matrix(self, model=None):
             modcopy = model[:]
             for nmag in self.fam[knb]['index']:
                 for seg in nmag:
-                    modcopy[seg].KL += delta/len(nmag)
+                    if self.idcs_out is None or seg not in self.idcs_out:
+                        modcopy[seg].KL += delta/len(nmag)
             nu = self.get_tunes(model=modcopy)
             tune_matrix[:, idx] = (nu-nu0)/delta
         return tune_matrix