From 1e9e1d017cd629edfbd2f15057a664a7ee031077 Mon Sep 17 00:00:00 2001 From: malina Date: Tue, 27 Feb 2024 16:15:23 +0100 Subject: [PATCH] Parallel orbit BBA for Maxwell --- pySC/correction/bba.py | 99 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 90 insertions(+), 9 deletions(-) diff --git a/pySC/correction/bba.py b/pySC/correction/bba.py index 6e90ceb..4255dc1 100644 --- a/pySC/correction/bba.py +++ b/pySC/correction/bba.py @@ -3,7 +3,7 @@ from pySC.core.beam import bpm_reading, all_elements_reading from pySC.core.classes import DotDict -from pySC.core.constants import TRACK_TBT, NUM_TO_AB, SETTING_REL, SETTING_ABS +from pySC.core.constants import TRACK_TBT, NUM_TO_AB, SETTING_REL, SETTING_ABS, SETTING_ADD from pySC.utils import at_wrapper, logging_tools, sc_tools, stats from pySC.correction import orbit_trajectory @@ -94,10 +94,54 @@ def trajectory_bba(SC, bpm_ords, mag_ords, **kwargs): return SC, bba_offsets, bba_offset_errors +def orbit_bba_parallel(SC, bpm_ords, mag_ords, **kwargs): + par = DotDict(dict(n_steps=10, fit_order=1, magnet_order=1, skewness=False, setpoint_method=SETTING_REL, + magnet_strengths=np.array([0.95, 1.05]), RMstruct=[], orbBumpWindow=5, BBABPMtarget=1E-3, + cm_ord_orbit=np.array([], dtype=int), cm_ord_setpoint=np.ones(1), + num_downstream_bpms=len(SC.ORD.BPM), dipole_compensation=True, + use_bpm_reading_for_orbit_bump_ref=False, plotResults=False)) + par.update(**kwargs) + par = _check_input(bpm_ords, mag_ords, par) + if SC.INJ.trackMode == TRACK_TBT: + raise ValueError('Beam-orbit-based alignment does not work in TBT mode. ' + 'Please set: SC.INJ.trackMode to "ORB" or "PORB".') + SC0 = SC.very_deep_copy() + bba_offsets = np.full(bpm_ords.shape, np.nan) + bba_offset_errors = np.full(bpm_ords.shape, np.nan) + + for n_dim in range(bpm_ords.shape[0]): + cm_ords, bpm_ranges = _response_scan(SC, n_dim, par) + LOGGER.info(f"Scanned plane {n_dim}") + with multiprocessing.Pool() as pool: + items = [(SC0, bpm_ords, cm_ords, j_bpm, mag_ords, n_dim, bpm_ranges, par) for j_bpm + in range(bpm_ords.shape[1])] + for results in pool.imap(single_plane_bba_orbit, items): + bba_offset, bba_offset_error, j_bpm = results + bba_offsets[n_dim, j_bpm], bba_offset_errors[n_dim, j_bpm] = bba_offset, bba_offset_error + SC = apply_bpm_offsets(SC, bpm_ords, bba_offsets, bba_offset_errors) + if par.plot_results: + plot_bba_results(SC, bpm_ords, bba_offsets, bba_offset_errors) + return SC, bba_offsets, bba_offset_errors + + +def single_plane_bba_orbit(inputs): + SC0, bpm_ords, cm_ords, j_bpm, mag_ords, n_dim, bpm_ranges, par = inputs + bpm_ind = np.where(bpm_ords[n_dim, j_bpm] == SC0.ORD.BPM)[0][0] + m_ord = mag_ords[n_dim, j_bpm] + mask = ~np.isnan(bpm_ranges[:, bpm_ind]) + if np.sum(mask) < 1: + return 0.0, 1.0, j_bpm + set_ind = np.argmax(bpm_ranges[mask, bpm_ind]) + bpm_pos, orbits = _data_measurement_orbit_simple(SC0, m_ord, bpm_ind, j_bpm, n_dim, par, cm_ords[set_ind], + par.cm_ord_setpoint) + bba_offsets, bba_offset_errors = _data_evaluation(SC0, bpm_pos, orbits, par.magnet_strengths[n_dim, j_bpm], n_dim, m_ord, par) + return bba_offsets, bba_offset_errors, j_bpm + def orbit_bba(SC, bpm_ords, mag_ords, **kwargs): par = DotDict(dict(n_steps=10, fit_order=1, magnet_order=1, skewness=False, setpoint_method=SETTING_REL, magnet_strengths=np.array([0.95, 1.05]), RMstruct=[], orbBumpWindow=5, BBABPMtarget=1E-3, - maxNumOfDownstreamBPMs=len(SC.ORD.BPM), dipole_compensation=True, + cm_ord_orbit=np.array([], dtype=int), cm_ord_setpoint=np.ones(1), + num_downstream_bpms=len(SC.ORD.BPM), dipole_compensation=True, use_bpm_reading_for_orbit_bump_ref=False, plotResults=False)) par.update(**kwargs) par = _check_input(bpm_ords, mag_ords, par) @@ -107,15 +151,22 @@ def orbit_bba(SC, bpm_ords, mag_ords, **kwargs): bba_offsets = np.full(bpm_ords.shape, np.nan) bba_offset_errors = np.full(bpm_ords.shape, np.nan) - for j_bpm in range(bpm_ords.shape[1]): # j_bpm: Index of BPM adjacent to magnet for BBA - LOGGER.info(f"BPM number {j_bpm}") - for n_dim in range(bpm_ords.shape[0]): - LOGGER.debug(f'BBA-BPM {j_bpm}/{bpm_ords.shape[1]}, n_dim = {n_dim}') + for n_dim in range(bpm_ords.shape[0]): + cm_ords, bpm_ranges = _response_scan(SC, n_dim, par) + LOGGER.info(f"Scanned plane {n_dim}") + for j_bpm in range(bpm_ords.shape[1]): # j_bpm: Index of BPM adjacent to magnet for BBA + LOGGER.info(f"BPM number {j_bpm}") bpm_ind = np.where(bpm_ords[n_dim, j_bpm] == SC.ORD.BPM)[0][0] m_ord = mag_ords[n_dim, j_bpm] + mask = ~np.isnan(bpm_ranges[:, bpm_ind]) + if np.sum(mask) < 1: + bba_offsets[n_dim, j_bpm], bba_offset_errors[n_dim, j_bpm] = 0.0, 1.0 + continue + set_ind = np.argmax(bpm_ranges[mask, bpm_ind]) SC0 = SC.very_deep_copy() - bpm_pos, orbits = _data_measurement_orb(SC, m_ord, bpm_ind, j_bpm, n_dim, par, - *_get_orbit_bump(SC, m_ord, bpm_ords[n_dim, j_bpm], n_dim, par)) + # bpm_pos, orbits = _data_measurement_orb(SC, m_ord, bpm_ind, j_bpm, n_dim, par, + # *_get_orbit_bump(SC, m_ord, bpm_ords[n_dim, j_bpm], n_dim, par)) + bpm_pos, orbits = _data_measurement_orbit_simple(SC, m_ord, bpm_ind, j_bpm, n_dim, par, cm_ords[set_ind], par.cm_ord_setpoint) bba_offsets[n_dim, j_bpm], bba_offset_errors[n_dim, j_bpm] = _data_evaluation(SC, bpm_pos, orbits, par.magnet_strengths[n_dim, j_bpm], n_dim, m_ord, par) SC = apply_bpm_offsets(SC, bpm_ords, bba_offsets, bba_offset_errors) if par.plot_results: @@ -219,7 +270,7 @@ def _data_evaluation(SC, bpm_pos, trajectories, mag_vec, n_dim, m_ord, par): center[j] = -p[1] / (par.fit_order * p[0]) # zero-crossing if linear, minimum is quadratic center_err[j] = np.sqrt(center[j] ** 2 * (pcov[0,0]/p[0]**2 + pcov[1,1]/p[1]**2 - 2 * pcov[0, 1] / p[0] / p[1])) mask = ~np.isnan(center) - offset_change = stats.weighted_mean(center[mask], center_err[mask]) + offset_change = stats.weighted_mean(center[mask], center_err[mask]) # TODO catch ZeroDivisionError offset_change_error = stats.weighted_error(center[mask]-offset_change, center_err[mask]) / np.sqrt(stats.effective_sample_size(center[mask], stats.weights_from_errors(center_err[mask]))) if not par.dipole_compensation and n_dim == 0 and SC.RING[m_ord].NomPolynomB[1] != 0: offset_change += getattr(SC.RING[m_ord], 'BendingAngle', 0) / SC.RING[m_ord].NomPolynomB[1] / SC.RING[m_ord].Length @@ -400,6 +451,36 @@ def _get_orbit_bump(SC, cm_ord, bpm_ord, n_dim, par): # TODO return cm_ords, cm_vec +def _data_measurement_orbit_simple(SC, m_ord, bpm_ind, j_bpm, n_dim, par, cm_ords, cm_steps): + meas_dim = 1 - n_dim if par.skewness else n_dim + factor = np.linspace(-1, 1, par.n_steps) + orbits = np.full((par.n_steps, len(par.magnet_strengths[n_dim, j_bpm]), len(SC.ORD.BPM)), np.nan) + bpm_pos = np.full((par.n_steps, len(par.magnet_strengths[n_dim, j_bpm])), np.nan) + init_kicks = SC.get_cm_setpoints(cm_ords, skewness=bool(n_dim)) + for n_q, setpoint_q in enumerate(par.magnet_strengths[n_dim, j_bpm]): + SC.set_magnet_setpoints(m_ord, setpoint_q, par.skewness, par.magnet_order, method=par.setpoint_method, + dipole_compensation=par.dipole_compensation) + for step in range(par.n_steps): + SC.set_cm_setpoints(cm_ords, init_kicks + factor[step] * cm_steps, bool(n_dim), method=SETTING_ABS) + bpm_readings = bpm_reading(SC)[0] + bpm_pos[step, n_q] = bpm_readings[n_dim, bpm_ind] + orbits[step, n_q, :] = bpm_readings[meas_dim, :] + return bpm_pos, orbits + + +def _response_scan(SC, n_dim, par): + n_setpoints = len(par.cm_ord_orbit) + bpm_ranges = np.zeros((n_setpoints, len(SC.ORD.BPM))) + for i, ord1 in enumerate(par.cm_ord_orbit): + SC.set_cm_setpoints(ord1, par.cm_ord_setpoint, skewness=bool(n_dim), method=SETTING_ADD) + bpm_plus = bpm_reading(SC)[0][n_dim, :] + SC.set_cm_setpoints(ord1, - 2 * par.cm_ord_setpoint, skewness=bool(n_dim), method=SETTING_ADD) + bpm_minus = bpm_reading(SC)[0][n_dim, :] + bpm_ranges[i, :] = np.abs(bpm_plus - bpm_minus) + SC.set_cm_setpoints(ord1, par.cm_ord_setpoint, skewness=bool(n_dim), method=SETTING_ADD) + return par.cm_ord_orbit, bpm_ranges + + def _plot_bba_step(SC, ax, bpm_ind, n_dim): s_pos = at_wrapper.findspos(SC.RING) bpm_readings, all_elements_positions = all_elements_reading(SC)