From 39877b2300ca2c1004a60145b5d33438f5cfac63 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Fri, 25 Oct 2024 09:01:43 -0300 Subject: [PATCH] add idcs_out argument --- apsuite/optics_analysis/base_correction.py | 21 ++++++++++++--------- apsuite/optics_analysis/tune_correction.py | 9 ++++++--- 2 files changed, 18 insertions(+), 12 deletions(-) mode change 100644 => 100755 apsuite/optics_analysis/base_correction.py mode change 100644 => 100755 apsuite/optics_analysis/tune_correction.py diff --git a/apsuite/optics_analysis/base_correction.py b/apsuite/optics_analysis/base_correction.py old mode 100644 new mode 100755 index c169e3df..f0e74a17 --- 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 old mode 100644 new mode 100755 index 5c09d369..6456bee7 --- 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