diff --git a/apsuite/commisslib/impedance_ivu_meas.py b/apsuite/commisslib/impedance_ivu_meas.py new file mode 100644 index 00000000..ccd0f4e1 --- /dev/null +++ b/apsuite/commisslib/impedance_ivu_meas.py @@ -0,0 +1,299 @@ +""".""" + +import numpy as np + +from siriuspy.devices import SOFB, IVU +from siriuspy.search import BPMSearch + +from apsuite.commisslib.meas_bpms_signals import ( + AcqBPMsSignals as _BaseAcq, + AcqBPMsSignalsParams as _BaseParams, +) +from apsuite.utils import ThreadedMeasBaseClass as _BaseThreaded + + +class ImpedanceIVUMeasParams(_BaseParams): + """.""" + + ADC_NSAMPLES_PER_TURN = 382 + HARM_NUM = 864 + + def __init__(self): + """.""" + super().__init__() + self.num_acquisitions = 10 + self.save_raw_data = False + self.num_buckets_to_process = 2 + self._nrturns = 0 + self.nrturns = 500 + self.bucket_hi_charge = 1 + self.bucket_lo_charge = 530 + self.acq_rate = 'ADCSwp' + self.signals2acq = 'ABCD' + self.timeout = 10 + self.event_mode = 'Injection' + self.timing_event = 'Study' + + def __str__(self): + """.""" + stg = 'AcqBPMsSignalsParams:\n' + stg += ''.join([f' {l}\n' for l in super().__str__().splitlines()]) + stg += '\nImpedanceIVUMeasParams:\n' + stg += f' num_acquisitions = {self.num_acquisitions}\n' + stg += f' save_raw_data = {self.save_raw_data}\n' + stg += f' num_buckets_to_process = {self.num_buckets_to_process}\n' + stg += f' nrturns = {self.nrturns}\n' + stg += f' bucket_hi_charge = {self.bucket_hi_charge}\n' + stg += f' bucket_lo_charge = {self.bucket_lo_charge}\n' + return stg + + @property + def nrturns(self): + """.""" + return self._nrturns + + @nrturns.setter + def nrturns(self, val): + self._nrturns = int(val) + self.nrpoints_after = self.nrturns * self.ADC_NSAMPLES_PER_TURN + self.nrpoints_before = 0 + + +class ImpedanceIVUMeas(_BaseThreaded, _BaseAcq): + """.""" + + def __init__(self, isonline=True, bpmtype='all'): + """.""" + bpmnames = BPMSearch.get_names(filters={'sec': 'SI', 'dev': 'BPM'}) + if bpmtype.startswith('odd'): + bpmnames = bpmnames[::2] + elif bpmtype.startswith('even'): + bpmnames = bpmnames[1::2] + _BaseThreaded.__init__(self, isonline=isonline, target=self._measure) + _BaseAcq.__init__(self, isonline=self.isonline, bpmnames=bpmnames) + self.params = ImpedanceIVUMeasParams() + + def create_devices(self, bpmnames=None): + """.""" + _BaseAcq.create_devices(self, bpmnames=bpmnames) + self.devices['sofb'] = SOFB(SOFB.DEVICES.SI) + self.devices['ivu18_08'] = IVU(IVU.DEVICES.IVU18_08SB) + self.devices['ivu18_14'] = IVU(IVU.DEVICES.IVU18_14SB) + + def get_data(self): + """.""" + data = super().get_data() + data['sofb_refx'] = self.devices['sofb'].refx + data['sofb_refy'] = self.devices['sofb'].refy + data['sofb_orbx'] = self.devices['sofb'].orbx + data['sofb_orby'] = self.devices['sofb'].orby + data['sofb_bpmxenbl'] = self.devices['sofb'].bpmxenbl + data['sofb_bpmyenbl'] = self.devices['sofb'].bpmyenbl + data['sofb_nr_points'] = self.devices['sofb'].nr_points + data['sofb_kickch'] = self.devices['sofb'].kickch + data['sofb_kickcv'] = self.devices['sofb'].kickcv + data['sofb_kickrf'] = self.devices['sofb'].kickrf + data['ivu18_08_gap'] = self.devices['ivu18_08'].gap + data['ivu18_14_gap'] = self.devices['ivu18_14'].gap + return data + + def load_and_apply(self, fname: str): + """.""" + return _BaseThreaded.load_and_apply(self, fname) + + load_and_apply.__doc__ = _BaseThreaded.load_and_apply.__doc__ + + def _measure(self): + data = [] + for i in range(self.params.num_acquisitions): + print( + f'Acquisition {i + 1:02d}/{self.params.num_acquisitions:02d}' + ) + if self._stopevt.is_set(): + break + self.acquire_data() + data.append(self.data) + + print('Finished!') + self.data = data + self.process_data() + self._filter_data_to_save() + + def process_data(self, idcs_to_discard=None, return_all=False): + """.""" + idcs_to_discard = idcs_to_discard or [] + data = self.data + if isinstance(data, dict): + data = [data] + for i, dt in enumerate(data): + if i in idcs_to_discard: + continue + dic = self._proc_single_data(dt, return_all=return_all) + dt.update(dic) + + def calc_delta_orbit_2_bunches(self): + """Calculate orbit variation between the two stored bunches. + + Raises: + RuntimeError: If there is no data acquired. + RuntimeError: If data is not processed yet. + + Returns: + dorb: orbit deviation between the two stored bunches. + orb1: orbit of the first stored bunch. + orb2: orbit of the second stored bunch. + """ + if not self.data: + raise RuntimeError('Get data First.') + orb1, orb2 = [], [] + for dt in self.data: + if 'b1_posx' not in dt or 'b2_posx' not in dt: + raise RuntimeError( + 'Missing bunch positions in data. Process data first.' + ) + orb1.append(np.vstack([dt['b1_posx'], dt['b1_posy']])) + orb2.append(np.vstack([dt['b2_posx'], dt['b2_posy']])) + orb1 = np.array(orb1) + orb2 = np.array(orb2) + dorb = orb1 - orb2 + return dorb, orb1, orb2 + + def calc_current_2_bunches(self): + """Calculate current of the two stored bunches. + + Raises: + RuntimeError: If there is no data acquired. + RuntimeError: If data is not processed yet. + + Returns: + curr1: current of the first stored bunch. + curr2: current of the second stored bunch. + + """ + if not self.data: + raise RuntimeError('Get data First.') + curr1, curr2 = [], [] + for dt in self.data: + if 'b1_curr' not in dt or 'b2_curr' not in dt: + raise RuntimeError( + 'Missing bunch currents in data. Process data first.' + ) + curr1.append(dt['b1_curr']) + curr2.append(dt['b2_curr']) + curr1 = np.array(curr1) + curr2 = np.array(curr2) + return curr1, curr2 + + def calc_sum_signal_2_bunches(self): + """Calculate the sum signal of the two stored bunches. + + Raises: + RuntimeError: If there is no data acquired. + RuntimeError: If data is not processed yet. + + Returns: + sumt: sum signal of the two bunches. + sum1: sum signal of the first stored bunch. + sum2: sum signal of the second stored bunch. + + """ + if not self.data: + raise RuntimeError('Get data First.') + sumt, sum1, sum2 = [], [], [] + for dt in self.data: + if 'b1_sum' not in dt or 'b2_sum' not in dt: + raise RuntimeError( + 'Missing sum signals in data. Process data first.' + ) + sumt.append(dt['bt_sum']) + sum1.append(dt['b1_sum']) + sum2.append(dt['b2_sum']) + sumt = np.array(sumt) + sum1 = np.array(sum1) + sum2 = np.array(sum2) + return sumt, sum1, sum2 + + def calc_sofb_orbit(self, isref=False): + """Calculate the SOFB orbit. + + Raises: + RuntimeError: If there is no data acquired. + RuntimeError: If data is not processed yet. + + Returns: + orb: The SOFB orbit. + """ + if not self.data: + raise RuntimeError('Get data First.') + orb = [] + prop = 'sofb_' + ('ref' if isref else 'orb') + for dt in self.data: + orb.append(np.hstack([dt[prop + 'x'], dt[prop + 'y']])) + return np.vstack(orb).T + + def _filter_data_to_save(self): + if self.params.save_raw_data: + return + for dt in self.data: + for ant in 'abcd': + dt.pop('ampl' + ant) + + def _proc_single_data(self, data, return_all=False): + b1_offset = 50 + windowp = 20 + windown = -10 + nbuc2proc = self.params.num_buckets_to_process + bhigh = self.params.bucket_hi_charge + blow = self.params.bucket_lo_charge + nsamp_pturn = ImpedanceIVUMeasParams.ADC_NSAMPLES_PER_TURN + hnum = ImpedanceIVUMeasParams.HARM_NUM + + ant_raw = np.array([data['ampl' + ant] for ant in 'abcd']) + ant_raw = ant_raw.swapaxes( + 1, 2 + ) # [4, 382 * N, 160] --> [4, 160, 382 * N] + curr = data['stored_current'] + + ant_abs = np.abs(ant_raw) + ant_amax = ant_abs[..., :nsamp_pturn].argmax(axis=-1) + + nsamp2keep = ant_raw.shape[-1] + nturn2keep = nsamp2keep // nsamp_pturn + idx = np.arange(nsamp2keep) + old_idx = (idx - b1_offset + ant_amax[..., None]) % nsamp2keep + + ant_raw2 = np.take_along_axis(ant_raw, old_idx, axis=-1) + ant_raw2 = ant_raw2.reshape(ant_raw2.shape[:2] + (nturn2keep, -1)) + + dic = {} + if return_all: + dic['ant_raw'] = ant_raw + dic['ant_amax'] = ant_amax + dic['ant_raw2'] = ant_raw2 + + b2_offset = abs(blow - bhigh) / hnum * nsamp_pturn + b2_offset = int(b2_offset) + b1_offset + slcs = [ + slice(b1_offset + windown, b1_offset + windowp), + slice(b2_offset + windown, b2_offset + windowp), + ] + pref = lambda x: f'b{x + 1}_' + for i in range(nbuc2proc): + b_sigs = ant_raw2[..., slcs[i]].std(axis=-1) + b_posx, b_posy = ImpedanceIVUMeas.calc_positions_from_amplitudes( + b_sigs + ) + b_sum = b_sigs.sum(axis=0) + + dic[pref(i) + 'posx'] = b_posx + dic[pref(i) + 'posy'] = b_posy + dic[pref(i) + 'sum'] = b_sum + if return_all: + dic[pref(i) + 'sigs'] = b_sigs + + bt_sum = sum([dic[pref(i) + 'sum'] for i in range(nbuc2proc)]) + dic['bt_sum'] = bt_sum + for i in range(nbuc2proc): + dic[pref(i) + 'curr'] = dic[pref(i) + 'sum'] * curr / bt_sum + + return dic diff --git a/apsuite/commisslib/meas_bpms_signals.py b/apsuite/commisslib/meas_bpms_signals.py index 238265e9..cfbe1219 100644 --- a/apsuite/commisslib/meas_bpms_signals.py +++ b/apsuite/commisslib/meas_bpms_signals.py @@ -90,13 +90,13 @@ class AcqBPMsSignals(_BaseClass): BPM_TRIGGER = "SI-Fam:TI-BPM" PSM_TRIGGER = "SI-Fam:TI-BPM-PsMtm" - def __init__(self, isonline=True, ispost_mortem=False): + def __init__(self, isonline=True, ispost_mortem=False, bpmnames=None): """.""" super().__init__(params=AcqBPMsSignalsParams(), isonline=isonline) self._ispost_mortem = ispost_mortem if self.isonline: - self.create_devices() + self.create_devices(bpmnames=bpmnames) calc_positions_from_amplitudes = staticmethod( FamBPMs.calc_positions_from_amplitudes) @@ -122,11 +122,12 @@ def load_and_apply(self, fname: str): self.data = data return ret - def create_devices(self): + def create_devices(self, bpmnames=None): """.""" self.devices["currinfo"] = CurrInfoSI() self.devices["fambpms"] = FamBPMs( devname=FamBPMs.DEVICES.SI, + bpmnames=bpmnames, ispost_mortem=self._ispost_mortem, props2init="acq", ) diff --git a/apsuite/utils.py b/apsuite/utils.py index eede6298..b514701d 100644 --- a/apsuite/utils.py +++ b/apsuite/utils.py @@ -132,7 +132,7 @@ class MeasBaseClass(DataBaseClass): def __init__(self, params=None, isonline=True): """.""" - super().__init__(params=params) + DataBaseClass.__init__(self, params=params) self.isonline = bool(isonline) self.devices = dict() self.analysis = dict() @@ -170,7 +170,7 @@ class ThreadedMeasBaseClass(MeasBaseClass): def __init__(self, params=None, target=None, isonline=True): """.""" - super().__init__(params=params, isonline=isonline) + MeasBaseClass.__init__(self, params=params, isonline=isonline) self._target = target self._stopevt = _Event() self._finished = _Event()