From e003dffd46082e96435a9659c1af3f561caec1e1 Mon Sep 17 00:00:00 2001 From: lmalina Date: Wed, 24 Jan 2024 17:12:11 +0100 Subject: [PATCH] Utils: imports of the whole modules and name changes (#52) --- pySC/core/beam.py | 30 +++--- pySC/core/simulated_commissioning.py | 96 +++++++++---------- pySC/correction/bba.py | 25 +++-- pySC/correction/injection_fit.py | 12 +-- pySC/correction/loco.py | 5 +- pySC/correction/orbit_trajectory.py | 11 +-- pySC/correction/ramp_errors.py | 5 +- pySC/correction/rf.py | 5 +- pySC/example.py | 44 ++++----- pySC/lattice_properties/apertures.py | 25 +++-- pySC/lattice_properties/magnet_orbit.py | 8 +- .../response_measurement.py | 6 +- pySC/lattice_properties/response_model.py | 11 +-- pySC/matlab_index.py | 17 ++-- pySC/plotting/plot_apertures.py | 6 +- pySC/plotting/plot_lattice.py | 8 +- pySC/plotting/plot_phase_space.py | 10 +- pySC/plotting/plot_support.py | 22 ++--- pySC/plotting/plot_synoptic.py | 4 +- pySC/utils/classdef_tools.py | 6 +- pySC/utils/sc_tools.py | 18 ++-- pySC/utils/stats.py | 2 +- tests/test_beam.py | 4 +- tests/test_example.py | 36 +++---- tests/test_injection_fit.py | 4 +- tests/test_lattice_setting.py | 21 ++-- tests/test_loco.py | 21 ++-- tests/test_pinv.py | 8 +- tests/test_plotting.py | 18 ++-- tests/test_simulated_commissioning.py | 10 +- 30 files changed, 240 insertions(+), 258 deletions(-) diff --git a/pySC/core/beam.py b/pySC/core/beam.py index b13712e..8299e6d 100644 --- a/pySC/core/beam.py +++ b/pySC/core/beam.py @@ -5,6 +5,7 @@ This module contains the 'beam-based' functions to interact with lattice under study. """ from typing import Tuple +import warnings import numpy as np from matplotlib import pyplot as plt @@ -12,10 +13,7 @@ from pySC.core.simulated_commissioning import SimulatedCommissioning from pySC.core.constants import TRACK_ORB, TRACK_PORB, TRACK_TBT -from pySC.utils.sc_tools import SCrandnc -from pySC.utils.at_wrapper import atgetfieldvalues, atpass, findorbit6, findspos, patpass -import warnings -from pySC.utils import logging_tools +from pySC.utils import at_wrapper, logging_tools, sc_tools warnings.filterwarnings("ignore", message='Mean of empty slice') LOGGER = logging_tools.get_logger(__name__) @@ -130,7 +128,7 @@ def beam_transmission(SC: SimulatedCommissioning, nParticles: int = None, nTurns if nTurns is None: nTurns = SC.INJ.nTurns LOGGER.debug(f'Calculating maximum beam transmission for {nParticles} particles and {nTurns} turns: ') - T = patpass(SC.RING, generate_bunches(SC, nParticles=nParticles), nTurns, np.array([len(SC.RING)]), keep_lattice=False) + T = at_wrapper.patpass(SC.RING, generate_bunches(SC, nParticles=nParticles), nTurns, np.array([len(SC.RING)]), keep_lattice=False) fraction_survived = np.mean(~np.isnan(T[0, :, :, :]), axis=(0, 1)) max_turns = np.sum(fraction_survived > 1 - SC.INJ.beamLostAt) if plot: @@ -161,10 +159,10 @@ def generate_bunches(SC: SimulatedCommissioning, nParticles=None) -> ndarray: """ if nParticles is None: nParticles = SC.INJ.nParticles - Z = np.tile(np.transpose(SC.INJ.randomInjectionZ * SCrandnc(2, (1, 6)) + SC.INJ.Z0), nParticles) + Z = np.tile(np.transpose(SC.INJ.randomInjectionZ * sc_tools.randnc(2, (1, 6)) + SC.INJ.Z0), nParticles) if nParticles != 1: V, L = np.linalg.eig(SC.INJ.beamSize) - Z += np.diag(np.sqrt(V)) @ L @ SCrandnc(3, (6, nParticles)) + Z += np.diag(np.sqrt(V)) @ L @ sc_tools.randnc(3, (6, nParticles)) return SC.INJ.postFun(Z) @@ -183,12 +181,12 @@ def plot_transmission(ax, fraction_survived, n_turns, beam_lost_at): def _real_bpm_reading(SC, track_mat, bpm_inds=None): # track_mat should be only x,y over all particles only at BPM positions n_bpms, nTurns = track_mat.shape[2:] bpm_ords = SC.ORD.BPM if bpm_inds is None else SC.ORD.BPM[bpm_inds] - bpm_noise = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, ('NoiseCO' if SC.INJ.trackMode == 'ORB' else "Noise"))) - bpm_noise = bpm_noise[:, :, np.newaxis] * SCrandnc(2, (2, n_bpms, nTurns)) - bpm_offset = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, 'Offset') + atgetfieldvalues(SC.RING, bpm_ords, 'SupportOffset')) - bpm_cal_error = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, 'CalError')) - bpm_roll = np.squeeze(atgetfieldvalues(SC.RING, bpm_ords, 'Roll') + atgetfieldvalues(SC.RING, bpm_ords, 'SupportRoll'), axis=1) - bpm_sum_error = np.transpose(atgetfieldvalues(SC.RING, bpm_ords, 'SumError'))[:, np.newaxis] * SCrandnc(2, (n_bpms, nTurns)) + bpm_noise = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, ('NoiseCO' if SC.INJ.trackMode == 'ORB' else "Noise"))) + bpm_noise = bpm_noise[:, :, np.newaxis] * sc_tools.randnc(2, (2, n_bpms, nTurns)) + bpm_offset = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'Offset') + at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SupportOffset')) + bpm_cal_error = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'CalError')) + bpm_roll = np.squeeze(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'Roll') + at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SupportRoll'), axis=1) + bpm_sum_error = np.transpose(at_wrapper.atgetfieldvalues(SC.RING, bpm_ords, 'SumError'))[:, np.newaxis] * sc_tools.randnc(2, (n_bpms, nTurns)) # averaging the X and Y positions at BPMs over particles mean_orbit = np.nanmean(track_mat, axis=1) transmission = np.mean(~np.isnan(track_mat[0, :, :, :]), axis=0) * (1 + bpm_sum_error) @@ -211,9 +209,9 @@ def _tracking(SC: SimulatedCommissioning, refs: ndarray) -> ndarray: # lattice_pass output: (6, N, R, T) coordinates of N particles at R reference points for T turns. # findorbit second output value: (R, 6) closed orbit vector at each specified location if SC.INJ.trackMode == TRACK_ORB: - pos = np.transpose(findorbit6(SC.RING, refs, keep_lattice=False)[1])[[0, 2], :].reshape(2, 1, len(refs), 1) + pos = np.transpose(at_wrapper.findorbit6(SC.RING, refs, keep_lattice=False)[1])[[0, 2], :].reshape(2, 1, len(refs), 1) else: - pos = atpass(SC.RING, generate_bunches(SC), SC.INJ.nTurns, refs, keep_lattice=False)[[0, 2], :, :, :] + pos = at_wrapper.atpass(SC.RING, generate_bunches(SC), SC.INJ.nTurns, refs, keep_lattice=False)[[0, 2], :, :, :] pos[1, np.isnan(pos[0, :, :, :])] = np.nan return pos @@ -235,7 +233,7 @@ def _reshape_3d_to_matlab_like_2d(mean_bpm_orbits_3d: ndarray) -> ndarray: def _plot_bpm_reading(SC, bpm_orbits_3d, bpm_inds=None, all_readings_5d=None): ap_ords, apers = _get_ring_aperture(SC) fig, ax = plt.subplots(num=1, nrows=2, ncols=1, sharex="all", figsize=(8, 6), dpi=100, facecolor="w") - s_pos = findspos(SC.RING) + s_pos = at_wrapper.findspos(SC.RING) circumference = s_pos[-1] bpms = SC.ORD.BPM if bpm_inds is None else SC.ORD.BPM[bpm_inds] if all_readings_5d is not None: diff --git a/pySC/core/simulated_commissioning.py b/pySC/core/simulated_commissioning.py index c956754..119d9ad 100644 --- a/pySC/core/simulated_commissioning.py +++ b/pySC/core/simulated_commissioning.py @@ -17,11 +17,7 @@ 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, SETTING_ABS, SETTING_REL, SETTING_ADD, NUM_TO_AB, SETPOINT) -from pySC.utils import logging_tools -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 +from pySC.utils import at_wrapper, classdef_tools, logging_tools, sc_tools LOGGER = logging_tools.get_logger(__name__) @@ -103,7 +99,7 @@ def register_bpms(self, ords: ndarray, **kwargs): Examples: Identify the ordinates of all elements named `BPM` and registers them as BPMs in `SC`:: - ords = SCgetOrds(SC.RING,'BPM') + ords = sc_tools.ords_from_regex(SC.RING,'BPM') SC.register_bpms(ords) Register the BPMs specified in `ords` in `SC` and set the uncertainty of the offset to `500um` in @@ -124,7 +120,7 @@ def register_bpms(self, ords: ndarray, **kwargs): SC.register_bpms(ords, Offset=500E-6*np.ones(2), SumError=0.2) See also: - *bpm_reading*, *SCgetOrds*, *SC.verify_structure*, *SC.apply_errors*, *SC.register_support*, *SC.update_support* + *bpm_reading*, *ords_from_regex*, *SC.verify_structure*, *SC.apply_errors*, *SC.register_support*, *SC.update_support* """ self._check_kwargs(kwargs, BPM_ERROR_FIELDS) self.ORD.BPM = np.unique(np.concatenate((self.ORD.BPM, ords))) @@ -166,7 +162,7 @@ def register_cavities(self, ords: ndarray, **kwargs): Examples: Identify the ordinates of all elements named `'CAV'` and register them as cavities in `SC`:: - ords = SCgetOrds(SC.RING, 'CAV') + ords = sc_tools.ords_from_regex(SC.RING, 'CAV') SC.register_cavities(ords) Register the cavities specified in `ords` in `SC` and sets the uncertainty of the frequency offset @@ -181,7 +177,7 @@ def register_cavities(self, ords: ndarray, **kwargs): SC.register_cavities(ords, FrequencyOffset=1E3, TimeLagOffset=0.3) See also: - *SCgetOrds*, *SC.verify_structure*, *SC.apply_errors* + *ords_from_regex*, *SC.verify_structure*, *SC.apply_errors* """ self._check_kwargs(kwargs, RF_ERROR_FIELDS) @@ -247,7 +243,7 @@ def register_magnets(self, ords: ndarray, **kwargs): # TODO docstring is too lo Examples: Identify the ordinates of all elements named `QF` and register them in `SC`:: - ords = SCgetOrds(SC.RING, 'QF') + ords = sc_tools.ords_from_regex(SC.RING, 'QF') SC.register_magnets(ords) Register the magnets specified in `ords` in `SC` and set the uncertainty of @@ -282,8 +278,8 @@ def register_magnets(self, ords: ndarray, **kwargs): # TODO docstring is too lo columns of `childOrds`. The uncertainty of the bending angle is set to 1E-4:: - masterOrds = SCgetOrds(SC.RING,'BENDa') - childOrds = numpy.vstack((SCgetOrds(SC.RING,'BENDb'), SCgetOrds(SC.RING,'BENDc'))) + masterOrds = sc_tools.ords_from_regex(SC.RING,'BENDa') + childOrds = numpy.vstack((sc_tools.ords_from_regex(SC.RING,'BENDb'), sc_tools.ords_from_regex(SC.RING,'BENDc'))) SC.register_magnets(masterOrds, BendingAngle=1E-4, MasterOf=childOrds) Register the magnets specified in `ords` in `SC` as combined function magnets @@ -314,7 +310,7 @@ def register_magnets(self, ords: ndarray, **kwargs): # TODO docstring is too lo CalErrorA=np.array([1E-2, 2E-3, 0])) See Also: - *SCgetOrds*, *SC.update_magnets*, *SC.verify_structure*, *SC.apply_errors*, *SC.register_support* + *ords_from_regex*, *SC.update_magnets*, *SC.verify_structure*, *SC.apply_errors*, *SC.register_support* """ self._check_kwargs(kwargs, MAGNET_TYPE_FIELDS + MAGNET_ERROR_FIELDS) @@ -418,7 +414,7 @@ def register_supports(self, support_ords: ndarray, support_type: str, **kwargs): SC.register_support(ords, "Girder", Offset=np.array([dX dY dZ]), Roll=np.array([az ax ay])); See Also: - *SCgetOrds*, *SC.update_support*, *SC.support_offset_and_roll*, *plot_support*, *SC.apply_errors*, + *ords_from_regex*, *SC.update_support*, *SC.support_offset_and_roll*, *plot_support*, *SC.apply_errors*, *SC.register_magnets*, *update_transformation* """ @@ -430,7 +426,7 @@ def register_supports(self, support_ords: ndarray, support_type: str, **kwargs): if upstream := np.sum(np.diff(support_ords, axis=0) < 0): LOGGER.warning(f"{upstream} {support_type} endpoints(s) may be upstream of startpoint(s).") # TODO check the dimensions of Roll and Offset values - self.ORD[support_type] = update_double_ordinates(self.ORD[support_type], support_ords) + self.ORD[support_type] = classdef_tools.update_double_ordinates(self.ORD[support_type], support_ords) for ind in np.ravel(support_ords): setattr(self.RING[ind], f"{support_type}Offset", np.zeros(3)) # [x,y,z] setattr(self.RING[ind], f"{support_type}Roll", np.zeros(3)) # [az,ax,ay] @@ -471,7 +467,7 @@ def set_systematic_multipole_errors(self, ords: ndarray, BA, order: int, skewnes Defines systematic multipole components for the 'QF' magnet and adds it to the field offsets of all magnets named 'QF':: - ords = SCgetOrds(SC.RING,'QF') + ords = sc_tools.ords_from_regex(SC.RING,'QF') BA = np.array([[1E-5, 0], [1E-4, 0], [0, 0], [1E-2, 0]]) RING = SC.set_systematic_multipole_errors(RING, ords, BA, 1, False) @@ -505,7 +501,7 @@ def set_random_multipole_errors(self, ords: ndarray, BA): Defines random multipole components for the 'QF' magnet and adds it to the field offsets of all magnets named 'QF':: - ords = SCgetOrds(SC.RING,'QF') + ords = sc_tools.ords_from_regex(SC.RING,'QF') BA = np.array([[1E-5, 0], [1E-4, 0], [0, 0], [1E-2, 0]]) SC.set_random_multipole_errors(ords, BA) @@ -515,11 +511,11 @@ def set_random_multipole_errors(self, ords: ndarray, BA): if BA.ndim != 2 or BA.shape[1] != 2: raise ValueError("BA has to be numpy.array of shape N x 2.") for ord in ords: - randBA = SCrandnc(2, BA.shape) * BA # TODO this should be registered in SC.SIG + randBA = sc_tools.randnc(2, BA.shape) * BA # TODO this should be registered in SC.SIG for ind, target in enumerate(("B", "A")): attr_name = f"Polynom{target}Offset" setattr(self.RING[ord], attr_name, - add_padded(getattr(self.RING[ord], attr_name), randBA[:, ind]) + classdef_tools.add_padded(getattr(self.RING[ord], attr_name), randBA[:, ind]) if hasattr(self.RING[ord], attr_name) else randBA[:, ind]) def apply_errors(self, nsigmas: float = 2): @@ -544,28 +540,28 @@ def apply_errors(self, nsigmas: float = 2): *SC.register_magnets*, *SC.register_support*, *SC.register_bpms*, *SC.register_cavities* """ # RF - for ind in intersect(self.ORD.RF, self.SIG.RF.keys()): + for ind in classdef_tools.intersect(self.ORD.RF, self.SIG.RF.keys()): for field in self.SIG.RF[ind]: - setattr(self.RING[ind], field, randn_cutoff(self.SIG.RF[ind][field], nsigmas)) + setattr(self.RING[ind], field, classdef_tools.randn_cutoff(self.SIG.RF[ind][field], nsigmas)) # BPM - for ind in intersect(self.ORD.BPM, self.SIG.BPM.keys()): + for ind in classdef_tools.intersect(self.ORD.BPM, self.SIG.BPM.keys()): for field in self.SIG.BPM[ind]: if re.search('Noise', field): setattr(self.RING[ind], field, self.SIG.BPM[ind][field]) else: - setattr(self.RING[ind], field, randn_cutoff(self.SIG.BPM[ind][field], nsigmas)) + setattr(self.RING[ind], field, classdef_tools.randn_cutoff(self.SIG.BPM[ind][field], nsigmas)) # Magnet - for ind in intersect(self.ORD.Magnet, self.SIG.Magnet.keys()): + for ind in classdef_tools.intersect(self.ORD.Magnet, self.SIG.Magnet.keys()): for field in self.SIG.Magnet[ind]: setattr(self.RING[ind], 'BendingAngleError' if field == 'BendingAngle' else field, - randn_cutoff(self.SIG.Magnet[ind][field], nsigmas)) + classdef_tools.randn_cutoff(self.SIG.Magnet[ind][field], nsigmas)) # Injection - self.INJ.Z0 = self.INJ.Z0ideal + self.SIG.staticInjectionZ * SCrandnc(nsigmas, (6,)) + self.INJ.Z0 = self.INJ.Z0ideal + self.SIG.staticInjectionZ * sc_tools.randnc(nsigmas, (6,)) self.INJ.randomInjectionZ = 1 * self.SIG.randomInjectionZ # Circumference if 'Circumference' in self.SIG.keys(): - circScaling = 1 + self.SIG.Circumference * SCrandnc(nsigmas, (1, 1)) - self.RING = SCscaleCircumference(self.RING, circScaling, 'rel') + circScaling = 1 + self.SIG.Circumference * sc_tools.randnc(nsigmas, (1, 1)) + self.RING = sc_tools.scale_circumference(self.RING, circScaling, 'rel') LOGGER.info('Circumference error applied.') # Misalignments self._apply_support_alignment_error(nsigmas) @@ -577,7 +573,7 @@ def apply_errors(self, nsigmas: float = 2): self.update_cavities() def _apply_support_alignment_error(self, nsigmas): - s_pos = findspos(self.RING) + s_pos = at_wrapper.findspos(self.RING) for support_type in SUPPORT_TYPES: for ordPair in self.ORD[support_type].T: if ordPair[0] not in self.SIG.Support.keys(): @@ -585,9 +581,9 @@ def _apply_support_alignment_error(self, nsigmas): for field, value in self.SIG.Support[ordPair[0]].items(): if support_type not in field: continue - setattr(self.RING[ordPair[0]], field, randn_cutoff(value, nsigmas)) + setattr(self.RING[ordPair[0]], field, classdef_tools.randn_cutoff(value, nsigmas)) setattr(self.RING[ordPair[1]], field, - randn_cutoff(value, nsigmas) if field in self.SIG.Support[ordPair[1]].keys() + classdef_tools.randn_cutoff(value, nsigmas) if field in self.SIG.Support[ordPair[1]].keys() else getattr(self.RING[ordPair[0]], field)) struct_length = np.remainder(np.diff(s_pos[ordPair]), s_pos[-1]) @@ -680,14 +676,14 @@ def update_supports(self, offset_bpms: bool = True, offset_magnets: bool = True) *SC.register_support*, *SC.support_offset_and_roll*, *plot_support* """ - s_pos = findspos(self.RING) + s_pos = at_wrapper.findspos(self.RING) if offset_magnets: if len(self.ORD.Magnet): offsets, rolls = self.support_offset_and_roll(s_pos[self.ORD.Magnet]) for i, ind in enumerate(self.ORD.Magnet): setattr(self.RING[ind], "SupportOffset", offsets[:, i]) setattr(self.RING[ind], "SupportRoll", rolls[:, i]) - self.RING[ind] = update_transformation(self.RING[ind]) + self.RING[ind] = sc_tools.update_transformation(self.RING[ind]) if hasattr(self.RING[ind], 'MasterOf'): for child_ind in self.RING[ind].MasterOf: for field in ("T1", "T2", "R1", "R2"): @@ -725,7 +721,7 @@ def support_offset_and_roll(self, s_locations: ndarray) -> Tuple[ndarray, ndarra See Also: *SC.register_support*, *SC.update_support*, *plot_support* """ - s_pos = findspos(self.RING) + s_pos = at_wrapper.findspos(self.RING) ring_length = s_pos[-1] off0 = np.zeros((3, len(s_pos))) roll0 = np.zeros((3, len(s_pos))) @@ -740,7 +736,7 @@ def support_offset_and_roll(self, s_locations: ndarray) -> Tuple[ndarray, ndarra tmpoff1[:, i] = off0[:, ord1[i]] + getattr(self.RING[ord1[i]], f"{suport_type}Offset") tmpoff2[:, i] = off0[:, ord2[i]] + getattr(self.RING[ord2[i]], f"{suport_type}Offset") for i in range(3): - off0[i, :] = s_interpolation(off0[i, :], s_pos, ord1, tmpoff1[i, :], ord2, tmpoff2[i, :]) + off0[i, :] = classdef_tools.s_interpolation(off0[i, :], s_pos, ord1, tmpoff1[i, :], ord2, tmpoff2[i, :]) for support_type in SUPPORT_TYPES: # Order has to be Section, Plinth, Girder for ords in self.ORD[support_type].T: @@ -797,12 +793,12 @@ def set_cavity_setpoints(self, ords: Union[int, List[int], ndarray], 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) + ords_1d, setpoints_1d = classdef_tools.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) + setpoints_1d *= at_wrapper.atgetfieldvalues(self.RING, ords_1d, setpoint_str) if method == SETTING_ADD: - setpoints_1d += atgetfieldvalues(self.RING, ords_1d, setpoint_str) + setpoints_1d += at_wrapper.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) @@ -825,7 +821,7 @@ def get_cm_setpoints(self, ords: Union[int, List[int], ndarray], skewness: bool) order = 0 ndim = int(skewness) letter = NUM_TO_AB[ndim] - setpoints = atgetfieldvalues(self.RING, ords_1d, f"{SETPOINT}{letter}", order) + setpoints = at_wrapper.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 @@ -863,7 +859,7 @@ def set_cm_setpoints(self, ords: Union[int, List[int], ndarray], """ # TODO corrector does not have PolynomA/B in at? - ords_1d, setpoints_1d = check_input_and_setpoints(method, ords, setpoints) + ords_1d, setpoints_1d = classdef_tools.check_input_and_setpoints(method, ords, setpoints) order = 0 ndim = int(skewness) letter = NUM_TO_AB[ndim] @@ -914,23 +910,23 @@ def set_magnet_setpoints(self, ords: Union[int, List[int], ndarray], Examples: Identify the ordinates of all elements named `'SF'` and switch their sextupole component off:: - ords = SCgetOrds(SC.RING,'SF') + ords = sc_tools.ords_from_regex(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') + ords = sc_tools.ords_from_regex(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) + ords_1d, setpoints_1d = classdef_tools.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) + setpoints_1d *= at_wrapper.atgetfieldvalues(self.RING, ords_1d, f"NomPolynom{letter}", order) if method == SETTING_ADD: - setpoints_1d += atgetfieldvalues(self.RING, ords_1d, f"{SETPOINT}{letter}", order) + setpoints_1d += at_wrapper.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') @@ -1061,7 +1057,7 @@ def _check_kwargs(kwargs, allowed_options): def _optional_magnet_fields(self, ind, MAGords, **kwargs): if 'CF' in kwargs.keys(): self.RING[ind].CombinedFunction = True - if intersect(("HCM", "VCM"), kwargs.keys()) and not hasattr(self.RING[ind], 'CMlimit'): + if classdef_tools.intersect(("HCM", "VCM"), kwargs.keys()) and not hasattr(self.RING[ind], 'CMlimit'): self.RING[ind].CMlimit = np.zeros(2) if 'HCM' in kwargs.keys(): self.RING[ind].CMlimit[0] = kwargs["HCM"] @@ -1076,17 +1072,17 @@ def _optional_magnet_fields(self, ind, MAGords, **kwargs): def _update_magnets(self, source_ord, target_ord): setpoints_a, setpoints_b = self.RING[source_ord].SetPointA, self.RING[source_ord].SetPointB - polynoms = dict(A=setpoints_a * add_padded(np.ones(len(setpoints_a)), self.RING[source_ord].CalErrorA), - B=setpoints_b * add_padded(np.ones(len(setpoints_b)), self.RING[source_ord].CalErrorB)) + polynoms = dict(A=setpoints_a * classdef_tools.add_padded(np.ones(len(setpoints_a)), self.RING[source_ord].CalErrorA), + B=setpoints_b * classdef_tools.add_padded(np.ones(len(setpoints_b)), self.RING[source_ord].CalErrorB)) for target in AB: new_polynom = polynoms[target][:] if hasattr(self.RING[target_ord], f'Polynom{target}Offset'): - new_polynom = add_padded(new_polynom, getattr(self.RING[target_ord], f'Polynom{target}Offset')) + new_polynom = classdef_tools.add_padded(new_polynom, getattr(self.RING[target_ord], f'Polynom{target}Offset')) for source in AB: if hasattr(self.RING[target_ord], f'SysPol{target}From{source}'): polynom_errors = getattr(self.RING[target_ord], f'SysPol{target}From{source}') for n in polynom_errors.keys(): - new_polynom = add_padded(new_polynom, polynoms[source][n] * polynom_errors[n]) + new_polynom = classdef_tools.add_padded(new_polynom, polynoms[source][n] * polynom_errors[n]) setattr(self.RING[target_ord], f"Polynom{target}", new_polynom) if hasattr(self.RING[source_ord], 'BendingAngleError'): diff --git a/pySC/correction/bba.py b/pySC/correction/bba.py index 925509d..0bdda8c 100644 --- a/pySC/correction/bba.py +++ b/pySC/correction/bba.py @@ -2,12 +2,9 @@ import numpy as np from pySC.core.beam import bpm_reading, all_elements_reading -from pySC.utils.sc_tools import SCrandnc -from pySC.utils import logging_tools from pySC.core.classes import DotDict from pySC.core.constants import TRACK_TBT, NUM_TO_AB, SETTING_REL, SETTING_ABS -from pySC.utils.stats import weighted_mean, weighted_error, effective_sample_size, weights_from_errors -from pySC.utils.at_wrapper import atgetfieldvalues, findspos +from pySC.utils import at_wrapper, logging_tools, sc_tools, stats from pySC.correction import orbit_trajectory @@ -27,7 +24,7 @@ def trajectory_bba(SC, bpm_ords, mag_ords, **kwargs): raise ValueError('Beam-trajectory-based alignment works in TBT mode with 2 turns. ' 'Please set: SC.INJ.nTurns = 2 and SC.INJ.trackMode = "TBT"') - q0 = atgetfieldvalues(SC.RING, par.q_ord_phase, "SetPointB", index=1) + q0 = at_wrapper.atgetfieldvalues(SC.RING, par.q_ord_phase, "SetPointB", index=1) 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]): # TODO currently assumes either horizontal or both planes @@ -97,7 +94,7 @@ def fake_bba(SC, bpm_ords, mag_ords, errors=None, fake_offset=None): fake_bpm_offset = (SC.RING[mag_ords[inds[0], inds[1]]].MagnetOffset[inds[0]] + SC.RING[mag_ords[inds[0], inds[1]]].SupportOffset[inds[0]] - SC.RING[bpm_ords[inds[0], inds[1]]].SupportOffset[inds[0]] - + fake_offset[inds[0]] * SCrandnc()) + + fake_offset[inds[0]] * sc_tools.randnc()) if not np.isnan(fake_bpm_offset): SC.RING[bpm_ords[inds[0], inds[1]]].Offset[inds[0]] = fake_bpm_offset else: @@ -134,10 +131,10 @@ def is_bba_errored(bba_offsets, bba_offset_errors): def _get_bpm_offset_from_mag(ring, bpm_ords, mag_ords): offset = np.full(bpm_ords.shape, np.nan) for n_dim in range(bpm_ords.shape[0]): - offset[n_dim, :] = (atgetfieldvalues(ring, bpm_ords[n_dim, :], 'Offset', n_dim) - + atgetfieldvalues(ring, bpm_ords[n_dim, :], 'SupportOffset', n_dim) - - atgetfieldvalues(ring, mag_ords[n_dim, :], 'MagnetOffset', n_dim) - - atgetfieldvalues(ring, mag_ords[n_dim, :], 'SupportOffset', n_dim)) + offset[n_dim, :] = (at_wrapper.atgetfieldvalues(ring, bpm_ords[n_dim, :], 'Offset', n_dim) + + at_wrapper.atgetfieldvalues(ring, bpm_ords[n_dim, :], 'SupportOffset', n_dim) + - at_wrapper.atgetfieldvalues(ring, mag_ords[n_dim, :], 'MagnetOffset', n_dim) + - at_wrapper.atgetfieldvalues(ring, mag_ords[n_dim, :], 'SupportOffset', n_dim)) return offset @@ -175,8 +172,8 @@ 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 = weighted_mean(center[mask], center_err[mask]) - offset_change_error = weighted_error(center[mask]-offset_change, center_err[mask]) / np.sqrt(effective_sample_size(center[mask], weights_from_errors(center_err[mask]))) + offset_change = stats.weighted_mean(center[mask], center_err[mask]) + 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 return offset_change, offset_change_error @@ -264,7 +261,7 @@ def _data_measurement_tbt(SC, m_ord, bpm_ind, j_bpm, n_dim, scaling, par): def _phase_advance_injection_scan(SC, n_dim, last_bpm_ind, par): - q0 = atgetfieldvalues(SC.RING, par.q_ord_phase, "SetPointB", index=1) + q0 = at_wrapper.atgetfieldvalues(SC.RING, par.q_ord_phase, "SetPointB", index=1) n_setpoints = len(par.q_ord_setpoints) bpm_ranges = np.zeros((n_setpoints, len(SC.ORD.BPM))) scalings = np.zeros(n_setpoints) @@ -357,7 +354,7 @@ def _get_orbit_bump(SC, cm_ord, bpm_ord, n_dim, par): # TODO def _plot_bba_step(SC, ax, bpm_ind, n_dim): - s_pos = findspos(SC.RING) + s_pos = at_wrapper.findspos(SC.RING) bpm_readings, all_elements_positions = all_elements_reading(SC) ax.plot(s_pos[SC.ORD.BPM], 1E3 * bpm_readings[n_dim, :len(SC.ORD.BPM)], marker='o') ax.plot(s_pos[SC.ORD.BPM[bpm_ind]], 1E3 * bpm_readings[n_dim, bpm_ind], marker='o', markersize=10, markerfacecolor='k') diff --git a/pySC/correction/injection_fit.py b/pySC/correction/injection_fit.py index 40e1602..356e2b4 100644 --- a/pySC/correction/injection_fit.py +++ b/pySC/correction/injection_fit.py @@ -1,18 +1,18 @@ import matplotlib.pyplot as plt import numpy as np from numpy import ndarray -from pySC.core.beam import bpm_reading -from pySC.utils.at_wrapper import atpass, findspos -from pySC.utils import logging_tools from scipy.optimize import fmin +from pySC.core.beam import bpm_reading from pySC.core.simulated_commissioning import SimulatedCommissioning +from pySC.utils import at_wrapper, logging_tools LOGGER = logging_tools.get_logger(__name__) PLANE_STR = ('Hor.', 'Ver.') + def fit_injection_drift(SC: SimulatedCommissioning, n_dims: ndarray = np.array([0, 1]), plot: bool = False): # uses SC.INJ.nShots - s_pos = findspos(SC.RING) + s_pos = at_wrapper.findspos(SC.RING) bpm_inds = np.arange(2, dtype=int) + len(SC.ORD.BPM) - 1 bref = bpm_reading(SC)[0][:, bpm_inds] s_bpm = np.array([s_pos[SC.ORD.BPM[-1]] - s_pos[-1], s_pos[SC.ORD.BPM[0]]]) @@ -36,7 +36,7 @@ def fit_injection_drift(SC: SimulatedCommissioning, n_dims: ndarray = np.array([ def fit_injection_trajectory(SC: SimulatedCommissioning, bpm_inds: ndarray = np.array([0, 1, 2]), plot: bool = False): # uses SC.INJ.nShots - s_pos = findspos(SC.RING) + s_pos = at_wrapper.findspos(SC.RING) bref = bpm_reading(SC)[0][:, bpm_inds] delta_z0 = np.zeros(6) delta_z0[0:4] = -fmin(_merit_function, np.zeros(4), args=(SC, bref, SC.ORD.BPM[bpm_inds])) @@ -68,7 +68,7 @@ def _fit_bpm_data(s_bpm, bref): def _merit_function(x, SC, bref, ords_used): - t = atpass(SC.IDEALRING, np.concatenate((x, np.zeros(2))), 1, ords_used, keep_lattice=False)[[0, 2], 0, :, 0] + t = at_wrapper.atpass(SC.IDEALRING, np.concatenate((x, np.zeros(2))), 1, ords_used, keep_lattice=False)[[0, 2], 0, :, 0] return np.sqrt(np.mean(bref - t) ** 2) diff --git a/pySC/correction/loco.py b/pySC/correction/loco.py index 2adf21a..52212a8 100644 --- a/pySC/correction/loco.py +++ b/pySC/correction/loco.py @@ -4,10 +4,9 @@ from pySC.lattice_properties.response_model import SCgetModelRM, SCgetModelDispersion from pySC.core.constants import SETTING_ADD, TRACK_ORB from pySC.core.beam import bpm_reading -from pySC.utils.sc_tools import SCgetPinv import matplotlib.pyplot as plt from scipy.optimize import least_squares -from pySC.utils import logging_tools +from pySC.utils import logging_tools, sc_tools LOGGER = logging_tools.get_logger(__name__) @@ -162,7 +161,7 @@ def get_inverse(jacobian, B, s_cut, weights, plot=False): n_resp_mats = len(jacobian) sum_corr = np.sum(jacobian, axis=2) # Sum over i and j for all planes matrix = np.dot(np.dot(sum_corr, weights), sum_corr.T) - inv_matrix = SCgetPinv(matrix, num_removed=n_resp_mats - min(n_resp_mats, s_cut), plot=plot) + inv_matrix = sc_tools.pinv(matrix, num_removed=n_resp_mats - min(n_resp_mats, s_cut), plot=plot) results = np.ravel(np.dot(inv_matrix, B)) # e = np.ravel(np.dot(matrix, results)) - np.ravel(B) return results diff --git a/pySC/correction/orbit_trajectory.py b/pySC/correction/orbit_trajectory.py index 1cf80eb..2d0d444 100644 --- a/pySC/correction/orbit_trajectory.py +++ b/pySC/correction/orbit_trajectory.py @@ -12,9 +12,8 @@ from typing import Tuple from pySC.core.beam import bpm_reading -from pySC.utils import logging_tools -from pySC.utils import sc_tools from pySC.core.constants import SETTING_ADD +from pySC.utils import logging_tools, sc_tools LOGGER = logging_tools.get_logger(__name__) NREPRO: int = 5 @@ -60,7 +59,7 @@ def first_turn(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, LOGGER.debug('First turn threading: Start') bpm_ords, cm_ords, reference = _check_ords(SC, response_matrix, reference, bpm_ords, cm_ords) bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging(SC, bpm_ords=bpm_ords) # Inject... - Mplus = sc_tools.SCgetPinv(response_matrix, **pinv_params) + Mplus = sc_tools.pinv(response_matrix, **pinv_params) for n in range(maxsteps): if transmission_history[-1] == 0: @@ -146,7 +145,7 @@ def stitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, n_b reference = reference.reshape(2, len(bpm_readings[0])) reference[:, len(bpm_ords):] = 0 reference = reference.reshape(response_matrix.shape[0]) - Mplus = sc_tools.SCgetPinv(response_matrix, **pinv_params) + Mplus = sc_tools.pinv(response_matrix, **pinv_params) # Main loop for steps in range(maxsteps): @@ -213,7 +212,7 @@ def balance(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, ep reference = reference.reshape(2, len(bpm_readings[0])) reference[:, len(bpm_ords):] = 0 reference = reference.reshape(response_matrix.shape[0]) - Mplus = sc_tools.SCgetPinv(response_matrix, **pinv_params) + Mplus = sc_tools.pinv(response_matrix, **pinv_params) # Main loop for steps in range(maxsteps): @@ -288,7 +287,7 @@ def correct(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, ep bpm_ords, cm_ords, reference = _check_ords(SC, response_matrix[:, :-1] if scaleDisp else response_matrix, reference, bpm_ords, cm_ords) bpm_readings, transmission_history, rms_orbit_history = _bpm_reading_and_logging(SC, bpm_ords=bpm_ords) # Inject ... - Mplus = sc_tools.SCgetPinv(response_matrix, **pinv_params) + Mplus = sc_tools.pinv(response_matrix, **pinv_params) # Main loop for steps in range(maxsteps): diff --git a/pySC/correction/ramp_errors.py b/pySC/correction/ramp_errors.py index e1c5f21..27ba066 100644 --- a/pySC/correction/ramp_errors.py +++ b/pySC/correction/ramp_errors.py @@ -2,8 +2,7 @@ from pySC.core.constants import SUPPORT_TYPES, RF_PROPERTIES from pySC.correction import orbit_trajectory from pySC.lattice_properties.response_model import SCgetModelRM -from pySC.utils.sc_tools import SCscaleCircumference -from pySC.utils import logging_tools +from pySC.utils import logging_tools, sc_tools LOGGER = logging_tools.get_logger(__name__) @@ -87,5 +86,5 @@ def scaleCircumference(SC, SC0, scale): for ord in range(len(SC0.RING)): D += SC0.RING[ord].Length D0 += SC0.IDEALRING[ord].Length - SC.RING = SCscaleCircumference(SC.RING, scale * (D - D0) + D0, 'abs') + SC.RING = sc_tools.scale_circumference(SC.RING, scale * (D - D0) + D0, 'abs') return SC diff --git a/pySC/correction/rf.py b/pySC/correction/rf.py index 2227494..069b286 100644 --- a/pySC/correction/rf.py +++ b/pySC/correction/rf.py @@ -8,10 +8,9 @@ import matplotlib.pyplot as plt from scipy.optimize import curve_fit, fsolve -from pySC.utils.at_wrapper import findorbit6 from pySC.core.beam import bpm_reading -from pySC.utils import logging_tools from pySC.core.constants import SPEED_OF_LIGHT +from pySC.utils import at_wrapper, logging_tools LOGGER = logging_tools.get_logger(__name__) MIN_TURNS_FOR_LINEAR_FIT = 3 @@ -99,7 +98,7 @@ def correct_rf_frequency(SC, cav_ords=None, bpm_ords=None, n_steps=15, f_range=( def phase_and_energy_error(SC, cav_ords): """This is not a simple observable in reality.""" lamb = SPEED_OF_LIGHT / SC.RING[cav_ords[0]].Frequency - orbit6 = findorbit6(SC.RING)[0] + orbit6 = at_wrapper.findorbit6(SC.RING)[0] phase_error = _fold_phase((orbit6[5] - SC.INJ.Z0[5]) / lamb, lamb) * 360 energy_error = SC.INJ.Z0[4] - orbit6[4] LOGGER.info(f'Static phase error {phase_error:.0f} deg') diff --git a/pySC/example.py b/pySC/example.py index 508e509..2cdcd4d 100644 --- a/pySC/example.py +++ b/pySC/example.py @@ -3,13 +3,11 @@ from pySC.correction.bba import trajectory_bba, fake_bba from at import Lattice -from pySC.utils.at_wrapper import atloco from pySC.core.simulated_commissioning import SimulatedCommissioning from pySC.correction import orbit_trajectory from pySC.core.beam import bpm_reading, beam_transmission 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.correction.loco_wrapper import (loco_model, loco_fit_parameters, apply_lattice_correction, loco_measurement, loco_bpm_structure, loco_cm_structure) from pySC.plotting.plot_phase_space import plot_phase_space @@ -17,7 +15,7 @@ from pySC.plotting.plot_lattice import plot_lattice 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 +from pySC.utils import at_wrapper, logging_tools, sc_tools LOGGER = logging_tools.get_logger(__name__) @@ -51,44 +49,44 @@ def _marker(name): LOGGER.info(f"{len(ring)=}") SC = SimulatedCommissioning(ring) # at.summary(ring) - ords = SCgetOrds(SC.RING, 'BPM') + ords = sc_tools.ords_from_regex(SC.RING, 'BPM') SC.register_bpms(ords, CalError=5E-2 * np.ones(2), # x and y, relative Offset=500E-6 * np.ones(2), # x and y, [m] Noise=10E-6 * np.ones(2), # x and y, [m] NoiseCO=1E-6 * np.ones(2), # x and y, [m] Roll=1E-3) # az, [rad] - ords = SCgetOrds(SC.RING, 'QF') + ords = sc_tools.ords_from_regex(SC.RING, 'QF') SC.register_magnets(ords, HCM=1E-3, # [rad] CalErrorB=np.array([5E-2, 1E-3]), # relative MagnetOffset=200E-6 * np.array([1, 1, 0]), # x, y and z, [m] MagnetRoll=200E-6 * np.array([1, 0, 0])) # az, ax and ay, [rad] - ords = SCgetOrds(SC.RING, 'QD') + ords = sc_tools.ords_from_regex(SC.RING, 'QD') SC.register_magnets(ords, VCM=1E-3, # [rad] CalErrorA=np.array([5E-2, 0]), # relative CalErrorB=np.array([0, 1E-3]), # relative MagnetOffset=200E-6 * np.array([1, 1, 0]), # x, y and z, [m] MagnetRoll=200E-6 * np.array([1, 0, 0])) # az, ax and ay, [rad] - ords = SCgetOrds(SC.RING, 'BEND') + ords = sc_tools.ords_from_regex(SC.RING, 'BEND') SC.register_magnets(ords, BendingAngle=1E-3, # relative MagnetOffset=200E-6 * np.array([1, 1, 0]), # x, y and z, [m] MagnetRoll=200E-6 * np.array([1, 0, 0])) # az, ax and ay, [rad] - ords = SCgetOrds(SC.RING, 'SF|SD') + ords = sc_tools.ords_from_regex(SC.RING, 'SF|SD') SC.register_magnets(ords, SkewQuad=0.1, # [1/m] CalErrorA=np.array([0, 1E-3, 0]), # relative CalErrorB=np.array([0, 0, 1E-3]), # relative MagnetOffset=200E-6 * np.array([1, 1, 0]), # x, y and z, [m] MagnetRoll=200E-6 * np.array([1, 0, 0])) # az, ax and ay, [rad] - ords = SCgetOrds(SC.RING, 'RFCav') + ords = sc_tools.ords_from_regex(SC.RING, 'RFCav') SC.register_cavities(ords, FrequencyOffset=5E3, # [Hz] VoltageOffset=5E3, # [V] TimeLagOffset=0.5) # [m] - ords = np.vstack((SCgetOrds(SC.RING, 'GirderStart'), SCgetOrds(SC.RING, 'GirderEnd'))) + ords = np.vstack((sc_tools.ords_from_regex(SC.RING, 'GirderStart'), sc_tools.ords_from_regex(SC.RING, 'GirderEnd'))) SC.register_supports(ords, "Girder", Offset=100E-6 * np.array([1, 1, 0]), # x, y and z, [m] Roll=200E-6 * np.array([1, 0, 0])) # az, ax and ay, [rad] - ords = np.vstack((SCgetOrds(SC.RING, 'SectionStart'), SCgetOrds(SC.RING, 'SectionEnd'))) + ords = np.vstack((sc_tools.ords_from_regex(SC.RING, 'SectionStart'), sc_tools.ords_from_regex(SC.RING, 'SectionEnd'))) SC.register_supports(ords, "Section", Offset=100E-6 * np.array([1, 1, 0])) # x, y and z, [m] SC.INJ.beamSize = np.diag(np.array([200E-6, 100E-6, 100E-6, 50E-6, 1E-3, 1E-4]) ** 2) @@ -96,9 +94,9 @@ def _marker(name): SC.SIG.staticInjectionZ = np.array([1E-3, 1E-4, 1E-3, 1E-4, 1E-3, 1E-3]) # [m; rad; m; rad; rel.; m] SC.SIG.Circumference = 2E-4 # relative SC.INJ.beamLostAt = 0.6 # relative - for ord in SCgetOrds(SC.RING, 'Drift'): + for ord in sc_tools.ords_from_regex(SC.RING, 'Drift'): SC.RING[ord].EApertures = 13E-3 * np.array([1, 1]) # [m] - for ord in SCgetOrds(SC.RING, 'QF|QD|BEND|SF|SD'): + for ord in sc_tools.ords_from_regex(SC.RING, 'QF|QD|BEND|SF|SD'): SC.RING[ord].EApertures = 10E-3 * np.array([1, 1]) # [m] SC.RING[SC.ORD.Magnet[50]].EApertures = np.array([6E-3, 3E-3]) # [m] @@ -108,7 +106,7 @@ def _marker(name): plot_support(SC) SC.RING = switch_cavity_and_radiation(SC.RING, 'cavityoff') - sextOrds = SCgetOrds(SC.RING, 'SF|SD') + sextOrds = sc_tools.ords_from_regex(SC.RING, 'SF|SD') 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) @@ -128,9 +126,9 @@ def _marker(name): # plot_cm_strengths(SC) # Performing trajectory BBA SC.INJ.nParticles = 1 - quadOrds = np.tile(SCgetOrds(SC.RING, 'QF|QD'), (2, 1)) + quadOrds = np.tile(sc_tools.ords_from_regex(SC.RING, 'QF|QD'), (2, 1)) BPMords = np.tile(SC.ORD.BPM, (2, 1)) - SC, bba_offsets, bba_offset_errors = trajectory_bba(SC, BPMords, quadOrds, q_ord_phase=SCgetOrds(SC.RING, 'QF|QD')[0], + SC, bba_offsets, bba_offset_errors = trajectory_bba(SC, BPMords, quadOrds, q_ord_phase=sc_tools.ords_from_regex(SC.RING, 'QF|QD')[0], q_ord_setpoints=np.array([0.8, 0.9, 1.0, 1.1, 1.2]), magnet_strengths=np.array([0.8, 0.9, 1.0, 1.1, 1.2]), dipole_compensation=True, plot_results=True) @@ -161,7 +159,7 @@ def _marker(name): [maxTurns, lostCount] = beam_transmission(SC, nParticles=100, nTurns=10) # Faking-BBA - quadOrds = np.tile(SCgetOrds(SC.RING, 'QF|QD'), (2, 1)) + quadOrds = np.tile(sc_tools.ords_from_regex(SC.RING, 'QF|QD'), (2, 1)) BPMords = np.tile(SC.ORD.BPM, (2, 1)) SC = fake_bba(SC, BPMords, quadOrds, fake_offset=np.array([50E-6, 50E-6])) @@ -183,7 +181,7 @@ def _marker(name): SC.RING = switch_cavity_and_radiation(SC.RING, 'cavityon') plot_phase_space(SC, nParticles=10, nTurns=1000) maxTurns, lostCount = beam_transmission(SC, nParticles=100, nTurns=200, plot=True) - SC, _, _, _ = tune_scan(SC, np.vstack((SCgetOrds(SC.RING, 'QF'), SCgetOrds(SC.RING, 'QD'))), + SC, _, _, _ = tune_scan(SC, np.vstack((sc_tools.ords_from_regex(SC.RING, 'QF'), sc_tools.ords_from_regex(SC.RING, 'QD'))), np.outer(np.ones(2), 1 + np.linspace(-0.01, 0.01, 51)), do_plot=False, nParticles=100, nTurns=200) CMstep = 1E-4 # [rad] # TODO later in the structure it is in mrad, ??? @@ -194,18 +192,18 @@ def _marker(name): cm_data = loco_cm_structure(SC, CMstep, FitKicks=True) loco_meas_data = loco_measurement(SC, CMstep, RFstep, SC.ORD.BPM, SC.ORD.CM) fit_parameters = loco_fit_parameters(SC, init.SC.RING, ring_data, RFstep, - [SCgetOrds(SC.RING, 'QF'), False, 'individual', 1E-3], + [sc_tools.ords_from_regex(SC.RING, 'QF'), False, 'individual', 1E-3], # {Ords, normal/skew, ind/fam, deltaK} - [SCgetOrds(SC.RING, 'QD'), False, 'individual', 1E-4]) + [sc_tools.ords_from_regex(SC.RING, 'QD'), False, 'individual', 1E-4]) for n in range(6): - _, bpm_data, cm_data, fit_parameters, loco_flags, ring_data = atloco(loco_meas_data, bpm_data, cm_data, + _, bpm_data, cm_data, fit_parameters, loco_flags, ring_data = at_wrapper.atloco(loco_meas_data, bpm_data, cm_data, fit_parameters, loco_flags, ring_data) SC = apply_lattice_correction(SC, fit_parameters) SC = orbit_trajectory.correct(SC, MCO, alpha=50, target=0, maxsteps=30) if n == 3: loco_flags.Coupling = True fit_parameters = loco_fit_parameters(SC, init.SC.RING, ring_data, RFstep, - [SCgetOrds(SC.RING, 'QF'), False, 'individual', 1E-3], - [SCgetOrds(SC.RING, 'QD'), False, 'individual', 1E-4], + [sc_tools.ords_from_regex(SC.RING, 'QF'), False, 'individual', 1E-3], + [sc_tools.ords_from_regex(SC.RING, 'QD'), False, 'individual', 1E-4], [SC.ORD.SkewQuad, True, 'individual', 1E-3]) diff --git a/pySC/lattice_properties/apertures.py b/pySC/lattice_properties/apertures.py index 6b64a89..53cf477 100644 --- a/pySC/lattice_properties/apertures.py +++ b/pySC/lattice_properties/apertures.py @@ -1,8 +1,7 @@ import matplotlib.pyplot as plt import numpy as np -from pySC.utils.at_wrapper import atpass, findorbit6, findorbit4, findspos -from pySC.utils import logging_tools +from pySC.utils import at_wrapper, logging_tools LOGGER = logging_tools.get_logger(__name__) @@ -16,9 +15,9 @@ def SCdynamicAperture(RING, dE, bounds=np.array([0, 1e-3]), nturns=1000, thetas= ZCO = np.zeros(6) if launchOnOrbit: if useOrbit6: - ZCO = findorbit6(RING) + ZCO = at_wrapper.findorbit6(RING) else: - tmp = findorbit4(RING, 0) + tmp = at_wrapper.findorbit4(RING, 0) if ~np.isnan(tmp[0]): ZCO[0:4] = tmp ZCO[5] = ZCO[5] + dE @@ -27,7 +26,7 @@ def SCdynamicAperture(RING, dE, bounds=np.array([0, 1e-3]), nturns=1000, thetas= for cntt in range(len(thetas)): # Loop over angles theta = thetas[cntt] limits = inibounds - atpass(RING, np.full(6, np.nan), 1, [1]) # Fake Track to initialize lattice + at_wrapper.atpass(RING, np.full(6, np.nan), 1, [1]) # Fake Track to initialize lattice scales = 0 while scales < 16: if _check_bounds(RING, ZCO, nturns, theta, limits): @@ -49,9 +48,9 @@ def SCdynamicAperture(RING, dE, bounds=np.array([0, 1e-3]), nturns=1000, thetas= DA = np.sum(np.sin(dthetas) * r0 * r1 / 2.) if centerOnOrbit: if useOrbit6: - tmp = findorbit6(RING) + tmp = at_wrapper.findorbit6(RING) else: - tmp = findorbit4(RING, 0) + tmp = at_wrapper.findorbit4(RING, 0) if not np.isnan(tmp[0]): x, y = pol2cart(RMAXs, thetas) x = x - tmp[0] @@ -64,7 +63,7 @@ def SCdynamicAperture(RING, dE, bounds=np.array([0, 1e-3]), nturns=1000, thetas= def SCmomentumAperture(RING, REFPTS, inibounds, nturns=1000, accuracy=1e-4, stepsize=1e-3, plot=0): dboundHI = np.zeros(len(REFPTS)) dboundLO = np.zeros(len(REFPTS)) - ZCOs = findorbit6(RING, REFPTS) + ZCOs = at_wrapper.findorbit6(RING, REFPTS) if any(~np.isfinite(ZCOs.flatten())): dbounds = np.array([dboundHI, dboundLO]).T LOGGER.info('Closed Orbit could not be determined during MA evaluation. MA set to zero.') @@ -85,7 +84,7 @@ def SCmomentumAperture(RING, REFPTS, inibounds, nturns=1000, accuracy=1e-4, step LOGGER.debug(f'ord: {ord}; Found: {local_bounds}') dbounds = np.array([dboundHI, dboundLO]).T if plot: - spos = findspos(RING)[REFPTS] + spos = at_wrapper.findspos(RING)[REFPTS] plt.figure(81222) plt.clf() plt.plot(spos, dboundHI, 'kx-') @@ -103,7 +102,7 @@ def _check_bounds(RING, ZCO, nturns, theta, boundsIn): Zmax = ZCO[:] Zmax[0] = boundsIn[1] * np.cos(theta) Zmax[2] = boundsIn[1] * np.sin(theta) - ROUT = atpass(RING, [Zmin, Zmax], nturns, len(RING), keep_lattice=True) # Track + ROUT = at_wrapper.atpass(RING, [Zmin, Zmax], nturns, len(RING), keep_lattice=True) # Track RLAST = ROUT[:, len(ROUT) - 1:len(ROUT)] # Get positions after last turn return ~np.isnan(RLAST[0, 0]) and np.isnan(RLAST[0, 1]) @@ -115,7 +114,7 @@ def _refine_bounds(RING, ZCO, nturns, theta, boundsIn): Z = ZCO[:] Z[0] = rmid * np.cos(theta) Z[2] = rmid * np.sin(theta) - ROUT = atpass(RING, Z, nturns, len(RING), keep_lattice=True) # Track + ROUT = at_wrapper.atpass(RING, Z, nturns, len(RING), keep_lattice=True) # Track RLAST = ROUT[:, len(ROUT) - 1] # Get positions after last turn return [rmin, rmid] if np.isnan(RLAST[0]) else [rmid, rmax] # Midpoint is outside or inside DA @@ -147,7 +146,7 @@ def refine_bounds(local_bounds, RING, ZCO, nturns): dmean = np.mean(local_bounds) Z0 = ZCO Z0[4] = Z0[4] + dmean - ROUT = atpass(RING, Z0, nturns, len(RING)) + ROUT = at_wrapper.atpass(RING, Z0, nturns, len(RING)) if np.isnan(ROUT[0]): # Particle dead :( local_bounds[1] = dmean # Set abs-upper bound to midpoint else: # Particle alive :) @@ -158,7 +157,7 @@ def refine_bounds(local_bounds, RING, ZCO, nturns): def check_bounds(local_bounds, RING, ZCO, nturns): Z = np.array([ZCO, ZCO]) Z[4, :] = Z[4, :] + local_bounds[:] - ROUT = atpass(RING, Z, nturns, len(RING)) + ROUT = at_wrapper.atpass(RING, Z, nturns, len(RING)) if np.isnan(ROUT[0, 0]) and not np.isnan(ROUT[0, 1]): LOGGER.warning('Closer-to-momentum particle is unstable. This shouldnt be!') return not np.isnan(ROUT[0, 0]) and np.isnan(ROUT[0, 1]) diff --git a/pySC/lattice_properties/magnet_orbit.py b/pySC/lattice_properties/magnet_orbit.py index 5423ba2..75423f0 100644 --- a/pySC/lattice_properties/magnet_orbit.py +++ b/pySC/lattice_properties/magnet_orbit.py @@ -1,14 +1,14 @@ import numpy as np import matplotlib.pyplot as plt -from pySC.utils.at_wrapper import findorbit6, findspos -from pySC.utils import logging_tools + +from pySC.utils import at_wrapper, logging_tools LOGGER = logging_tools.get_logger(__name__) def SCgetCOD(SC, ords=None, plot=False): if ords is None: ords = SC.ORD.Magnet - T = findorbit6(SC.RING, ords) + T = at_wrapper.findorbit6(SC.RING, ords) if any(np.isnan(T)): LOGGER.error('Closed orbit could not be found.') COD = np.nan(2, len(ords)) @@ -18,7 +18,7 @@ def SCgetCOD(SC, ords=None, plot=False): magOffset[:, i] = SC.RING[ord].T2([1, 3]) COD = T[[1, 3], :] - magOffset if plot: - sPos = findspos(SC.RING)[ords] + sPos = at_wrapper.findspos(SC.RING)[ords] ylabelStr = ['$\Delta x$ [mm]', '$\Delta y$ [mm]'] plt.figure(784) plt.clf() diff --git a/pySC/lattice_properties/response_measurement.py b/pySC/lattice_properties/response_measurement.py index 1e37652..33dbb60 100644 --- a/pySC/lattice_properties/response_measurement.py +++ b/pySC/lattice_properties/response_measurement.py @@ -1,8 +1,8 @@ import numpy as np from pySC.core.beam import bpm_reading -from pySC.utils import logging_tools -from pySC.utils.at_wrapper import atgetfieldvalues +from pySC.utils import at_wrapper, logging_tools + LOGGER = logging_tools.get_logger(__name__) @@ -68,7 +68,7 @@ def dispersion(SC, rf_step, bpm_ords=None, cav_ords=None, n_steps=2): for n_cav, cav_ord in enumerate(cav_ords): rf_steps[n_cav, :] = SC.RING[cav_ord].FrequencySetPoint + np.linspace(-rf_step, rf_step, n_steps) dB = np.zeros((n_steps, *np.shape(bref))) - rf0 = atgetfieldvalues(SC.RING, cav_ords, "FrequencySetPoint") + rf0 = at_wrapper.atgetfieldvalues(SC.RING, cav_ords, "FrequencySetPoint") for nStep in range(n_steps): SC.set_cavity_setpoints(cav_ords, rf_steps[:, nStep], 'Frequency', 'abs') dB[nStep, :] = np.ravel(bpm_reading(SC, bpm_ords=bpm_ords)[0]) - bref diff --git a/pySC/lattice_properties/response_model.py b/pySC/lattice_properties/response_model.py index a6a6c08..3e3f49c 100644 --- a/pySC/lattice_properties/response_model.py +++ b/pySC/lattice_properties/response_model.py @@ -1,11 +1,10 @@ +import copy import numpy as np from at import Lattice from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.utils.at_wrapper import atpass, findorbit6 from pySC.core.constants import NUM_TO_AB, RF_PROPERTIES -import copy -from pySC.utils import logging_tools +from pySC.utils import at_wrapper, logging_tools LOGGER = logging_tools.get_logger(__name__) @@ -50,7 +49,7 @@ def SCgetModelRM(SC, BPMords, CMords, trackMode='TBT', Z0=np.zeros(6), nTurns=1, """ LOGGER.info('Calculating model response matrix') - track_methods = dict(TBT=atpass, ORB=orbpass) + track_methods = dict(TBT=at_wrapper.atpass, ORB=orbpass) if trackMode not in track_methods.keys(): ValueError(f'Unknown track mode {trackMode}. Valid values are {track_methods.keys()}') ring = SC.IDEALRING.deepcopy() if useIdealRing else SCgetModelRING(SC) @@ -88,7 +87,7 @@ def SCgetModelRM(SC, BPMords, CMords, trackMode='TBT', Z0=np.zeros(6), nTurns=1, def orbpass(RING, Z0, nTurns, REFPTS): - return np.transpose(findorbit6(RING, REFPTS)[1])[[0,1,2,3], :].reshape(4, 1, len(REFPTS), 1) + return np.transpose(at_wrapper.findorbit6(RING, REFPTS)[1])[[0,1,2,3], :].reshape(4, 1, len(REFPTS), 1) def SCgetModelDispersion(SC, BPMords, CAVords, trackMode='ORB', Z0=np.zeros(6), nTurns=1, rfStep=1E3, useIdealRing=True): @@ -122,7 +121,7 @@ def SCgetModelDispersion(SC, BPMords, CAVords, trackMode='ORB', Z0=np.zeros(6), """ LOGGER.info('Calculating model dispersion') - track_methods = dict(TBT=atpass, ORB=orbpass) + track_methods = dict(TBT=at_wrapper.atpass, ORB=orbpass) if trackMode not in track_methods.keys(): ValueError(f'Unknown track mode {trackMode}. Valid values are {track_methods.keys()}') ring = SC.IDEALRING.deepcopy() if useIdealRing else SCgetModelRING(SC) diff --git a/pySC/matlab_index.py b/pySC/matlab_index.py index c7dd4ac..082f0c9 100644 --- a/pySC/matlab_index.py +++ b/pySC/matlab_index.py @@ -34,9 +34,8 @@ from pySC.plotting.plot_lattice import plot_lattice, plot_cm_strengths from pySC.plotting.plot_phase_space import plot_phase_space from pySC.plotting.plot_support import plot_support -from pySC.utils import logging_tools -from pySC.utils.sc_tools import SCgetOrds as get_ords, SCgetPinv as get_pinv, SCrandnc as randnc, \ - SCscaleCircumference as scale_circumference, update_transformation, SCmultipolesRead as read_multipoles +from pySC.utils import logging_tools, sc_tools + LOGGER = logging_tools.get_logger(__name__) LOGGER.warn("Matlab_index imported: \n" @@ -149,11 +148,11 @@ def SCgetModelRM(SC, BPMords, CMords, /, *, trackMode='TBT', Z0=np.zeros(6), nTu def SCgetOrds(ring: Lattice, regex: str, /, *, verbose: bool = False) -> ndarray: - return get_ords(ring=ring, regex=regex) + return sc_tools.ords_from_regex(ring=ring, regex=regex) def SCgetPinv(matrix: ndarray, /, *, N: int = 0, alpha: float = 0, damping: float = 1, plot: bool = False) -> ndarray: - return get_pinv(matrix=matrix, num_removed=N, alpha=alpha, damping=damping, plot=plot) + return sc_tools.pinv(matrix=matrix, num_removed=N, alpha=alpha, damping=damping, plot=plot) def SCgetRespMat(SC, Amp, BPMords, CMords, /, *, mode='fixedKick', nSteps=2, fit='linear', verbose=0): @@ -177,7 +176,7 @@ def SCgetTransformation(dx, dy, dz, ax, ay, az, magTheta, magLength, refPoint='c dict(Length=magLength, BendingAngle=magTheta, SupportOffset=np.zeros(3), SupportRoll=np.zeros(3), MagnetOffset=np.array([dx, dy, dz]), MagnetRoll=np.roll(np.array([ax, ay, az]), 1))) # inside the function, it will be rolled back - fake_element = update_transformation(fake_element) + fake_element = sc_tools.update_transformation(fake_element) return fake_element.T1, fake_element.T2, fake_element.R1, fake_element.R2 @@ -194,7 +193,7 @@ def SCmomentumAperture(RING, REFPTS, inibounds, /, *, nturns=1000, accuracy=1e-4 def SCmultipolesRead(fname): - return read_multipoles(fname) + return sc_tools.read_multipoles(fname) def SCparticlesIn3D(*args): @@ -236,7 +235,7 @@ def SCrampUpErrors(SC, /, *, nStepsRamp=10, eps=1e-5, target=0, alpha=10, maxste def SCrandnc(cut_off: float = 2, shape: tuple = (1,)) -> ndarray: - return randnc(cut_off, shape) + return sc_tools.randnc(cut_off, shape) def SCregisterBPMs(SC: SimulatedCommissioning, BPMords: ndarray, **kwargs) -> SimulatedCommissioning: @@ -265,7 +264,7 @@ def SCSanityCheck(SC: SimulatedCommissioning) -> None: def SCscaleCircumference(RING, circ, /, *, mode='abs'): - return scale_circumference(RING, circ, mode=mode) + return sc_tools.scale_circumference(RING, circ, mode=mode) def SCsetCavs2SetPoints(SC: SimulatedCommissioning, CAVords: ndarray, type: str, setpoints: ndarray, /, *, diff --git a/pySC/plotting/plot_apertures.py b/pySC/plotting/plot_apertures.py index 87d3c88..09f152a 100644 --- a/pySC/plotting/plot_apertures.py +++ b/pySC/plotting/plot_apertures.py @@ -5,7 +5,7 @@ from at import Lattice, Refpts from matplotlib import pyplot as plt -from pySC.utils.at_wrapper import findspos, atlinopt +from pySC.utils import at_wrapper SLICES = 400 @@ -18,7 +18,7 @@ def plot_data_apertures(ring: Lattice, refpts: Refpts, **kwargs): ea_all, ra_all = _apertures(ring) # Extract the plot data - s_pos = findspos(ring)[refpts] + s_pos = at_wrapper.findspos(ring)[refpts] ea = ea_all[refpts] * to_mm ra = ra_all[refpts] * to_mm if "cut_off" in kwargs.keys(): @@ -38,7 +38,7 @@ def plot_data_apertures(ring: Lattice, refpts: Refpts, **kwargs): def plot_data_beta_disp(ring: Lattice, refpts, **kwargs): """Generates data for plotting beta functions and dispersion""" - _, _, data = atlinopt(ring, refpts=refpts, get_chrom=True, **kwargs) + _, _, data = at_wrapper.atlinopt(ring, refpts=refpts, get_chrom=True, **kwargs) s_pos = data['s_pos'] betax = data['beta'][:, 0] betay = data['beta'][:, 1] diff --git a/pySC/plotting/plot_lattice.py b/pySC/plotting/plot_lattice.py index f05a587..4e9fd08 100644 --- a/pySC/plotting/plot_lattice.py +++ b/pySC/plotting/plot_lattice.py @@ -5,7 +5,7 @@ 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 -from pySC.utils.at_wrapper import findspos +from pySC.utils import at_wrapper warnings.filterwarnings("ignore", message='Legend does not support handles for PatchCollection instances.') @@ -97,7 +97,7 @@ def plot_cm_strengths(SC: SimulatedCommissioning): init_font = plt.rcParams["font.size"] plt.rcParams.update({'font.size': 18}) f, ax = plt.subplots(nrows=2, num=86, figsize=(9, 7.5), facecolor="w") - s_pos = findspos(SC.RING) + s_pos = at_wrapper.findspos(SC.RING) for n_dim in range(2): setpoints = 1E6 * SC.get_cm_setpoints(SC.ORD.CM[n_dim], skewness=bool(n_dim)) count, bins_count = np.histogram(setpoints, bins=len(setpoints)) @@ -120,8 +120,8 @@ def _get_s_range(ring, s_range, indices, n_sectors): if s_range is not None: return s_range if indices is not None: - s_pos = findspos(ring) + s_pos = at_wrapper.findspos(ring) return s_pos[np.min(indices)], s_pos[np.max(indices)] if n_sectors is not None: - return 0, findspos(ring)[-1] / n_sectors + return 0, at_wrapper.findspos(ring)[-1] / n_sectors return ring.s_range diff --git a/pySC/plotting/plot_phase_space.py b/pySC/plotting/plot_phase_space.py index 56e5c79..3ee8943 100644 --- a/pySC/plotting/plot_phase_space.py +++ b/pySC/plotting/plot_phase_space.py @@ -2,7 +2,7 @@ import numpy as np from pySC.core.constants import SPEED_OF_LIGHT -from pySC.utils.at_wrapper import atpass, findorbit6, findspos +from pySC.utils import at_wrapper from pySC.core.beam import generate_bunches @@ -10,22 +10,22 @@ def plot_phase_space(SC, ords=np.zeros(1, dtype=int), custom_bunch=None, nPartic init_font = plt.rcParams["font.size"] plt.rcParams.update({'font.size': 18}) z_in, n_particles, n_turns = _check_input(SC, custom_bunch, nParticles, nTurns) - T = atpass(SC.RING, z_in, n_turns, ords, keep_lattice=False) + T = at_wrapper.atpass(SC.RING, z_in, n_turns, ords, keep_lattice=False) T[:, np.isnan(T[0, :])] = np.nan label_str = [r'$\Delta x$ [$\mu$m]', r"$\Delta x'$ [$\mu$rad]", r'$\Delta y$ [$\mu$m]', r"$\Delta y'$ [$\mu$rad]", r'$\Delta S$ [m]', r'$\delta E$ $[\%]$'] title_str = ['Horizontal', 'Vertical', 'Longitudinal'] cav_ord = SC.ORD.RF[0] if SC.RING[cav_ord].PassMethod == 'RFCavityPass': - circumference = findspos(SC.RING)[-1] + circumference = at_wrapper.findspos(SC.RING)[-1] length_slippage = SPEED_OF_LIGHT * SC.RING[cav_ord].HarmNumber / SC.RING[cav_ord].Frequency - circumference T[5, :, :, :] = T[5, :, :, :] - length_slippage * np.arange(n_turns)[np.newaxis, np.newaxis, :] label_str[4] = r'$\Delta S_{act}$ [m]' if plotCO: - CO = np.squeeze(findorbit6(SC.RING, ords)[1]) + CO = np.squeeze(at_wrapper.findorbit6(SC.RING, ords)[1]) if np.isnan(CO[0]): start_point_guess = np.nanmean(T, axis=(1, 2, 3)) - CO = np.squeeze(findorbit6(SC.RING, ords, guess=start_point_guess)[1]) + CO = np.squeeze(at_wrapper.findorbit6(SC.RING, ords, guess=start_point_guess)[1]) if np.isnan(CO[0]): CO = np.full(6, np.nan) else: diff --git a/pySC/plotting/plot_support.py b/pySC/plotting/plot_support.py index 1378b61..1b9bffd 100644 --- a/pySC/plotting/plot_support.py +++ b/pySC/plotting/plot_support.py @@ -1,7 +1,7 @@ import matplotlib.pyplot as plt import numpy as np import copy -from pySC.utils.at_wrapper import findspos, atgetfieldvalues +from pySC.utils import at_wrapper from pySC.core.constants import SUPPORT_TYPES from typing import Tuple from pySC.core.simulated_commissioning import SimulatedCommissioning @@ -28,7 +28,7 @@ def plot_support(SC: SimulatedCommissioning, font_size: int = 8, x_lim: Tuple[fl init_font = plt.rcParams["font.size"] plt.rcParams.update({'font.size': font_size}) # Get s - positions along the lattice - s_pos = findspos(SC.RING) + s_pos = at_wrapper.findspos(SC.RING) circumference = s_pos[-1] if x_lim is None: x_lim = (0, circumference) @@ -53,20 +53,20 @@ def plot_support(SC: SimulatedCommissioning, font_size: int = 8, x_lim: Tuple[fl # Magnet offsets and rolls off_support_line, roll_support_line = SC.support_offset_and_roll(s) - off_mag_support = atgetfieldvalues(SC.RING, SC.ORD.Magnet, "SupportOffset") - roll_mag_support = atgetfieldvalues(SC.RING, SC.ORD.Magnet, "SupportRoll") - off_mag_individual = atgetfieldvalues(SC.RING, SC.ORD.Magnet, "MagnetOffset") - roll_mag_individual = atgetfieldvalues(SC.RING, SC.ORD.Magnet, "MagnetRoll") - off_mag_total = atgetfieldvalues(SC.RING, SC.ORD.Magnet, "T2")[:, [0, 2, 5]] + off_mag_support = at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.Magnet, "SupportOffset") + roll_mag_support = at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.Magnet, "SupportRoll") + off_mag_individual = at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.Magnet, "MagnetOffset") + roll_mag_individual = at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.Magnet, "MagnetRoll") + off_mag_total = at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.Magnet, "T2")[:, [0, 2, 5]] roll_mag_total = roll_mag_support + roll_mag_individual # BPM offsets and rolls # Longitudinal offsets and Pitch and Yaw angles not supported for BPMs pad_off, pad_roll = ((0, 0), (0, 1)), ((0, 0), (0, 2)) - off_bpm = np.pad(atgetfieldvalues(SC.RING, SC.ORD.BPM, "Offset"), pad_off) - roll_bpm = np.pad(atgetfieldvalues(SC.RING, SC.ORD.BPM, "Roll"), pad_roll) - off_bpm_support = np.pad(atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportOffset"), pad_off) - roll_bpm_support = np.pad(atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportRoll"), pad_roll) + off_bpm = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "Offset"), pad_off) + roll_bpm = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "Roll"), pad_roll) + off_bpm_support = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportOffset"), pad_off) + roll_bpm_support = np.pad(at_wrapper.atgetfieldvalues(SC.RING, SC.ORD.BPM, "SupportRoll"), pad_roll) # create figure fig, ax = plt.subplots(nrows=9, ncols=2, num=1213, sharex="all", figsize=(10, 15)) diff --git a/pySC/plotting/plot_synoptic.py b/pySC/plotting/plot_synoptic.py index 9f60b60..166cee1 100644 --- a/pySC/plotting/plot_synoptic.py +++ b/pySC/plotting/plot_synoptic.py @@ -6,7 +6,7 @@ from matplotlib.collections import PatchCollection from matplotlib.patches import Polygon -from pySC.utils.at_wrapper import findspos +from pySC.utils import at_wrapper MONITOR = "monitor" NORMAL_KWARGS = ("dipole", "quadrupole", "sextupole", "multipole", MONITOR) @@ -27,7 +27,7 @@ def plot_synoptic(SC, axes: matplotlib.axes.Axes = None, correctors: bool = Fals if axes is None: axsyn.set_zorder(-0.2) # Put synoptic in the background - s_pos = findspos(ring)[:-1] + s_pos = at_wrapper.findspos(ring)[:-1] if famnames: # get fam name of each main multipole diff --git a/pySC/utils/classdef_tools.py b/pySC/utils/classdef_tools.py index 782fd14..9641d74 100644 --- a/pySC/utils/classdef_tools.py +++ b/pySC/utils/classdef_tools.py @@ -1,7 +1,7 @@ import numpy as np from pySC.core.constants import SETTING_METHODS -from pySC.utils.sc_tools import SCrandnc +from pySC.utils import sc_tools def update_double_ordinates(ords1, ords2): @@ -16,8 +16,8 @@ def intersect(primary, secondary): def randn_cutoff(field, default_cut_off): if isinstance(field, list): - return field[0] * SCrandnc(field[1], np.shape(field[0])) - return field * SCrandnc(default_cut_off, np.shape(field)) + return field[0] * sc_tools.randnc(field[1], np.shape(field[0])) + return field * sc_tools.randnc(default_cut_off, np.shape(field)) def add_padded(v1, v2): diff --git a/pySC/utils/sc_tools.py b/pySC/utils/sc_tools.py index 93fe674..24f7a86 100644 --- a/pySC/utils/sc_tools.py +++ b/pySC/utils/sc_tools.py @@ -1,17 +1,17 @@ import re import numpy as np -from at import Lattice from matplotlib import pyplot as plt from numpy import ndarray +from at import Lattice + +from pySC.utils import at_wrapper, logging_tools -from pySC.utils.at_wrapper import findspos -from pySC.utils import logging_tools LOGGER = logging_tools.get_logger(__name__) -def SCrandnc(cut_off: float = 2, shape: tuple = (1, )) -> ndarray: +def randnc(cut_off: float = 2, shape: tuple = (1,)) -> ndarray: """ Generates an array of random number(s) from normal distribution with a cut-off. @@ -36,7 +36,7 @@ def SCrandnc(cut_off: float = 2, shape: tuple = (1, )) -> ndarray: return out.reshape(out_shape) -def SCgetOrds(ring: Lattice, regex: str) -> ndarray: +def ords_from_regex(ring: Lattice, regex: str) -> ndarray: """ Returns the indices of the elements in the ring whose names match the regex. @@ -58,7 +58,7 @@ def SCgetOrds(ring: Lattice, regex: str) -> ndarray: return indices -def SCgetPinv(matrix: ndarray, num_removed: int = 0, alpha: float = 0, damping: float = 1, plot: bool = False) -> ndarray: +def pinv(matrix: ndarray, num_removed: int = 0, alpha: float = 0, damping: float = 1, plot: bool = False) -> ndarray: """ Computes the pseudo-inverse of a matrix using the Singular Value Decomposition (SVD) method. @@ -93,11 +93,11 @@ def SCgetPinv(matrix: ndarray, num_removed: int = 0, alpha: float = 0, damping: return matrix_inv -def SCscaleCircumference(RING, circ, mode='abs'): # TODO +def scale_circumference(RING, circ, mode='abs'): # TODO allowed_modes = ("abs", "rel") if mode not in allowed_modes: raise ValueError(f'Unsupported circumference scaling mode: ``{mode}``. Allowed are {allowed_modes}.') - C = findspos(RING)[-1] + C = at_wrapper.findspos(RING)[-1] D = 0 for ind in range(len(RING)): if RING[ind].PassMethod == 'DriftPass': @@ -142,7 +142,7 @@ def update_transformation(element): return element -def SCmultipolesRead(fname): # TODO sample of the input anywhere? +def read_multipoles(fname): # TODO sample of the input anywhere? f = open(fname, 'r') tab = np.array(f.read().split()).astype(float) f.close() diff --git a/pySC/utils/stats.py b/pySC/utils/stats.py index 38e9766..fafc5d8 100644 --- a/pySC/utils/stats.py +++ b/pySC/utils/stats.py @@ -16,7 +16,7 @@ CONFIDENCE_LEVEL = (1 + erf(1 / np.sqrt(2))) / 2 PI2: float = 2 * np.pi -PI2I: float = 2j * np.pi +PI2I: complex = 2j * np.pi def circular_mean(data, period=PI2, errors=None, axis=None): diff --git a/tests/test_beam.py b/tests/test_beam.py index a1b9ce3..e0af495 100644 --- a/tests/test_beam.py +++ b/tests/test_beam.py @@ -2,8 +2,8 @@ import numpy as np from numpy.testing import assert_equal from pySC.core.beam import bpm_reading, all_elements_reading, beam_transmission, generate_bunches -from pySC.utils.sc_tools import SCgetOrds from pySC.core.simulated_commissioning import SimulatedCommissioning +from pySC.utils import sc_tools from tests.test_at_wrapper import at_lattice @@ -55,7 +55,7 @@ def test_generate_bunches_basic(sc): @pytest.fixture def sc(at_lattice): SC = SimulatedCommissioning(at_lattice) - SC.register_bpms(SCgetOrds(SC.RING, 'BPM'), CalError=5E-2 * np.ones(2), + SC.register_bpms(sc_tools.ords_from_regex(SC.RING, 'BPM'), CalError=5E-2 * np.ones(2), Offset=500E-6 * np.ones(2), Noise=10E-6 * np.ones(2), NoiseCO=0E-6 * np.ones(2), diff --git a/tests/test_example.py b/tests/test_example.py index 51e6a7c..7341507 100644 --- a/tests/test_example.py +++ b/tests/test_example.py @@ -7,7 +7,7 @@ from pySC.correction.tune import tune_scan from pySC.correction.bba import trajectory_bba, fake_bba, _get_bpm_offset_from_mag from pySC.lattice_properties.response_model import SCgetModelRM, SCgetModelDispersion -from pySC.utils.sc_tools import SCgetOrds +from pySC.utils import sc_tools from pySC.core.lattice_setting import switch_cavity_and_radiation from pySC.correction.rf import correct_rf_phase, correct_rf_frequency @@ -15,40 +15,42 @@ def test_example(at_lattice): np.random.seed(12345678) sc = SimulatedCommissioning(at_lattice) - sc.register_bpms(SCgetOrds(sc.RING, 'BPM'), + sc.register_bpms(sc_tools.ords_from_regex(sc.RING, 'BPM'), CalError=5E-2 * np.ones(2), Offset=500E-6 * np.ones(2), Noise=10E-6 * np.ones(2), NoiseCO=1E-6 * np.ones(2), Roll=1E-3) - sc.register_magnets(SCgetOrds(sc.RING, 'QF'), + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'QF'), HCM=1E-3, CalErrorB=np.array([5E-2, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - sc.register_magnets(SCgetOrds(sc.RING, 'QD'), VCM=1E-3, + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'QD'), VCM=1E-3, CalErrorA=np.array([5E-2, 0]), CalErrorB=np.array([0, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - sc.register_magnets(SCgetOrds(sc.RING, 'BEND'), + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'BEND'), BendingAngle=1E-3, # relative MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - sc.register_magnets(SCgetOrds(sc.RING, 'SF|SD'), + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'SF|SD'), SkewQuad=0.1, CalErrorA=np.array([0, 1E-3, 0]), CalErrorB=np.array([0, 0, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - sc.register_cavities(SCgetOrds(sc.RING, 'RFCav'), FrequencyOffset=5E3, + sc.register_cavities(sc_tools.ords_from_regex(sc.RING, 'RFCav'), FrequencyOffset=5E3, VoltageOffset=5E3, TimeLagOffset=0.5) - sc.register_supports(np.vstack((SCgetOrds(sc.RING, 'GirderStart'), SCgetOrds(sc.RING, 'GirderEnd'))), + sc.register_supports(np.vstack((sc_tools.ords_from_regex(sc.RING, 'GirderStart'), + sc_tools.ords_from_regex(sc.RING, 'GirderEnd'))), "Girder", Offset=100E-6 * np.array([1, 1, 0]), Roll=200E-6 * np.array([1, 0, 0])) - sc.register_supports(np.vstack((SCgetOrds(sc.RING, 'SectionStart'), SCgetOrds(sc.RING, 'SectionEnd'))), + sc.register_supports(np.vstack((sc_tools.ords_from_regex(sc.RING, 'SectionStart'), + sc_tools.ords_from_regex(sc.RING, 'SectionEnd'))), "Section", Offset=100E-6 * np.array([1, 1, 0])) sc.INJ.beamSize = np.diag(np.array([200E-6, 100E-6, 100E-6, 50E-6, 1E-3, 1E-4]) ** 2) @@ -56,16 +58,16 @@ def test_example(at_lattice): sc.SIG.staticInjectionZ = np.array([1E-3, 1E-4, 1E-3, 1E-4, 1E-3, 1E-3]) sc.SIG.Circumference = 2E-4 sc.INJ.beamLostAt = 0.6 - for ord in SCgetOrds(sc.RING, 'Drift'): + for ord in sc_tools.ords_from_regex(sc.RING, 'Drift'): sc.RING[ord].EApertures = 13E-3 * np.array([1, 1]) - for ord in SCgetOrds(sc.RING, 'QF|QD|BEND|SF|SD'): + for ord in sc_tools.ords_from_regex(sc.RING, 'QF|QD|BEND|SF|SD'): sc.RING[ord].EApertures = 10E-3 * np.array([1, 1]) sc.RING[sc.ORD.Magnet[50]].EApertures = np.array([6E-3, 3E-3]) sc.apply_errors() sc.verify_structure() sc.RING = switch_cavity_and_radiation(sc.RING, 'cavityoff') - sext_ords = SCgetOrds(sc.RING, 'SF|SD') + sext_ords = sc_tools.ords_from_regex(sc.RING, 'SF|SD') 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) @@ -82,11 +84,11 @@ def test_example(at_lattice): # Performing BBA sc.INJ.nParticles = 1 - quadOrds = np.tile(SCgetOrds(sc.RING, 'QF|QD'), (2, 1)) + quadOrds = np.tile(sc_tools.ords_from_regex(sc.RING, 'QF|QD'), (2, 1)) BPMords = np.tile(sc.ORD.BPM, (2, 1)) init_offsets = _get_bpm_offset_from_mag(sc.RING, BPMords, quadOrds) - sc, bba_offsets, bba_offset_errors = trajectory_bba(sc, np.tile(sc.ORD.BPM, (2, 1)), np.tile(SCgetOrds(sc.RING, 'QF|QD'), (2, 1)), - q_ord_phase=SCgetOrds(sc.RING, 'QF|QD')[0], + sc, bba_offsets, bba_offset_errors = trajectory_bba(sc, np.tile(sc.ORD.BPM, (2, 1)), np.tile(sc_tools.ords_from_regex(sc.RING, 'QF|QD'), (2, 1)), + q_ord_phase=sc_tools.ords_from_regex(sc.RING, 'QF|QD')[0], q_ord_setpoints=np.array([0.8, 0.9, 1.0, 1.1, 1.2]), plot_results=True, magnet_strengths=np.array([0.8, 0.9, 1.0, 1.1, 1.2]), @@ -107,7 +109,7 @@ def test_example(at_lattice): sc.INJ.nTurns = 15 sc = correct_rf_frequency(sc, n_steps=15, f_range=4E3 * np.array([-1, 1])) - sc = fake_bba(sc, np.tile(sc.ORD.BPM, (2, 1)), np.tile(SCgetOrds(sc.RING, 'QF|QD'), (2, 1)), + sc = fake_bba(sc, np.tile(sc.ORD.BPM, (2, 1)), np.tile(sc_tools.ords_from_regex(sc.RING, 'QF|QD'), (2, 1)), fake_offset=np.array([50E-6, 50E-6])) # Orbit correction @@ -127,7 +129,7 @@ def test_example(at_lattice): sc = cur sc.RING = switch_cavity_and_radiation(sc.RING, 'cavityon') max_turns, fraction_survived = beam_transmission(sc, nParticles=100, nTurns=200, plot=True) - sc, _, _, _ = tune_scan(sc, np.vstack((SCgetOrds(sc.RING, 'QF'), SCgetOrds(sc.RING, 'QD'))), + sc, _, _, _ = tune_scan(sc, np.vstack((sc_tools.ords_from_regex(sc.RING, 'QF'), sc_tools.ords_from_regex(sc.RING, 'QD'))), np.outer(np.ones(2), 1 + np.linspace(-0.01, 0.01, 51)), do_plot=False, nParticles=50, nTurns=100, target=0.95) max_turns, fraction_survived = beam_transmission(sc, nParticles=100, nTurns=200, plot=True) diff --git a/tests/test_injection_fit.py b/tests/test_injection_fit.py index 74c50c0..e8af03e 100644 --- a/tests/test_injection_fit.py +++ b/tests/test_injection_fit.py @@ -1,6 +1,6 @@ import pytest import numpy as np -from pySC.utils.sc_tools import SCgetOrds +from pySC.utils import sc_tools from pySC.correction.injection_fit import fit_injection_trajectory, fit_injection_drift from pySC.core.simulated_commissioning import SimulatedCommissioning from tests.test_at_wrapper import at_lattice @@ -20,7 +20,7 @@ def test_fit_injection_drift(sc): @pytest.fixture def sc(at_lattice): SC = SimulatedCommissioning(at_lattice) - SC.register_bpms(SCgetOrds(SC.RING, 'BPM'), CalError=5E-2 * np.ones(2), + SC.register_bpms(sc_tools.ords_from_regex(SC.RING, 'BPM'), CalError=5E-2 * np.ones(2), Offset=200E-6 * np.ones(2), Noise=10E-6 * np.ones(2), Roll=1E-3) diff --git a/tests/test_lattice_setting.py b/tests/test_lattice_setting.py index cfafe5a..5ee51ea 100644 --- a/tests/test_lattice_setting.py +++ b/tests/test_lattice_setting.py @@ -4,17 +4,16 @@ import numpy as np from numpy.testing import assert_equal, assert_allclose from tests.test_at_wrapper import at_lattice -from pySC.utils.sc_tools import SCgetOrds from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.utils.classdef_tools import check_input_and_setpoints from pySC.core.constants import SETTING_ABS, SETTING_REL, SETTING_ADD +from pySC.utils import sc_tools, classdef_tools 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(indices, setpoints, True, method="add") + sc.set_cm_setpoints(indices, setpoints, True, method=SETTING_ADD) assert_equal(indices, np.arange(11, 450, 22, dtype=int)) assert_equal(setpoints, 1e-4 * np.ones(len(indices))) @@ -26,7 +25,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(indices, setpoints, True, 1, method="add") + sc.set_magnet_setpoints(indices, setpoints, True, 1, method=SETTING_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 +38,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(indices, setpoints, "Frequency", method="add") + sc.set_cavity_setpoints(indices, setpoints, "Frequency", method=SETTING_ADD) assert_equal(indices, np.zeros(1, dtype=int)) assert_equal(setpoints, np.ones(1)) assert sc.RING.__repr__() != lattice_copy.__repr__() @@ -49,19 +48,19 @@ 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 = classdef_tools.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 = classdef_tools.check_input_and_setpoints(SETTING_ABS, os, ss) assert_equal(ords1d, wanted_ords) assert_equal(setps1d, wanted_setpoints) # check also empty ord wanted_ord, wanted_setpoint = np.array([], dtype=int), np.array([]) - ord1d, setp1d = check_input_and_setpoints(SETTING_REL, [], 1.0) + ord1d, setp1d = classdef_tools.check_input_and_setpoints(SETTING_REL, [], 1.0) assert_equal(ord1d, wanted_ord) assert_equal(setp1d, wanted_setpoint) @@ -112,16 +111,16 @@ def unit_sc(): @pytest.fixture def sc(at_lattice): SC = SimulatedCommissioning(at_lattice) - SC.register_magnets(SCgetOrds(SC.RING, 'QF'), HCM=1E-3, + SC.register_magnets(sc_tools.ords_from_regex(SC.RING, 'QF'), HCM=1E-3, CalErrorB=np.array([5E-2, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - SC.register_magnets(SCgetOrds(SC.RING, 'QD'), VCM=1E-3, + SC.register_magnets(sc_tools.ords_from_regex(SC.RING, 'QD'), VCM=1E-3, CalErrorA=np.array([5E-2, 0]), CalErrorB=np.array([0, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - SC.register_cavities(SCgetOrds(SC.RING, 'RFCav'), + SC.register_cavities(sc_tools.ords_from_regex(SC.RING, 'RFCav'), FrequencyOffset=5E3, VoltageOffset=5E3, TimeLagOffset=0.5) diff --git a/tests/test_loco.py b/tests/test_loco.py index 9422ca0..392e107 100644 --- a/tests/test_loco.py +++ b/tests/test_loco.py @@ -3,8 +3,7 @@ import pytest from pathlib import Path from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.utils.sc_tools import SCgetOrds -from pySC.utils import logging_tools +from pySC.utils import logging_tools, sc_tools from pySC.correction import loco from pySC.lattice_properties.response_model import SCgetModelRM, SCgetModelDispersion @@ -15,21 +14,21 @@ def test_loco_hmba(at_ring): np.random.seed(12345678) sc = SimulatedCommissioning(at_ring) - sc.register_bpms(SCgetOrds(sc.RING, 'BPM'), Roll=0.0, CalError=1E-2 * np.ones(2), NoiseCO=np.array([1e-7, 1E-7])) - sc.register_magnets(SCgetOrds(sc.RING, 'QF|QD'), CalErrorB=np.array([0, 1E-2])) # relative - sc.register_magnets(SCgetOrds(sc.RING, 'CXY'), CalErrorA=np.array([1E-4, 0]), CalErrorB=np.array([1E-4, 0])) - sc.register_magnets(SCgetOrds(sc.RING, 'BEND')) - sc.register_magnets(SCgetOrds(sc.RING, 'SF|SD')) # [1/m] - sc.register_cavities(SCgetOrds(sc.RING, 'RFC')) + sc.register_bpms(sc_tools.ords_from_regex(sc.RING, 'BPM'), Roll=0.0, CalError=1E-2 * np.ones(2), NoiseCO=np.array([1e-7, 1E-7])) + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'QF|QD'), CalErrorB=np.array([0, 1E-2])) # relative + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'CXY'), CalErrorA=np.array([1E-4, 0]), CalErrorB=np.array([1E-4, 0])) + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'BEND')) + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'SF|SD')) # [1/m] + sc.register_cavities(sc_tools.ords_from_regex(sc.RING, 'RFC')) sc.apply_errors() - cor_ords = SCgetOrds(sc.RING, 'CXY') + cor_ords = sc_tools.ords_from_regex(sc.RING, 'CXY') used_correctors1 = loco.select_equally_spaced_elements(cor_ords, 10) used_correctors2 = loco.select_equally_spaced_elements(cor_ords, 10) used_cor_ords = [used_correctors1, used_correctors2] used_bpms_ords = loco.select_equally_spaced_elements(sc.ORD.BPM, len(sc.ORD.BPM)) - cav_ords = SCgetOrds(sc.RING, 'RFC') - quads_ords = [SCgetOrds(sc.RING, 'QF'), SCgetOrds(sc.RING, 'QD')] + cav_ords = sc_tools.ords_from_regex(sc.RING, 'RFC') + quads_ords = [sc_tools.ords_from_regex(sc.RING, 'QF'), sc_tools.ords_from_regex(sc.RING, 'QD')] CMstep = np.array([1.e-4]) # correctors change [rad] dk = 1.e-4 # quads change diff --git a/tests/test_pinv.py b/tests/test_pinv.py index 91ced9f..81a45c6 100644 --- a/tests/test_pinv.py +++ b/tests/test_pinv.py @@ -1,15 +1,15 @@ import pytest import numpy as np -from pySC.utils.sc_tools import SCgetPinv +from pySC.utils.sc_tools import pinv def test_pinv(matrix): - a = SCgetPinv(matrix, num_removed=0, alpha=0, damping=1, plot=False) - a2 = SCgetPinv(matrix, num_removed=0, alpha=0, damping=0.9, plot=True) + a = pinv(matrix, num_removed=0, alpha=0, damping=1, plot=False) + a2 = pinv(matrix, num_removed=0, alpha=0, damping=0.9, plot=True) b = np.linalg.pinv(matrix) assert np.allclose(a, b) assert np.allclose(a2, b * 0.9) - a3 = SCgetPinv(matrix, num_removed=0, alpha=0.1, damping=0.9, plot=True) + a3 = pinv(matrix, num_removed=0, alpha=0.1, damping=0.9, plot=True) @pytest.fixture diff --git a/tests/test_plotting.py b/tests/test_plotting.py index ebb8eb2..037d055 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -4,8 +4,7 @@ from pySC.plotting.plot_lattice import plot_lattice, plot_cm_strengths, _get_s_range from pySC.plotting.plot_support import plot_support from pySC.plotting.plot_phase_space import plot_phase_space -from pySC.utils.sc_tools import SCgetOrds -from pySC.utils.at_wrapper import findspos +from pySC.utils import at_wrapper, sc_tools from pySC.core.simulated_commissioning import SimulatedCommissioning from tests.test_at_wrapper import at_lattice @@ -18,7 +17,7 @@ def test_plot_lattice(sc): def test_get_s_range(sc): s_range = np.array([0, 20]) indices = np.array([30, 130], dtype=int) - s_indices = findspos(sc.RING)[indices] + s_indices = at_wrapper.findspos(sc.RING)[indices] n_sectors = 10 assert_equal(_get_s_range(sc.RING, s_range, indices, n_sectors), s_range) assert_equal(_get_s_range(sc.RING, None, indices, n_sectors), s_indices) @@ -41,29 +40,30 @@ def test_plot_support(sc): @pytest.fixture def sc(at_lattice): SC = SimulatedCommissioning(at_lattice) - SC.register_bpms(SCgetOrds(SC.RING, 'BPM'), CalError=5E-2 * np.ones(2), + SC.register_bpms(sc_tools.ords_from_regex(SC.RING, 'BPM'), CalError=5E-2 * np.ones(2), Offset=500E-6 * np.ones(2), Noise=10E-6 * np.ones(2), NoiseCO=0E-6 * np.ones(2), Roll=1E-3) - SC.register_magnets(SCgetOrds(SC.RING, 'QF'), + SC.register_magnets(sc_tools.ords_from_regex(SC.RING, 'QF'), HCM=1E-3, CalErrorB=np.array([5E-2, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - SC.register_magnets(SCgetOrds(SC.RING, 'QD'), VCM=1E-3, + SC.register_magnets(sc_tools.ords_from_regex(SC.RING, 'QD'), VCM=1E-3, CalErrorA=np.array([5E-2, 0]), CalErrorB=np.array([0, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), MagnetRoll=200E-6 * np.array([1, 0, 0])) - SC.register_cavities(SCgetOrds(SC.RING, 'RFCav'), FrequencyOffset=5E3, + SC.register_cavities(sc_tools.ords_from_regex(SC.RING, 'RFCav'), FrequencyOffset=5E3, VoltageOffset=5E3, TimeLagOffset=0.5) - SC.register_supports(np.vstack((SCgetOrds(SC.RING, 'GirderStart'), SCgetOrds(SC.RING, 'GirderEnd'))), + SC.register_supports(np.vstack((sc_tools.ords_from_regex(SC.RING, 'GirderStart'), + sc_tools.ords_from_regex(SC.RING, 'GirderEnd'))), "Girder", Offset=100E-6 * np.array([1, 1, 0]), Roll=200E-6 * np.array([1, 0, 0])) - for e_ord in SCgetOrds(SC.RING, 'Drift'): + for e_ord in sc_tools.ords_from_regex(SC.RING, 'Drift'): SC.RING[e_ord].EApertures = 13E-3 * np.array([1, 1]) np.random.seed(123) SC.apply_errors() diff --git a/tests/test_simulated_commissioning.py b/tests/test_simulated_commissioning.py index 362708d..5443deb 100644 --- a/tests/test_simulated_commissioning.py +++ b/tests/test_simulated_commissioning.py @@ -5,7 +5,7 @@ from numpy.testing import assert_allclose, assert_equal import at from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.utils.sc_tools import SCgetOrds +from pySC.utils import sc_tools def test_simulated_commissioning_data_structure(at_cell): @@ -55,7 +55,7 @@ def test_register_cavities(at_cell): def test_register_magnets(at_cell): sc = SimulatedCommissioning(at_cell) mag_dict = dict(CalErrorB=np.array([5E-2, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]),) - indices = SCgetOrds(sc.RING, "Q") + indices = sc_tools.ords_from_regex(sc.RING, "Q") sc.register_magnets(np.repeat(indices, 2), CalErrorB=np.array([5E-2, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]),) update_dict = dict(CalErrorA=5E-2) sc.register_magnets(indices[2::-1], CalErrorA=5E-2) @@ -68,7 +68,7 @@ def test_register_magnets(at_cell): def test_register_supports(at_cell): sc = SimulatedCommissioning(at_cell) mag_dict = dict(SectionOffset=200E-6 * np.array([1, 1, 0]),) - indices = np.vstack((SCgetOrds(sc.RING, "SF"), SCgetOrds(sc.RING, "SD"))) + indices = np.vstack((sc_tools.ords_from_regex(sc.RING, "SF"), sc_tools.ords_from_regex(sc.RING, "SD"))) sc.register_supports(np.repeat(indices, 2, axis=1), "Section", Offset=200E-6 * np.array([1, 1, 0]),) update_dict1 = dict(SectionRoll=5E-2 * np.array([1, 0, 0]), SectionOffset=100E-6 * np.array([1, 1, 0])) update_dict2 = dict(SectionRoll=5E-2 * np.array([1, 0, 0])) @@ -105,9 +105,9 @@ def test_apply_errors(at_cell): sc.register_bpms(np.array([8, 21]), CalError=5E-2 * np.ones(2), Offset=500E-6 * np.ones(2), Noise=10E-6 * np.ones(2), NoiseCO=1E-6 * np.ones(2), Roll=1E-3) sc.register_cavities(np.array([0]), FrequencyOffset=5E3, VoltageOffset=5E3, TimeLagOffset=0.5) - sc.register_magnets(SCgetOrds(sc.RING, "Q"), CalErrorB=np.array([5E-2, 1E-3]), + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, "Q"), CalErrorB=np.array([5E-2, 1E-3]), MagnetOffset=200E-6 * np.array([1, 1, 0]), ) - indices = np.vstack((SCgetOrds(sc.RING, 'SectionStart'), SCgetOrds(sc.RING, 'SectionEnd'))) + indices = np.vstack((sc_tools.ords_from_regex(sc.RING, 'SectionStart'), sc_tools.ords_from_regex(sc.RING, 'SectionEnd'))) sc.register_supports(indices, "Section", Offset=200E-6 * np.array([1, 1, 0]), Roll=20E-6 * np.array([1, 0, 0])) sc.apply_errors() # BPM at index 8