diff --git a/pySC/core/constants.py b/pySC/core/constants.py index d73fe26..63b65aa 100644 --- a/pySC/core/constants.py +++ b/pySC/core/constants.py @@ -23,3 +23,4 @@ SETTING_ADD: str = "add" SETTING_METHODS: Tuple[str, str, str] = (SETTING_ABS, SETTING_REL, SETTING_ADD) SPEED_OF_LIGHT = 299792458 +SETPOINT: str = "SetPoint" diff --git a/pySC/core/lattice_setting.py b/pySC/core/lattice_setting.py index deedad0..bc1df09 100644 --- a/pySC/core/lattice_setting.py +++ b/pySC/core/lattice_setting.py @@ -5,66 +5,11 @@ This module contains the 'machine-based' functions to interact with lattice under study. """ import numpy as np -from typing import Union, List - from at import Lattice from numpy import ndarray - -from pySC.core.constants import SETTING_METHODS, SETTING_ABS, SETTING_REL, SETTING_ADD, NUM_TO_AB -from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.utils.at_wrapper import atgetfieldvalues from pySC.utils import logging_tools LOGGER = logging_tools.get_logger(__name__) -SETPOINT = "SetPoint" - - -def set_cavity_setpoints(SC: SimulatedCommissioning, - ords: Union[int, List[int], ndarray], - setpoints: Union[float, List[float], ndarray], - param: str, method: str = SETTING_ABS) -> SimulatedCommissioning: - """ - Set RF properties to setpoints - - Set the setpoints of `Voltage`, `Frequency` or `TimeLag` as specified in "param" of the rf - cavities specified in `ords`. If only a single setpoint is given for multiple cavities, - the setpoint is applied to all cavities. - - Args: - SC: SimulatedCommissioning class instance - ords: Array of cavity ordinates in the lattice structure (SC.ORD.RF) - setpoints: Setpoints (array or single value for all cavities) - param: String ('Voltage', 'Frequency' or 'TimeLag') specifying which cavity field should be set. - method: 'abs' (default), Use absolute setpoint - 'rel', Use relative setpoint to nominal value - 'add', Add setpoints to current value - - Returns: - The modified SC structure. - - Examples: - Sets the time lag of all cavities registered in SC to zero:: - - SC = set_cavity_setpoints(SC, ords=SC.ORD.Cavity, - setpoints=0, param='TimeLag') - - Adds 1kHz to the frequency of the first cavity:: - - SC = set_cavity_setpoints(SC, ords=SC.ORD.Cavity(1), - setpoints=1E3, param='Frequency', - method='add') - - """ - ords_1d, setpoints_1d = _check_input_and_setpoints(method, ords, setpoints) - setpoint_str = f"{param}{SETPOINT}" - if method == SETTING_REL: - setpoints_1d *= atgetfieldvalues(SC.RING, ords_1d, setpoint_str) - if method == SETTING_ADD: - setpoints_1d += atgetfieldvalues(SC.RING, ords_1d, setpoint_str) - for i, ord in enumerate(ords_1d): - setattr(SC.RING[ord], setpoint_str, setpoints_1d[i]) - SC.update_cavities(ords_1d) - return SC def switch_rf(ring: Lattice, ords: ndarray, state: bool) -> Lattice: @@ -74,165 +19,6 @@ def switch_rf(ring: Lattice, ords: ndarray, state: bool) -> Lattice: return ring -def get_cm_setpoints(SC: SimulatedCommissioning, ords: Union[int, List[int], ndarray], skewness: bool) -> ndarray: - """ - - Return current dipole Corrector Magnets (CM) setpoints - - Reads the setpoints of the CMs specified in `ords` in the dimension `skewness`. - - Args: - SC: SimulatedCommissioning class instance - ords: Array of CM ordinates in the lattice structure (ex: SC.ORD.CM[0]) - skewness: boolean specifying CM dimension ([False|True] -> [hor|ver]) - - Returns: - CM setpoints [rad] - """ - ords_1d = np.ravel(np.array([ords], dtype=int)) - order = 0 - ndim = int(skewness) - letter = NUM_TO_AB[ndim] - setpoints = atgetfieldvalues(SC.RING, ords_1d, f"{SETPOINT}{letter}", order) - for i, ord1d in enumerate(ords_1d): - if SC.RING[ord1d].PassMethod != 'CorrectorPass': - # positive setpoint -> positive kick -> negative horizontal field - setpoints[i] *= (-1) ** (ndim + 1) * SC.RING[ord1d].Length - return setpoints - - -def set_cm_setpoints(SC: SimulatedCommissioning, - ords: Union[int, List[int], ndarray], - setpoints: Union[float, List[float], ndarray], - skewness: bool, method: str = SETTING_ABS) -> SimulatedCommissioning: - """ - Sets dipole corrector magnets to different setpoints - - Sets horizontal or vertical CMs as specified in `ords` and `skewness`, respectively, to `setpoints` - [rad] and updates the magnetic fields. If the corresponding setpoint exceeds the CM limit - specified in the corresponding lattice field `CMlimit`, the CM is clipped to that value - and a warning is being printed (to switch off, use `warning('off','SC:CM1'))`. Positive setpoints - will results in kicks in the positive horizontal or vertical direction. - - Args: - SC: SimulatedCommissioning class instance - ords: Array of CM ordinates in the lattice structure (ex: SC.ORD.CM[0]) - setpoints: CM setpoints (array or single value for all CMs) [rad] - skewness: boolean specifying CM dimension ([False|True] -> [hor|ver]) - method: 'abs' (default), Use absolute setpoint - 'rel', Use relative setpoint to current value - 'add', Add setpoints to current value - - Returns: - - The lattice structure with modified and applied setpoints - - Examples: - Set all registered horizontal CMs to zero:: - - SC = set_cm_setpoints(SC, ords=SC.ORD.HCM, - skewness=False, setpoints=0.0) - - Add 10urad to the fourth registered vertical CM:: - - SC = set_cm_setpoints(SC, ords=SC.ORD.VCM[4], - setpoints=1E-5, skewness=True, - method='add') - - """ - # TODO corrector does not have PolynomA/B in at? - ords_1d, setpoints_1d = _check_input_and_setpoints(method, ords, setpoints) - order = 0 - ndim = int(skewness) - letter = NUM_TO_AB[ndim] - for i, ord in enumerate(ords_1d): - # positive setpoint -> positive kick -> negative horizontal field - norm_by = (-1) ** (ndim + 1) * SC.RING[ord].Length if SC.RING[ord].PassMethod != 'CorrectorPass' else 1 - if method == SETTING_REL: - setpoints_1d[i] *= getattr(SC.RING[ord], f"{SETPOINT}{letter}")[order] * norm_by - if method == SETTING_ADD: - setpoints_1d[i] += getattr(SC.RING[ord], f"{SETPOINT}{letter}")[order] * norm_by - if hasattr(SC.RING[ord], 'CMlimit') and abs(setpoints_1d[i]) > abs(SC.RING[ord].CMlimit[ndim]): - LOGGER.info(f'CM (ord: {ord} / dim: {ndim}) is clipping') - setpoints_1d[i] = np.sign(setpoints_1d[i]) * SC.RING[ord].CMlimit[ndim] - getattr(SC.RING[ord], f"{SETPOINT}{letter}")[order] = setpoints_1d[i] / norm_by - - SC.update_magnets(ords_1d) - return SC - - -def set_magnet_setpoints(SC: SimulatedCommissioning, - ords: Union[int, List[int], ndarray], - setpoints: Union[float, List[float], ndarray], - skewness: bool, order: int, method: str = SETTING_ABS, - dipole_compensation: bool = False) -> SimulatedCommissioning: - """ - Sets magnets to setpoints - - Sets magnets (except CMs) as specified in `ords` to `setpoints` while `order` and `skewness` defines - which field entry should be used (see below). The setpoints may be given relative to their nominal - value or in absolute terms. If the considered quadrupole is a combined function magnet with - non-zero bending angle and the kick compensation flag 'dipole_compensation'=True, the appropriate bending - angle difference is calculated and the horizontal CM setpoint is changed accordingly to compensate - for that dipole kick difference. - If the setpoint of a skew quadrupole exceeds the limit specified in the corresponding lattice - field `SkewQuadLimit`, the setpoint is clipped to that value and a warning is logged. - - Args: - SC: SimulatedCommissioning class instance - ords: Array of magnets ordinates in the lattice structure (ex: SC.ORD.HCM) (numpy.array() or list of int [int,int,..]) - setpoints: magnets setpoints (array or single value for all magnets). - setpoints are assigned to the given order and skewness, i.e. once updated through - SimulatedCommissioning.apply_errors, they correspond to a single element of PolynomA or PolynomB - skewness: boolean specifying magnet plane ([False|True] -> [PolynomB|PolynomA]) - method: 'abs' (default), Use absolute setpoint - 'rel', Use relative setpoint to nominal value - 'add', Add setpoints to current value - order: Numeric value defining the order of the considered magnet: [0,1,2,...] => [dip,quad,sext,...] - dipole_compensation: (default = False) Used for combined function magnets. If this flag is set and if there is a horizontal CM - registered in the considered magnet, the CM is used to compensate the bending angle difference - if the applied quadrupole setpoints differs from the design value. - - Returns: - The SimulatedCommissioning class instance containing lattice with modified and applied setpoints. - - Examples: - Identify the ordinates of all elements named `'SF'` and switch their sextupole component off:: - - ords = SCgetOrds(SC.RING,'SF') - SC.register_magnets(ords) - SC = set_magnet_setpoints(SC, ords=ords, - skewness=False, order=2, setpoints=0.0, - method='abs') - - Identify the ordinates of all elements named `QF` and `QD` and set their quadrupole component to 99% of their design value:: - - ords = SCgetOrds(SC.RING,'QF|QD') - SC.register_magnets(ords) - SC = set_magnet_setpoints(SC, ords=ords, - skewness=False, order=1, setpoints=0.99, - method='rel') - - """ - ords_1d, setpoints_1d = _check_input_and_setpoints(method, ords, setpoints) - letter = NUM_TO_AB[int(skewness)] - if method == SETTING_REL: - setpoints_1d *= atgetfieldvalues(SC.RING, ords_1d, f"NomPolynom{letter}", order) - if method == SETTING_ADD: - setpoints_1d += atgetfieldvalues(SC.RING, ords_1d, f"{SETPOINT}{letter}", order) - for i, ord in enumerate(ords_1d): - if skewness and order == 1 and getattr(SC.RING[ord], 'SkewQuadLimit', np.inf) < np.abs(setpoints_1d[i]): - LOGGER.info(f'SkewLim \n Skew quadrupole (ord: {ord}) is clipping') - setpoints_1d[i] = np.sign(setpoints_1d[i]) * SC.RING[ord].SkewQuadLimit - # TODO should check CF magnets - if dipole_compensation and order == 1: # quad # TODO check also skewness? - SC = _dipole_compensation(SC, ord, setpoints_1d[i]) - getattr(SC.RING[ord], f"{SETPOINT}{letter}")[order] = setpoints_1d[i] - - SC.update_magnets(ords_1d) - return SC - - def switch_cavity_and_radiation(ring: Lattice, *args: str) -> Lattice: # TODO some at methods do that? """ switch cavity / radiation to on / off @@ -280,20 +66,3 @@ def switch_cavity_and_radiation(ring: Lattice, *args: str) -> Lattice: # TODO s return switch_rf(ring, np.arange(len(ring)), True) return ring - -def _dipole_compensation(SC, ord, setpoint): - if getattr(SC.RING[ord], 'BendingAngle', 0) != 0 and ord in SC.ORD.HCM: - return set_cm_setpoints( - SC, ord, (setpoint - SC.RING[ord].SetPointB[1]) / SC.RING[ord].NomPolynomB[1] * SC.RING[ord].BendingAngle, - skewness=False, method=SETTING_ADD) - return SC - - -def _check_input_and_setpoints(method, ords, setpoints): - if method not in SETTING_METHODS: - raise ValueError(f'Unsupported setpoint method: {method}. Allowed options are: {SETTING_METHODS}.') - ords_1d = np.ravel(np.array([ords], dtype=int)) - setpoints_1d = np.ravel(np.array([setpoints])) - if len(setpoints_1d) not in (1, len(ords_1d)): - raise ValueError(f'Setpoints have to have length of 1 or matching to the length or ordinates.') - return ords_1d, (np.repeat(setpoints_1d, len(ords_1d)) if len(setpoints_1d) == 1 else setpoints_1d) diff --git a/pySC/core/simulated_commissioning.py b/pySC/core/simulated_commissioning.py index ffbc6f5..7312ba6 100644 --- a/pySC/core/simulated_commissioning.py +++ b/pySC/core/simulated_commissioning.py @@ -7,7 +7,7 @@ """ import copy import re -from typing import Tuple +from typing import Tuple, Union, List import numpy as np from at import Lattice @@ -15,10 +15,12 @@ from pySC.core.classes import Injection, Sigmas, Indices, DotDict from pySC.core.constants import (BPM_ERROR_FIELDS, RF_ERROR_FIELDS, RF_PROPERTIES, MAGNET_TYPE_FIELDS, - MAGNET_ERROR_FIELDS, AB, SUPPORT_TYPES, SUPPORT_ERROR_FIELDS) + MAGNET_ERROR_FIELDS, AB, SUPPORT_TYPES, SUPPORT_ERROR_FIELDS, SETTING_ABS, SETTING_REL, + SETTING_ADD, NUM_TO_AB, SETPOINT) from pySC.utils import logging_tools -from pySC.utils.at_wrapper import findspos -from pySC.utils.classdef_tools import update_double_ordinates, add_padded, intersect, randn_cutoff, s_interpolation +from pySC.utils.at_wrapper import findspos, atgetfieldvalues +from pySC.utils.classdef_tools import update_double_ordinates, add_padded, intersect, randn_cutoff, s_interpolation, \ + check_input_and_setpoints from pySC.utils.sc_tools import SCrandnc, SCscaleCircumference, update_transformation LOGGER = logging_tools.get_logger(__name__) @@ -766,6 +768,186 @@ def support_offset_and_roll(self, s_locations: ndarray) -> Tuple[ndarray, ndarra return off, roll return off0, roll0 + def set_cavity_setpoints(self, ords: Union[int, List[int], ndarray], + setpoints: Union[float, List[float], ndarray], + param: str, method: str = SETTING_ABS): + """ + Set RF properties to setpoints + + Set the setpoints of `Voltage`, `Frequency` or `TimeLag` as specified in "param" of the rf + cavities specified in `ords`. If only a single setpoint is given for multiple cavities, + the setpoint is applied to all cavities. + + Args: + ords: Array of cavity ordinates in the lattice structure (SC.ORD.RF) + setpoints: Setpoints (array or single value for all cavities) + param: String ('Voltage', 'Frequency' or 'TimeLag') specifying which cavity field should be set. + method: 'abs' (default), Use absolute setpoint + 'rel', Use relative setpoint to nominal value + 'add', Add setpoints to current value + + Examples: + Sets the time lag of all cavities registered in SC to zero:: + + SC.set_cavity_setpoints(ords=SC.ORD.RF, setpoints=0.0, param='TimeLag') + + Adds 1kHz to the frequency of the first cavity:: + + SC.set_cavity_setpoints(ords=SC.ORD.RF[0], setpoints=1E3, param='Frequency', method='add') + + """ + ords_1d, setpoints_1d = check_input_and_setpoints(method, ords, setpoints) + setpoint_str = f"{param}{SETPOINT}" + if method == SETTING_REL: + setpoints_1d *= atgetfieldvalues(self.RING, ords_1d, setpoint_str) + if method == SETTING_ADD: + setpoints_1d += atgetfieldvalues(self.RING, ords_1d, setpoint_str) + for i, ord in enumerate(ords_1d): + setattr(self.RING[ord], setpoint_str, setpoints_1d[i]) + self.update_cavities(ords_1d) + + def get_cm_setpoints(self, ords: Union[int, List[int], ndarray], skewness: bool) -> ndarray: + """ + + Return current dipole Corrector Magnets (CM) setpoints + + Reads the setpoints of the CMs specified in `ords` in the dimension `skewness`. + + Args: + ords: Array of CM ordinates in the lattice structure (ex: SC.ORD.CM[0]) + skewness: boolean specifying CM dimension ([False|True] -> [hor|ver]) + + Returns: + CM setpoints [rad] + """ + ords_1d = np.ravel(np.array([ords], dtype=int)) + order = 0 + ndim = int(skewness) + letter = NUM_TO_AB[ndim] + setpoints = atgetfieldvalues(self.RING, ords_1d, f"{SETPOINT}{letter}", order) + for i, ord1d in enumerate(ords_1d): + if self.RING[ord1d].PassMethod != 'CorrectorPass': + # positive setpoint -> positive kick -> negative horizontal field + setpoints[i] *= (-1) ** (ndim + 1) * self.RING[ord1d].Length + return setpoints + + def set_cm_setpoints(self, ords: Union[int, List[int], ndarray], + setpoints: Union[float, List[float], ndarray], + skewness: bool, method: str = SETTING_ABS): + """ + Sets dipole corrector magnets to different setpoints + + Sets horizontal or vertical CMs as specified in `ords` and `skewness`, respectively, to `setpoints` + [rad] and updates the magnetic fields. If the corresponding setpoint exceeds the CM limit + specified in the corresponding lattice field `CMlimit`, the CM is clipped to that value + and a warning is being printed (to switch off, use `warning('off','SC:CM1'))`. Positive setpoints + will result in kicks in the positive horizontal or vertical direction. + + Args: + ords: Array of CM ordinates in the lattice structure (ex: SC.ORD.CM[0]) + setpoints: CM setpoints (array or single value for all CMs) [rad] + skewness: boolean specifying CM dimension ([False|True] -> [hor|ver]) + method: 'abs' (default), Use absolute setpoint + 'rel', Use relative setpoint to current value + 'add', Add setpoints to current value + + Examples: + Set all registered horizontal CMs to zero:: + + SC.set_cm_setpoints(ords=SC.ORD.HCM, skewness=False, setpoints=0.0) + + Add 10urad to the fourth registered vertical CM:: + + SC.set_cm_setpoints(ords=SC.ORD.VCM[4], setpoints=1E-5, skewness=True, method='add') + + """ + # TODO corrector does not have PolynomA/B in at? + ords_1d, setpoints_1d = check_input_and_setpoints(method, ords, setpoints) + order = 0 + ndim = int(skewness) + letter = NUM_TO_AB[ndim] + for i, ord in enumerate(ords_1d): + # positive setpoint -> positive kick -> negative horizontal field + norm_by = (-1) ** (ndim + 1) * self.RING[ord].Length if self.RING[ord].PassMethod != 'CorrectorPass' else 1 + if method == SETTING_REL: + setpoints_1d[i] *= getattr(self.RING[ord], f"{SETPOINT}{letter}")[order] * norm_by + if method == SETTING_ADD: + setpoints_1d[i] += getattr(self.RING[ord], f"{SETPOINT}{letter}")[order] * norm_by + if hasattr(self.RING[ord], 'CMlimit') and abs(setpoints_1d[i]) > abs(self.RING[ord].CMlimit[ndim]): + LOGGER.info(f'CM (ord: {ord} / dim: {ndim}) is clipping') + setpoints_1d[i] = np.sign(setpoints_1d[i]) * self.RING[ord].CMlimit[ndim] + getattr(self.RING[ord], f"{SETPOINT}{letter}")[order] = setpoints_1d[i] / norm_by + + self.update_magnets(ords_1d) + + def set_magnet_setpoints(self, ords: Union[int, List[int], ndarray], + setpoints: Union[float, List[float], ndarray], + skewness: bool, order: int, method: str = SETTING_ABS, + dipole_compensation: bool = False): + """ + Sets magnets to setpoints + + Sets magnets (except CMs) as specified in `ords` to `setpoints` while `order` and `skewness` defines + which field entry should be used (see below). The setpoints may be given relative to their nominal + value or in absolute terms. If the considered quadrupole is a combined function magnet with + non-zero bending angle and the kick compensation flag 'dipole_compensation'=True, the appropriate bending + angle difference is calculated and the horizontal CM setpoint is changed accordingly to compensate + for that dipole kick difference. + If the setpoint of a skew quadrupole exceeds the limit specified in the corresponding lattice + field `SkewQuadLimit`, the setpoint is clipped to that value and a warning is logged. + + Args: + ords: Array of magnets ordinates in the lattice structure (ex: SC.ORD.HCM) (numpy.array() or list of int [int,int,..]) + setpoints: magnets setpoints (array or single value for all magnets). + setpoints are assigned to the given order and skewness, i.e. once updated through + SimulatedCommissioning.apply_errors, they correspond to a single element of PolynomA or PolynomB + skewness: boolean specifying magnet plane ([False|True] -> [PolynomB|PolynomA]) + method: 'abs' (default), Use absolute setpoint + 'rel', Use relative setpoint to nominal value + 'add', Add setpoints to current value + order: Numeric value defining the order of the considered magnet: [0,1,2,...] => [dip,quad,sext,...] + dipole_compensation: (default = False) Used for combined function magnets. If this flag is set and if there is a horizontal CM + registered in the considered magnet, the CM is used to compensate the bending angle difference + if the applied quadrupole setpoints differs from the design value. + + Examples: + Identify the ordinates of all elements named `'SF'` and switch their sextupole component off:: + + ords = SCgetOrds(SC.RING,'SF') + SC.register_magnets(ords) + SC.set_magnet_setpoints(ords=ords, skewness=False, order=2, setpoints=0.0, method='abs') + + Identify the ordinates of all elements named `QF` and `QD` and set their quadrupole component to 99% of their design value:: + + ords = SCgetOrds(SC.RING,'QF|QD') + SC.register_magnets(ords) + SC.set_magnet_setpoints(ords=ords, skewness=False, order=1, setpoints=0.99, method='rel') + + """ + ords_1d, setpoints_1d = check_input_and_setpoints(method, ords, setpoints) + letter = NUM_TO_AB[int(skewness)] + if method == SETTING_REL: + setpoints_1d *= atgetfieldvalues(self.RING, ords_1d, f"NomPolynom{letter}", order) + if method == SETTING_ADD: + setpoints_1d += atgetfieldvalues(self.RING, ords_1d, f"{SETPOINT}{letter}", order) + for i, ord in enumerate(ords_1d): + if skewness and order == 1 and getattr(self.RING[ord], 'SkewQuadLimit', np.inf) < np.abs(setpoints_1d[i]): + LOGGER.info(f'SkewLim \n Skew quadrupole (ord: {ord}) is clipping') + setpoints_1d[i] = np.sign(setpoints_1d[i]) * self.RING[ord].SkewQuadLimit + # TODO should check CF magnets + if dipole_compensation and order == 1: # quad # TODO check also skewness? + self._dipole_compensation(ord, setpoints_1d[i]) + getattr(self.RING[ord], f"{SETPOINT}{letter}")[order] = setpoints_1d[i] + + self.update_magnets(ords_1d) + + def _dipole_compensation(self, ord, setpoint): + if getattr(self.RING[ord], 'BendingAngle', 0) != 0 and ord in self.ORD.HCM: + self.set_cm_setpoints( + ord, + (setpoint - self.RING[ord].SetPointB[1]) / self.RING[ord].NomPolynomB[1] * self.RING[ord].BendingAngle, + skewness=False, method=SETTING_ADD) + def verify_structure(self): """ Verifies the integrity of SimilatedCommissionining class instance and warns if things look fishy. diff --git a/pySC/correction/chroma.py b/pySC/correction/chroma.py index dc6088d..15f3c0f 100644 --- a/pySC/correction/chroma.py +++ b/pySC/correction/chroma.py @@ -1,7 +1,6 @@ import numpy as np from scipy.optimize import fmin -from pySC.core.lattice_setting import set_magnet_setpoints from pySC.utils.at_wrapper import atlinopt from pySC.utils import logging_tools @@ -17,7 +16,7 @@ def fit_chroma(SC, s_ords, target_chroma=None, init_step_size=np.array([2, 2]), return SC if tune_knobs_ords is not None and tune_knobs_delta_k is not None: for nFam in range(len(tune_knobs_ords)): - SC = set_magnet_setpoints(SC, tune_knobs_ords[nFam], tune_knobs_delta_k[nFam], False, 1, + SC.set_magnet_setpoints(tune_knobs_ords[nFam], tune_knobs_delta_k[nFam], False, 1, method='add') # TODO quads here? LOGGER.debug(f'Fitting chromaticities from {atlinopt(SC.RING, 0, [])[2]} to {target_chroma}.') # first two elements SP0 = np.zeros((len(s_ords), len(s_ords[0]))) # TODO can the lengts vary @@ -26,12 +25,12 @@ def fit_chroma(SC, s_ords, target_chroma=None, init_step_size=np.array([2, 2]), SP0[nFam][n] = SC.RING[s_ords[nFam][n]].SetPointB[2] fun = lambda x: _fit_chroma_fun(SC, s_ords, x, SP0, target_chroma) sol = fmin(fun, init_step_size, xtol=xtol, ftol=ftol) - SC = set_magnet_setpoints(SC, s_ords, sol + SP0, False, 1, method='abs', dipole_compensation=True) + SC.set_magnet_setpoints(s_ords, sol + SP0, False, 1, method='abs', dipole_compensation=True) LOGGER.debug(f' Final chromaticity: {atlinopt(SC.RING, 0, [])[2]}\n Setpoints change: {sol}.') # first two elements return SC def _fit_chroma_fun(SC, q_ords, setpoints, init_setpoints, target): - SC = set_magnet_setpoints(SC, q_ords, setpoints + init_setpoints, False, 2, method='abs', dipole_compensation=True) + SC.set_magnet_setpoints(q_ords, setpoints + init_setpoints, False, 2, method='abs', dipole_compensation=True) _, _, nu = atlinopt(SC.RING, 0, []) return np.sqrt(np.mean((nu - target) ** 2)) diff --git a/pySC/correction/loco_wrapper.py b/pySC/correction/loco_wrapper.py index 80afabb..8f4e428 100644 --- a/pySC/correction/loco_wrapper.py +++ b/pySC/correction/loco_wrapper.py @@ -2,7 +2,6 @@ from pySC.core.classes import DotDict from pySC.lattice_properties.response_measurement import response_matrix, dispersion -from pySC.core.lattice_setting import set_magnet_setpoints from pySC.utils import logging_tools LOGGER = logging_tools.get_logger(__name__) @@ -103,9 +102,9 @@ def apply_lattice_correction(SC, fit_parameters, dipole_compensation=True, dampi setpoint = fit_parameters.OrigValues[n_group] + damping * ( fit_parameters.IdealValues[n_group] - fit_parameters.Values[n_group]) if field == 'SetPointB': # Normal quadrupole - SC = set_magnet_setpoints(SC, ord, setpoint, False, 1, dipole_compensation=dipole_compensation) + SC.set_magnet_setpoints(ord, setpoint, False, 1, dipole_compensation=dipole_compensation) elif field == 'SetPointA': # Skew quadrupole - SC = set_magnet_setpoints(SC, ord, setpoint, True, 1) + SC.set_magnet_setpoints(ord, setpoint, True, 1) SC = SC.update_magnets(SC.ORD.Magnet) return SC diff --git a/pySC/correction/orbit_trajectory.py b/pySC/correction/orbit_trajectory.py index 4338343..27fe52b 100644 --- a/pySC/correction/orbit_trajectory.py +++ b/pySC/correction/orbit_trajectory.py @@ -7,7 +7,6 @@ import numpy as np from pySC.core.beam import bpm_reading -from pySC.core.lattice_setting import set_cavity_setpoints, set_cm_setpoints from pySC.utils import logging_tools from pySC.utils.sc_tools import SCrandnc @@ -70,8 +69,8 @@ def SCfeedbackFirstTurn(SC, Mplus, reference=None, CMords=None, BPMords=None, dphi = np.dot(Mplus, (measurement - reference)) lastCMh = _get_last_cm(transmission_history[-1]-1, 1, BPMords, CMords[0])[1][0] lastCMv = _get_last_cm(transmission_history[-1]-1, 1, BPMords, CMords[1])[1][0] - SC = set_cm_setpoints(SC, CMords[0][:lastCMh + 1], -dphi[:lastCMh + 1], skewness=False, method='add') - SC = set_cm_setpoints(SC, CMords[1][:lastCMv + 1], -dphi[len(CMords[0]):len(CMords[0]) + lastCMv + 1], skewness=True, method='add') + SC.set_cm_setpoints(CMords[0][:lastCMh + 1], -dphi[:lastCMh + 1], skewness=False, method='add') + SC.set_cm_setpoints(CMords[1][:lastCMv + 1], -dphi[len(CMords[0]):len(CMords[0]) + lastCMv + 1], skewness=True, method='add') bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging( SC, BPMords=BPMords, ind_history=transmission_history, orb_history=rms_orbit_history) # Inject... @@ -156,8 +155,8 @@ def SCfeedbackStitch(SC, Mplus, reference=None, CMords=None, BPMords=None, nBPMs # Correction step dphi = np.dot(Mplus, (measurement - reference)) - SC = set_cm_setpoints(SC, CMords[0], -dphi[:len(CMords[0])], skewness=False, method="add") - SC = set_cm_setpoints(SC, CMords[1], -dphi[len(CMords[0]):], skewness=True, method="add") + SC.set_cm_setpoints(CMords[0], -dphi[:len(CMords[0])], skewness=False, method="add") + SC.set_cm_setpoints(CMords[1], -dphi[len(CMords[0]):], skewness=True, method="add") bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging( SC, BPMords=BPMords, ind_history=transmission_history, orb_history=rms_orbit_history) @@ -221,8 +220,8 @@ def SCfeedbackBalance(SC, Mplus, reference=None, CMords=None, BPMords=None, eps= # Correction step dphi = np.dot(Mplus, (measurement - reference)) - SC = set_cm_setpoints(SC, CMords[0], -dphi[:len(CMords[0])], skewness=False, method="add") - SC = set_cm_setpoints(SC, CMords[1], -dphi[len(CMords[0]):], skewness=True, method="add") + SC.set_cm_setpoints(CMords[0], -dphi[:len(CMords[0])], skewness=False, method="add") + SC.set_cm_setpoints(CMords[1], -dphi[len(CMords[0]):], skewness=True, method="add") bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging( SC, BPMords=BPMords, ind_history=transmission_history, orb_history=rms_orbit_history) @@ -295,10 +294,10 @@ def SCfeedbackRun(SC, Mplus, reference=None, CMords=None, BPMords=None, eps=1e-4 # Correction step dphi = np.dot(Mplus, (measurement - reference)) if scaleDisp != 0: # TODO this is weight - SC = set_cavity_setpoints(SC, SC.ORD.RF, -scaleDisp * dphi[-1], "Frequency", method="add") + SC.set_cavity_setpoints(SC.ORD.RF, -scaleDisp * dphi[-1], "Frequency", method="add") dphi = dphi[:-1] # TODO the last setpoint is cavity frequency - SC = set_cm_setpoints(SC, CMords[0], -dphi[:len(CMords[0])], skewness=False, method="add") - SC = set_cm_setpoints(SC, CMords[1], -dphi[len(CMords[0]):], skewness=True, method="add") + SC.set_cm_setpoints(CMords[0], -dphi[:len(CMords[0])], skewness=False, method="add") + SC.set_cm_setpoints(CMords[1], -dphi[len(CMords[0]):], skewness=True, method="add") bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging( SC, BPMords=BPMords, ind_history=transmission_history, orb_history=rms_orbit_history) # Inject ... @@ -391,8 +390,8 @@ def _wiggling(SC, BPMords, CMords, transmission_limit, angle_range=(50E-6, 200E- tmpCMordsV = _get_last_cm(transmission_history[-1] - 1, nWiggleCM, BPMords, CMords[1])[0] # Last CMs in vert for i in range(dpts.shape[1]): - SC = set_cm_setpoints(SC, tmpCMordsH, dpts[0, i], skewness=False, method='add') - SC = set_cm_setpoints(SC, tmpCMordsV, dpts[1, i], skewness=True, method='add') + SC.set_cm_setpoints(tmpCMordsH, dpts[0, i], skewness=False, method='add') + SC.set_cm_setpoints(tmpCMordsV, dpts[1, i], skewness=True, method='add') bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging( SC, BPMords=BPMords, ind_history=transmission_history, orb_history=rms_orbit_history) if transmission_history[-1] >= transmission_limit: diff --git a/pySC/correction/rf.py b/pySC/correction/rf.py index 015c01e..2227494 100644 --- a/pySC/correction/rf.py +++ b/pySC/correction/rf.py @@ -10,7 +10,6 @@ from pySC.utils.at_wrapper import findorbit6 from pySC.core.beam import bpm_reading -from pySC.core.lattice_setting import set_cavity_setpoints from pySC.utils import logging_tools from pySC.core.constants import SPEED_OF_LIGHT @@ -31,9 +30,9 @@ def _sin_fit_fun(x, a, b, c): test_vec = lamb * np.linspace(-0.5, 0.5, n_steps) for step in range(n_steps): - SC = set_cavity_setpoints(SC, cav_ords, test_vec[step], 'TimeLag', 'add') + SC.set_cavity_setpoints(cav_ords, test_vec[step], 'TimeLag', 'add') bpm_shift[step], bpm_shift_err[step], mean_tbt_orbit = _get_tbt_energy_shift(SC, bpm_ords) - SC = set_cavity_setpoints(SC, cav_ords, -test_vec[step], 'TimeLag', 'add') + SC.set_cavity_setpoints(cav_ords, -test_vec[step], 'TimeLag', 'add') if plot_progress: _plot_progress((test_vec, bpm_shift, bpm_shift_err), mean_tbt_orbit, phase=True) @@ -55,7 +54,7 @@ def _sin_fit_fun(x, a, b, c): if np.isnan(delta_phi): # TODO does this ever happen? raise RuntimeError('RF phase correction unsuccessful: NaN phase') - SC = set_cavity_setpoints(SC, cav_ords, delta_phi, 'TimeLag', 'add') + SC.set_cavity_setpoints(cav_ords, delta_phi, 'TimeLag', 'add') LOGGER.info(f'Time lag correction step: {delta_phi[0]:.3f} m\n') if plot_results: _plot_results((test_vec, bpm_shift, bpm_shift_err), sol, delta_phi, x_scale=360 / lamb, phase=True) @@ -73,9 +72,9 @@ def correct_rf_frequency(SC, cav_ords=None, bpm_ords=None, n_steps=15, f_range=( # Main loop for step in range(n_steps): - SC = set_cavity_setpoints(SC, cav_ords, test_vec[step], 'Frequency', 'add') + SC.set_cavity_setpoints(cav_ords, test_vec[step], 'Frequency', 'add') bpm_shift[step], bpm_shift_err[step], mean_tbt_orbit = _get_tbt_energy_shift(SC, bpm_ords) - SC = set_cavity_setpoints(SC, cav_ords, -test_vec[step], 'Frequency', 'add') + SC.set_cavity_setpoints(cav_ords, -test_vec[step], 'Frequency', 'add') if plot_progress: _plot_progress((test_vec, bpm_shift, bpm_shift_err), mean_tbt_orbit, phase=False) @@ -89,7 +88,7 @@ def correct_rf_frequency(SC, cav_ords=None, bpm_ords=None, n_steps=15, f_range=( if np.isnan(delta_f): raise RuntimeError('RF frequency correction unsuccessful: NaN frequency') - SC = set_cavity_setpoints(SC, cav_ords, delta_f, 'Frequency', 'add') + SC.set_cavity_setpoints(cav_ords, delta_f, 'Frequency', 'add') LOGGER.info(f'Frequency correction step: {1E-3 * delta_f:.2f} kHz') sol = lambda x: x * param[0] + param[1] if plot_results: diff --git a/pySC/correction/tune.py b/pySC/correction/tune.py index 817f061..b42a816 100644 --- a/pySC/correction/tune.py +++ b/pySC/correction/tune.py @@ -10,7 +10,6 @@ from scipy.optimize import fmin from pySC.core.beam import beam_transmission, plot_transmission -from pySC.core.lattice_setting import set_magnet_setpoints from pySC.utils import logging_tools from pySC.utils.at_wrapper import atlinopt @@ -32,10 +31,10 @@ def tune_scan(SC, quad_ords, rel_quad_changes, target=1, n_points=60, do_plot=Fa ords = np.hstack(quad_ords) for q1, q2 in inds.T: q_setpoints = np.hstack((np.ones(nq[0]) * rel_quad_changes[0][q1], np.ones(nq[1]) * rel_quad_changes[1][q2])) - SC = set_magnet_setpoints(SC, ords, q_setpoints, False, 1, method='rel') + SC.set_magnet_setpoints(ords, q_setpoints, False, 1, method='rel') max_turns[q1, q2], surviving_fraction = beam_transmission(SC, nParticles=nParticles, nTurns=nTurns) transmission[q1, q2, :] = 1 * surviving_fraction - SC = set_magnet_setpoints(SC, ords, 1 / q_setpoints, False, 1, method='rel') + SC.set_magnet_setpoints(ords, 1 / q_setpoints, False, 1, method='rel') if do_plot: f, ax = plot_scan(transmission[:, :, -1], max_turns, first_quads, rel_quad_changes) @@ -48,7 +47,7 @@ def tune_scan(SC, quad_ords, rel_quad_changes, target=1, n_points=60, do_plot=Fa LOGGER.info(f'Transmission target reached with:\n' f' {first_quads[0]} SetPoint: {setpoints[0]:.4f}\n' f' {first_quads[1]} SetPoint: {setpoints[1]:.4f}') - SC = set_magnet_setpoints(SC, ords, q_setpoints, False, 1, method='rel') + SC.set_magnet_setpoints(ords, q_setpoints, False, 1, method='rel') return SC, setpoints, max_turns, transmission testTrans = np.zeros(n_points) @@ -71,7 +70,7 @@ def tune_scan(SC, quad_ords, rel_quad_changes, target=1, n_points=60, do_plot=Fa f' {first_quads[0]} SetPoint: {setpoints[0]:.4f}\n' f' {first_quads[1]} SetPoint: {setpoints[0]:.4f}') q_setpoints = np.hstack((setpoints[0] * np.ones(nq[0]), setpoints[1] * np.ones(nq[1]))) - SC = set_magnet_setpoints(SC, ords, q_setpoints, False, 1, method='rel') + SC.set_magnet_setpoints(ords, q_setpoints, False, 1, method='rel') return SC, setpoints, max_turns, transmission @@ -118,13 +117,13 @@ def fit_tune(SC, q_ords, target_tune=None, xtol=1E-4, ftol=1E-3, fit_integer=Tru SP0[nFam][n] = SC.RING[q_ords[nFam][n]].SetPointB[1] fun = lambda x: _fit_tune_fun(SC, q_ords, x, SP0, target_tune, fit_integer) sol = fmin(fun, xtol=xtol, ftol=ftol) - SC = set_magnet_setpoints(SC, q_ords, sol + SP0, False, 1, method='abs', dipole_compensation=True) + SC.set_magnet_setpoints(q_ords, sol + SP0, False, 1, method='abs', dipole_compensation=True) LOGGER.debug(f' Final tune: [{tune(SC, fit_integer)}]\n Setpoints change: [{sol}]') return SC def _fit_tune_fun(SC, q_ords, setpoints, init_setpoints, target, fit_integer): - SC = set_magnet_setpoints(SC, q_ords, setpoints + init_setpoints, False, 1, method='abs', dipole_compensation=True) + SC.set_magnet_setpoints(q_ords, setpoints + init_setpoints, False, 1, method='abs', dipole_compensation=True) nu = tune(SC, fit_integer) return np.sqrt(np.mean((nu - target) ** 2)) diff --git a/pySC/example.py b/pySC/example.py index 542d20b..640d370 100644 --- a/pySC/example.py +++ b/pySC/example.py @@ -14,7 +14,7 @@ from pySC.plotting.plot_phase_space import plot_phase_space from pySC.plotting.plot_support import plot_support from pySC.plotting.plot_lattice import plot_lattice -from pySC.core.lattice_setting import set_magnet_setpoints, switch_cavity_and_radiation +from pySC.core.lattice_setting import switch_cavity_and_radiation from pySC.correction.rf import correct_rf_phase, correct_rf_frequency, phase_and_energy_error from pySC.utils import logging_tools @@ -108,7 +108,7 @@ def _marker(name): SC.RING = switch_cavity_and_radiation(SC.RING, 'cavityoff') sextOrds = SCgetOrds(SC.RING, 'SF|SD') - SC = set_magnet_setpoints(SC, sextOrds, 0.0, False, 2, method='abs') + SC.set_magnet_setpoints(sextOrds, 0.0, False, 2, method='abs') RM1 = SCgetModelRM(SC, SC.ORD.BPM, SC.ORD.CM, nTurns=1) RM2 = SCgetModelRM(SC, SC.ORD.BPM, SC.ORD.CM, nTurns=2) Minv1 = SCgetPinv(RM1, alpha=50) @@ -128,7 +128,7 @@ def _marker(name): # Turning on the sextupoles for rel_setting in np.linspace(0.1, 1, 5): - SC = set_magnet_setpoints(SC, sextOrds, rel_setting, False, 2, method='rel') + SC.set_magnet_setpoints(sextOrds, rel_setting, False, 2, method='rel') try: SC = SCfeedbackBalance(SC, Minv2, maxsteps=32, eps=eps) except RuntimeError: diff --git a/pySC/lattice_properties/response_measurement.py b/pySC/lattice_properties/response_measurement.py index 8bc0a20..1e37652 100644 --- a/pySC/lattice_properties/response_measurement.py +++ b/pySC/lattice_properties/response_measurement.py @@ -1,7 +1,6 @@ import numpy as np from pySC.core.beam import bpm_reading -from pySC.core.lattice_setting import set_cm_setpoints, get_cm_setpoints, set_cavity_setpoints from pySC.utils import logging_tools from pySC.utils.at_wrapper import atgetfieldvalues @@ -27,7 +26,7 @@ def response_matrix(SC, amp, bpm_ords, cm_ords, mode='fixedKick', n_steps=2, fit raise ValueError('No closed orbit found.') i = 0 for n_dim in range(2): - cmstart = get_cm_setpoints(SC, cm_ords[n_dim], bool(n_dim)) + cmstart = SC.get_cm_setpoints(cm_ords[n_dim], bool(n_dim)) for nCM in range(len(cm_ords[n_dim])): max_step, gradient = _kick_amplitude(SC, bref, bpm_ords, cm_ords[n_dim][nCM], amp[n_dim][nCM], bool(n_dim), mode) cm_step_vec = np.linspace(-max_step, max_step, n_steps) @@ -36,8 +35,8 @@ def response_matrix(SC, amp, bpm_ords, cm_ords, mode='fixedKick', n_steps=2, fit gradient = np.vstack((np.zeros((n_steps - 1, len(bref))), gradient.T)) for n_step in range(n_steps): if cm_step_vec[n_step] != 0 and cm_step_vec[n_step] != max_step: - SC = set_cm_setpoints(SC, cm_ords[n_dim][nCM], cmstart[nCM] + cm_step_vec[n_step], skewness=bool(n_dim)) - real_cm_setpoint[n_step] = get_cm_setpoints(SC, cm_ords[n_dim][nCM], skewness=bool(n_dim)) + SC.set_cm_setpoints(cm_ords[n_dim][nCM], cmstart[nCM] + cm_step_vec[n_step], skewness=bool(n_dim)) + real_cm_setpoint[n_step] = SC.get_cm_setpoints(cm_ords[n_dim][nCM], skewness=bool(n_dim)) gradient[n_step, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - bref dCM = real_cm_setpoint - cmstart[nCM] cm_steps[n_dim][:, nCM] = dCM @@ -51,7 +50,7 @@ def response_matrix(SC, amp, bpm_ords, cm_ords, mode='fixedKick', n_steps=2, fit rm[nBPM, i] = np.polyfit(x, y, fit_order)[fit_order - 1] error[nBPM, i] = np.sqrt(np.mean((rm[nBPM, i] * x - y).T ** 2)) i = i + 1 - SC = set_cm_setpoints(SC, cm_ords[n_dim][nCM], cmstart[nCM], skewness=bool(n_dim)) + SC.set_cm_setpoints(cm_ords[n_dim][nCM], cmstart[nCM], skewness=bool(n_dim)) rm[np.isnan(rm)] = 0 LOGGER.debug(' done.') return rm, error, cm_steps @@ -61,9 +60,9 @@ def dispersion(SC, rf_step, bpm_ords=None, cav_ords=None, n_steps=2): bpm_ords, cav_ords = _check_ords(SC, bpm_ords, cav_ords) bref = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) if n_steps == 2: - SC = set_cavity_setpoints(SC, cav_ords, rf_step, 'Frequency', 'add') + SC.set_cavity_setpoints(cav_ords, rf_step, 'Frequency', 'add') B = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - SC = set_cavity_setpoints(SC, cav_ords, -rf_step, 'Frequency', 'add') + SC.set_cavity_setpoints(cav_ords, -rf_step, 'Frequency', 'add') return (B - bref) / rf_step rf_steps = np.zeros((len(cav_ords), n_steps)) for n_cav, cav_ord in enumerate(cav_ords): @@ -71,9 +70,9 @@ def dispersion(SC, rf_step, bpm_ords=None, cav_ords=None, n_steps=2): dB = np.zeros((n_steps, *np.shape(bref))) rf0 = atgetfieldvalues(SC.RING, cav_ords, "FrequencySetPoint") for nStep in range(n_steps): - SC = set_cavity_setpoints(SC, cav_ords, rf_steps[:, nStep], 'Frequency', 'abs') + SC.set_cavity_setpoints(cav_ords, rf_steps[:, nStep], 'Frequency', 'abs') dB[nStep, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - bref - SC = set_cavity_setpoints(SC, cav_ords, rf0, 'Frequency', 'abs') + SC.set_cavity_setpoints(cav_ords, rf0, 'Frequency', 'abs') return np.linalg.lstsq(np.linspace(-rf_step, rf_step, n_steps), dB)[0] @@ -86,7 +85,7 @@ def _check_ords(SC, bpm_ords, cav_ords): def _kick_amplitude(SC, bref, bpm_ords, cm_ord, amp, skewness: bool, mode): - cmstart = get_cm_setpoints(SC, cm_ord, skewness) + cmstart = SC.get_cm_setpoints(cm_ord, skewness) max_step = amp params = dict(fixedOffset=(4, 0.5), fixedKick=(20, 0.9)) n_iter, decrease_factor = params[mode] @@ -107,10 +106,10 @@ def _kick_amplitude(SC, bref, bpm_ords, cm_ord, amp, skewness: bool, mode): def _try_setpoint(SC, bpm_ords, cm_ord, cmstart, max_step, skewness): - SC = set_cm_setpoints(SC, cm_ord, cmstart + max_step, skewness) - real_cm_setpoint = get_cm_setpoints(SC, cm_ord, skewness) + SC.set_cm_setpoints(cm_ord, cmstart + max_step, skewness) + real_cm_setpoint = SC.get_cm_setpoints(cm_ord, skewness) if real_cm_setpoint != (cmstart + max_step): LOGGER.debug('CM clipped. Using different CM direction.') max_step *= -1 - SC = set_cm_setpoints(SC, cm_ord, cmstart + max_step, skewness) + SC.set_cm_setpoints(cm_ord, cmstart + max_step, skewness) return SC, max_step, np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) diff --git a/pySC/matlab_index.py b/pySC/matlab_index.py index 656e42e..b7a04f4 100644 --- a/pySC/matlab_index.py +++ b/pySC/matlab_index.py @@ -16,9 +16,9 @@ from pySC.core.beam import beam_transmission, bpm_reading, generate_bunches from pySC.core.simulated_commissioning import SimulatedCommissioning + from pySC.core.classes import DotDict -from pySC.core.lattice_setting import (set_cavity_setpoints, set_magnet_setpoints, - set_cm_setpoints, get_cm_setpoints, switch_cavity_and_radiation) +from pySC.core.lattice_setting import switch_cavity_and_radiation from pySC.correction.bba import SCBBA as bba from pySC.correction.injection_fit import fit_injection_trajectory, fit_injection_drift from pySC.correction.orbit_trajectory import SCfeedbackFirstTurn as first_turn, SCfeedbackStitch as stitch, \ @@ -118,7 +118,7 @@ def SCgetCMSetPoints(SC: SimulatedCommissioning, CMords: ndarray, nDim: int) -> "Transition concerns nDim 1 -> horizontal, 2-> vertical.") if nDim not in (1, 2): raise ValueError("Function expects nDim 1 (hor) or 2 (ver)") - return get_cm_setpoints(SC, ords=CMords, skewness=(nDim == 2)) + return SC.get_cm_setpoints(ords=CMords, skewness=(nDim == 2)) def SCgetCOD(SC, /, *, ords=None, plot=False): @@ -266,7 +266,8 @@ def SCscaleCircumference(RING, circ, /, *, mode='abs'): def SCsetCavs2SetPoints(SC: SimulatedCommissioning, CAVords: ndarray, type: str, setpoints: ndarray, /, *, mode: str = 'abs') -> SimulatedCommissioning: - return set_cavity_setpoints(SC, ords=CAVords, setpoints=setpoints, param=type, method=mode) + SC.set_cavity_setpoints(ords=CAVords, setpoints=setpoints, param=type, method=mode) + return SC def SCsetCMs2SetPoints(SC: SimulatedCommissioning, CMords: ndarray, setpoints: ndarray, nDim: int, /, *, @@ -275,8 +276,8 @@ def SCsetCMs2SetPoints(SC: SimulatedCommissioning, CMords: ndarray, setpoints: n "Transition concerns nDim 1 -> horizontal, 2-> vertical.") if nDim not in (1, 2): raise ValueError("Function expects nDim 1 (hor) or 2 (ver)") - SC = set_cm_setpoints(SC, ords=CMords, setpoints=setpoints, skewness=(nDim == 2), method=mode) - return SC, get_cm_setpoints(SC, ords=CMords, skewness=(nDim == 2)) + SC.set_cm_setpoints(ords=CMords, setpoints=setpoints, skewness=(nDim == 2), method=mode) + return SC, SC.get_cm_setpoints(ords=CMords, skewness=(nDim == 2)) def SCsetMags2SetPoints(SC: SimulatedCommissioning, MAGords: ndarray, type: int, order: int, setpoints: ndarray, /, *, @@ -286,10 +287,9 @@ def SCsetMags2SetPoints(SC: SimulatedCommissioning, MAGords: ndarray, type: int, " order 1 -> dipole, 2-> quadrupole, 3-> sextupole, ...") if type not in (1, 2): raise ValueError("Function expects type 1 (skew) or 2 (normal)") - - return set_magnet_setpoints(SC, ords=MAGords, setpoints=setpoints, skewness=(type == 1), order=order - 1, - method=method, dipole_compensation=dipCompensation) - + SC.set_magnet_setpoints(ords=MAGords, setpoints=setpoints, skewness=(type == 1), order=order - 1, + method=method, dipole_compensation=dipCompensation) + return SC def SCsetMultipoles(SC: SimulatedCommissioning, ords: ndarray, AB, /, *, method: str = 'rnd', order: int = None, type: int = None): LOGGER.warn("Function SCsetMultipoles contains non-trivial transition between Matlab and Python code.\n" diff --git a/pySC/plotting/plot_lattice.py b/pySC/plotting/plot_lattice.py index abb98a8..f05a587 100644 --- a/pySC/plotting/plot_lattice.py +++ b/pySC/plotting/plot_lattice.py @@ -2,7 +2,6 @@ from matplotlib import pyplot as plt from numpy import ndarray import warnings -from pySC.core.lattice_setting import get_cm_setpoints from pySC.core.simulated_commissioning import SimulatedCommissioning from pySC.plotting.plot_synoptic import plot_synoptic from pySC.plotting.plot_apertures import plot_data_apertures, plot_data_beta_disp, baseplot @@ -100,7 +99,7 @@ def plot_cm_strengths(SC: SimulatedCommissioning): f, ax = plt.subplots(nrows=2, num=86, figsize=(9, 7.5), facecolor="w") s_pos = findspos(SC.RING) for n_dim in range(2): - setpoints = 1E6 * get_cm_setpoints(SC, SC.ORD.CM[n_dim], skewness=bool(n_dim)) + setpoints = 1E6 * SC.get_cm_setpoints(SC.ORD.CM[n_dim], skewness=bool(n_dim)) count, bins_count = np.histogram(setpoints, bins=len(setpoints)) ax[0].bar(s_pos[SC.ORD.CM[n_dim]], setpoints) ax[1].plot(bins_count[1:], np.cumsum(count / np.sum(count)), lw=3) diff --git a/pySC/utils/classdef_tools.py b/pySC/utils/classdef_tools.py index 18d32e3..782fd14 100644 --- a/pySC/utils/classdef_tools.py +++ b/pySC/utils/classdef_tools.py @@ -1,5 +1,6 @@ import numpy as np +from pySC.core.constants import SETTING_METHODS from pySC.utils.sc_tools import SCrandnc @@ -48,3 +49,13 @@ def s_interpolation(off, s, ord1, f1, ord2, f2): off[ind1] = np.interp(C + s[ind1], np.array([s1[n], s2[n] + C]), np.array([f1[n], f2[n]])) off[ord1[n]:] = np.interp(s[ord1[n]:], np.array([s1[n], s2[n] + C]), np.array([f1[n], f2[n]])) return off + + +def check_input_and_setpoints(method, ords, setpoints): + if method not in SETTING_METHODS: + raise ValueError(f'Unsupported setpoint method: {method}. Allowed options are: {SETTING_METHODS}.') + ords_1d = np.ravel(np.array([ords], dtype=int)) + setpoints_1d = np.ravel(np.array([setpoints])) + if len(setpoints_1d) not in (1, len(ords_1d)): + raise ValueError(f'Setpoints have to have length of 1 or matching to the length or ordinates.') + return ords_1d, (np.repeat(setpoints_1d, len(ords_1d)) if len(setpoints_1d) == 1 else setpoints_1d) diff --git a/tests/test_example.py b/tests/test_example.py index baee93a..3582e21 100644 --- a/tests/test_example.py +++ b/tests/test_example.py @@ -8,7 +8,7 @@ from pySC.correction.tune import tune_scan from pySC.lattice_properties.response_model import SCgetModelRM, SCgetModelDispersion from pySC.utils.sc_tools import SCgetOrds, SCgetPinv -from pySC.core.lattice_setting import set_magnet_setpoints, switch_cavity_and_radiation +from pySC.core.lattice_setting import switch_cavity_and_radiation from pySC.correction.rf import correct_rf_phase, correct_rf_frequency @@ -66,7 +66,7 @@ def test_example(at_lattice): sc.RING = switch_cavity_and_radiation(sc.RING, 'cavityoff') sext_ords = SCgetOrds(sc.RING, 'SF|SD') - sc = set_magnet_setpoints(sc, sext_ords, 0.0, False, 2, method='abs') + sc.set_magnet_setpoints(sext_ords, 0.0, False, 2, method='abs') rm1 = SCgetModelRM(sc, sc.ORD.BPM, sc.ORD.CM, nTurns=1) rm2 = SCgetModelRM(sc, sc.ORD.BPM, sc.ORD.CM, nTurns=2) minv1 = SCgetPinv(rm1, alpha=50) @@ -84,7 +84,7 @@ def test_example(at_lattice): # Turning on the sextupoles for rel_setting in np.linspace(0.1, 1, 5): - sc = set_magnet_setpoints(sc, sext_ords, rel_setting, False, 2, method='rel') + sc.set_magnet_setpoints(sext_ords, rel_setting, False, 2, method='rel') sc = SCfeedbackBalance(sc, minv2, maxsteps=32, eps=eps) sc.RING = switch_cavity_and_radiation(sc.RING, 'cavityon') diff --git a/tests/test_lattice_setting.py b/tests/test_lattice_setting.py index 88a050f..a9cd02f 100644 --- a/tests/test_lattice_setting.py +++ b/tests/test_lattice_setting.py @@ -6,7 +6,7 @@ from tests.test_at_wrapper import at_lattice from pySC.utils.sc_tools import SCgetOrds from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.core.lattice_setting import set_cm_setpoints, set_magnet_setpoints, set_cavity_setpoints, get_cm_setpoints, _check_input_and_setpoints +from pySC.utils.classdef_tools import check_input_and_setpoints from pySC.core.constants import SETTING_ABS, SETTING_REL, SETTING_ADD @@ -14,7 +14,7 @@ def test_set_cm_setpoints_side_effects(sc): indices = np.arange(11, 450, 22, dtype=int) setpoints = 1e-4 * np.ones(len(indices)) for _ in range(2): - sc = set_cm_setpoints(sc, indices, setpoints, True, method="add") + sc.set_cm_setpoints(indices, setpoints, True, method="add") assert_equal(indices, np.arange(11, 450, 22, dtype=int)) assert_equal(setpoints, 1e-4 * np.ones(len(indices))) @@ -26,7 +26,7 @@ def test_set_magnet_setpoints_side_effects(sc): indices = np.arange(11, 450, 22, dtype=int) setpoints = 1e-4 * np.ones(len(indices)) for _ in range(2): - sc = set_magnet_setpoints(sc, indices, setpoints, True, 1, method="add") + sc.set_magnet_setpoints(indices, setpoints, True, 1, method="add") assert_equal(indices, np.arange(11, 450, 22, dtype=int)) assert_equal(setpoints, 1e-4 * np.ones(len(indices))) assert sc.RING.__repr__() != lattice_copy.__repr__() @@ -39,7 +39,7 @@ def test_set_cavity_setpoints_side_effects(sc): indices = np.zeros(1, dtype=int) setpoints = np.ones(1) for _ in range(2): - sc = set_cavity_setpoints(sc, indices, setpoints, "Frequency", method="add") + sc.set_cavity_setpoints(indices, setpoints, "Frequency", method="add") assert_equal(indices, np.zeros(1, dtype=int)) assert_equal(setpoints, np.ones(1)) assert sc.RING.__repr__() != lattice_copy.__repr__() @@ -49,35 +49,35 @@ def test_check_input_and_setpoints(): wanted_ord, wanted_setpoint = np.array([2], dtype=int), np.array([3.5]) for o, s in zip((2, [2], np.array([2])), (3.5, np.array([3.5]), [3.5])): - ord1d, setp1d = _check_input_and_setpoints(SETTING_ABS, o, s) + ord1d, setp1d = check_input_and_setpoints(SETTING_ABS, o, s) assert_equal(ord1d, wanted_ord) assert_equal(setp1d, wanted_setpoint) wanted_ords, wanted_setpoints = np.array([2, 6], dtype=int), np.array([3.5, 3.5]) for os, ss in zip(((2, 6), [2, 6], np.array([2, 6])), (3.5, np.array([3.5, 3.5]), [3.5])): - ords1d, setps1d = _check_input_and_setpoints(SETTING_ABS, os, ss) + ords1d, setps1d = check_input_and_setpoints(SETTING_ABS, os, ss) assert_equal(ords1d, wanted_ords) assert_equal(setps1d, wanted_setpoints) def test_set_magnet_setpoints(unit_sc): - unit_sc = set_magnet_setpoints(unit_sc, [2, 3], 1.1, False, 1, method=SETTING_REL) + unit_sc.set_magnet_setpoints([2, 3], 1.1, False, 1, method=SETTING_REL) assert_equal(unit_sc.RING[2].PolynomB, np.array([0, 1.32])) assert_equal(unit_sc.RING[3].PolynomB, np.array([0, -1.32])) - unit_sc = set_magnet_setpoints(unit_sc, 2, 0.1, False, 1, method=SETTING_ADD) + unit_sc.set_magnet_setpoints(2, 0.1, False, 1, method=SETTING_ADD) assert_allclose(unit_sc.RING[2].PolynomB, np.array([0, 1.42])) assert_equal(unit_sc.RING[3].PolynomB, np.array([0, -1.32])) def test_set_cm_setpoints(unit_sc): - unit_sc = set_cm_setpoints(unit_sc, [2, 3], 1e-4, False, method=SETTING_ABS) + unit_sc.set_cm_setpoints([2, 3], 1e-4, False, method=SETTING_ABS) assert_equal(unit_sc.RING[2].PolynomB, np.array([-2e-4, 1.2])) # 2 due to 0.5 m long assert_equal(unit_sc.RING[3].PolynomB, np.array([-2e-4, -1.2])) - unit_sc = set_cm_setpoints(unit_sc, [2, 3], 1e-1, False, method=SETTING_ABS) # clipping at 1 mrad + unit_sc.set_cm_setpoints([2, 3], 1e-1, False, method=SETTING_ABS) # clipping at 1 mrad assert_equal(unit_sc.RING[2].PolynomB, np.array([-2e-3, 1.2])) assert_equal(unit_sc.RING[3].PolynomB, np.array([-2e-3, -1.2])) - unit_sc = set_cm_setpoints(unit_sc, 2, 1e-4, True, method=SETTING_ADD) + unit_sc.set_cm_setpoints(2, 1e-4, True, method=SETTING_ADD) assert_equal(unit_sc.RING[2].PolynomB, np.array([-2e-3, 1.2])) # PolynomB unchanged assert_equal(unit_sc.RING[3].PolynomB, np.array([-2e-3, -1.2])) assert_equal(unit_sc.RING[2].PolynomA, np.array([2e-4, 0])) @@ -118,4 +118,4 @@ def sc(at_lattice): VoltageOffset=5E3, TimeLagOffset=0.5) SC.apply_errors() - return SC \ No newline at end of file + return SC