From a17b76b35f054b05e87e8f3bb1fb65a6f2074102 Mon Sep 17 00:00:00 2001 From: malina Date: Wed, 24 Jan 2024 14:28:49 +0100 Subject: [PATCH] Fixing documentation, imports and names --- pySC/correction/bba.py | 6 +- pySC/correction/orbit_trajectory.py | 86 ++++++++++++++--------------- pySC/correction/ramp_errors.py | 7 +-- pySC/example.py | 22 +++----- pySC/matlab_index.py | 21 ++++--- tests/test_example.py | 19 +++---- 6 files changed, 76 insertions(+), 85 deletions(-) diff --git a/pySC/correction/bba.py b/pySC/correction/bba.py index 4b092c5..925509d 100644 --- a/pySC/correction/bba.py +++ b/pySC/correction/bba.py @@ -8,7 +8,7 @@ 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.correction.orbit_trajectory import SCfeedbackRun +from pySC.correction import orbit_trajectory LOGGER = logging_tools.get_logger(__name__) @@ -344,8 +344,8 @@ def _get_orbit_bump(SC, cm_ord, bpm_ord, n_dim, par): # TODO W0[n_dim, max(1, tmpBPMind - par.orbBumpWindow):(tmpBPMind - 1)] = 0 W0[n_dim, (tmpBPMind + 1):min(len(par.RMstruct.BPMords), tmpBPMind + par.orbBumpWindow)] = 0 - CUR = SCfeedbackRun(SC, par.RMstruct.MinvCO, reference=R0, cm_ords=cm_ords, bpm_ords=par.RMstruct.BPMords, eps=1E-6, - target=0, maxsteps=50, scaleDisp=par.RMstruct.scaleDisp, ) + CUR = orbit_trajectory.correct(SC, par.RMstruct.RM, reference=R0, cm_ords=cm_ords, bpm_ords=par.RMstruct.BPMords, eps=1E-6, + target=0, maxsteps=50, scaleDisp=par.RMstruct.scaleDisp, ) cm_vec = [] factor = np.linspace(-1, 1, par.n_steps) for n_dim in range(2): diff --git a/pySC/correction/orbit_trajectory.py b/pySC/correction/orbit_trajectory.py index e60b020..1cf80eb 100644 --- a/pySC/correction/orbit_trajectory.py +++ b/pySC/correction/orbit_trajectory.py @@ -23,20 +23,20 @@ WIGGLE_ANGLE_RANGE: Tuple[float, float] = (500E-6, 1000E-6) -def SCfeedbackFirstTurn(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, maxsteps=100, **pinv_params): +def first_turn(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, maxsteps=100, **pinv_params): """ Achieves one-turn transmission Achieves a first turn in `SC.RING`. This algorithm is based on the idea that repeated trajectory corrections calculated via a suitably regularized - pseudo-inverse trajectory-response matrix `Mplus` will drive the BPM readings + pseudo-inverse of `response matrix` will drive the BPM readings and CM settings to a fixed point. - lim_{n->oo} B_n = const. , with B_{n+1} = Phi(Mplus . B_{n} ), (1) + lim_{n->oo} B_n = const. , with B_{n+1} = Phi(response_matrix^{-1} . B_{n} ), (1) where mapping Phi maps corrector kicks to BPM-readings. The RMS-values of both, BPM readings and CM settings, are determined by the - regularization of Mplus. Successively - during the course of repeated + regularization of pseudo-inverted response matrix. Successively - during the course of repeated application of (1) - more and more transmission is achieved throughout the ring, more magnets are traversed near their magnetic center (which is hopefully at least somewhere near the BPM zero-point), resulting in decreased kicks. @@ -46,7 +46,7 @@ def SCfeedbackFirstTurn(SC, response_matrix, reference=None, cm_ords=None, bpm_o Args: SC: SimulatedCommissioning class instance. - Mplus: Pseudo-inverse trajectory-response matrix. + response_matrix: Trajectory-response matrix. reference: (None) target orbit in the format `[x_1 ... x_n y_1 ...y_n]`, where [x_i,y_i]` is the target position at the i-th BPM. cm_ords: List of CM ordinates to be used for correction (SC.ORD.CM) @@ -57,14 +57,14 @@ def SCfeedbackFirstTurn(SC, response_matrix, reference=None, cm_ords=None, bpm_o SimulatedCommissioning class instance with corrected `SC.RING` """ - LOGGER.debug('SCfeedbackFirstTurn: Start') + 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) for n in range(maxsteps): if transmission_history[-1] == 0: - raise RuntimeError('SCfeedbackFirstTurn: FAIL (no BPM reading to begin with)') + raise RuntimeError('First turn threading: FAIL (no BPM reading to begin with)') # Set BPM readings measurement = bpm_readings[:, :].reshape(reference.shape) @@ -82,15 +82,15 @@ def SCfeedbackFirstTurn(SC, response_matrix, reference=None, cm_ords=None, bpm_o # Check stopping criteria if _is_repro(transmission_history, NREPRO) and transmission_history[-1] == bpm_readings.shape[1]: # last three the same and full transmission - LOGGER.debug('SCfeedbackFirstTurn: Success') + LOGGER.debug('First turn threading: Success') return SC if _is_repro(transmission_history, WIGGLE_AFTER): SC = _wiggling(SC, bpm_ords, cm_ords, transmission_history[-1] + 1) - raise RuntimeError('SCfeedbackFirstTurn: FAIL (maxsteps reached)') + raise RuntimeError('First turn threading: FAIL (maxsteps reached)') -def SCfeedbackStitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, nBPMs=4, maxsteps=30, **pinv_params): +def stitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, n_bpms=4, maxsteps=30, **pinv_params): """ Achieves 2-turn transmission @@ -104,11 +104,12 @@ def SCfeedbackStitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords Args: SC: SimulatedCommissioning class instance. - Mplus: Pseudo-inverse trajectory-response matrix. + response_matrix: Trajectory-response matrix. reference: (None) target orbit in the format `[x_1 ... x_n y_1 ...y_n]`, where [x_i,y_i]` is the target position at the i-th BPM. cm_ords: List of CM ordinates to be used for correction (SC.ORD.CM) bpm_ords: List of BPM ordinates at which the reading should be evaluated (SC.ORD.BPM) + n_bpms: Number of BPMs to match the trajectory in first and second turn maxsteps: break, if this number of correction steps have been performed (default = 100) Returns: @@ -120,17 +121,16 @@ def SCfeedbackStitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords three BPMs, for a maximum of 20 steps:: RM2 = SCgetModelRM(SC, SC.ORD.BPM, SC.ORD.CM, nTurns=2) - Minv2 = SCgetPinv(RM2, alpha=10) SC.INJ.nTurns = 2 - SC = SCfeedbackStitch(SC, Minv2, nBPMs=3, maxsteps=20) + SC = stitch(SC, RM2, alpha=10, nBPMs=3, maxsteps=20) """ - LOGGER.debug('SCfeedbackStitch: Start') + LOGGER.debug('Second turn stitching: Start') if SC.INJ.nTurns != 2: raise ValueError("Stitching works only with two turns.") 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... - transmission_limit = len(bpm_ords) + nBPMs + transmission_limit = len(bpm_ords) + n_bpms if transmission_history[-1] < len(bpm_ords): raise RuntimeError("Stitching works only with full 1st turn transmission.") @@ -151,9 +151,9 @@ def SCfeedbackStitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords # Main loop for steps in range(maxsteps): # Set BPM readings - delta_b = np.zeros((2,len(bpm_ords))) - delta_b[0][:nBPMs] = bpm_readings[0][len(bpm_ords):len(bpm_ords) + nBPMs] - bpm_readings[0][:nBPMs] - delta_b[1][:nBPMs] = bpm_readings[1][len(bpm_ords):len(bpm_ords) + nBPMs] - bpm_readings[1][:nBPMs] + delta_b = np.zeros((2, len(bpm_ords))) + delta_b[0][:n_bpms] = bpm_readings[0][len(bpm_ords):len(bpm_ords) + n_bpms] - bpm_readings[0][:n_bpms] + delta_b[1][:n_bpms] = bpm_readings[1][len(bpm_ords):len(bpm_ords) + n_bpms] - bpm_readings[1][:n_bpms] measurement = np.concatenate((bpm_readings[:, :len(bpm_ords)], delta_b), axis=1).ravel() measurement[np.isnan(measurement)] = 0 @@ -166,14 +166,14 @@ def SCfeedbackStitch(SC, response_matrix, reference=None, cm_ords=None, bpm_ords # Check stopping criteria if transmission_history[-1] < transmission_history[-2]: - RuntimeError('SCfeedbackStitch: FAIL Setback') + RuntimeError('Second turn stitching: FAIL Setback') if _is_repro(transmission_history, NREPRO) and transmission_history[-1] == bpm_readings.shape[1]: - LOGGER.debug('SCfeedbackStitch: Success') + LOGGER.debug('Second turn stitching: Success') return SC - raise RuntimeError('SCfeedbackStitch: FAIL Reached maxsteps') + raise RuntimeError('Second turn stitching: FAIL Reached maxsteps') -def SCfeedbackBalance(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, eps=1e-4, maxsteps=10, **pinv_params): +def balance(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, eps=1e-4, maxsteps=10, **pinv_params): """ Balance two-turn BPM readings @@ -189,7 +189,7 @@ def SCfeedbackBalance(SC, response_matrix, reference=None, cm_ords=None, bpm_ord Args: SC: SimulatedCommissioning class instance. - Mplus: Pseudo-inverse trajectory-response matrix. + response_matrix: Trajectory-response matrix. reference: (None) target orbit in the format `[x_1 ... x_n y_1 ...y_n]`, where [x_i,y_i]` is the target position at the i-th BPM. cm_ords: List of CM ordinates to be used for correction (SC.ORD.CM) @@ -201,7 +201,7 @@ def SCfeedbackBalance(SC, response_matrix, reference=None, cm_ords=None, bpm_ord SimulatedCommissioning class instance with corrected `SC.RING` """ - LOGGER.debug('SCfeedbackBalance: Start') + LOGGER.debug('Balancing two turns: Start') if SC.INJ.nTurns != 2: raise ValueError("Balancing works only with two turns.") bpm_ords, cm_ords, reference = _check_ords(SC, response_matrix, reference, bpm_ords, cm_ords) @@ -231,32 +231,32 @@ def SCfeedbackBalance(SC, response_matrix, reference=None, cm_ords=None, bpm_ord # Check stopping criteria if transmission_history[-1] < bpm_readings.shape[1]: - raise RuntimeError('SCfeedbackBalance: FAIL (lost transmission)') + raise RuntimeError('Balancing two turns: FAIL (lost transmission)') if _is_stable_or_converged(NREPRO, eps, rms_orbit_history): - LOGGER.debug(f'SCfeedbackBalance: Success (converged after {steps} steps)') + LOGGER.debug(f'Balancing two turns: Success (converged after {steps} steps)') return SC - raise RuntimeError('SCfeedbackBalance: FAIL (maxsteps reached, unstable)') + raise RuntimeError('Balancing two turns: FAIL (maxsteps reached, unstable)') -def SCfeedbackRun(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, eps=1e-4, target=0, maxsteps=30, scaleDisp=0, **pinv_params): +def correct(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=None, eps=1e-4, target=0, maxsteps=30, scaleDisp=0, **pinv_params): """ - iterative orbit correction + Iterative orbit/trajectory correction Iteratively applies orbit corrections using the pseudoinverse of the - trajectory response matrix `Mplus`, until a break-condition specified by one + trajectory `response_matrix`, until a break-condition specified by one of 'eps', 'target' or 'maxsteps' is met. The dispersion can be included, thus the rf frequency as a correction - parameter. If the dispersion is to be included, `Mplus` has to have the size - `(len(SC.ORD.HCM) + len(SC.ORD.VCM) + 1) x len(SC.ORD.BPM)`, otherwise the size - `(len(SC.ORD.HCM) + len(SC.ORD.VCM)) x len(SC.ORD.BPM)`, or correspondingly if the CM + parameter. If the dispersion is to be included, `response_matrix` has to have the size + `(2 * len(SC.ORD.BPM)) x (len(SC.ORD.HCM) + len(SC.ORD.VCM) + 1)`, otherwise the size + `(2 * len(SC.ORD.BPM)) x (len(SC.ORD.HCM) + len(SC.ORD.VCM))`, or correspondingly if the CM and/or BPM ordinates for the correction is explicitly given (see options below). `SC.RING` is assumed to be a lattice with transmission through all considered turns. Raises RuntimeError if transmission is lost. Args: SC: SimulatedCommissioning class instance. - Mplus: Pseudo-inverse trajectory-response matrix. + response_matrix: Trajectory/orbit-response matrix. reference: (None) target orbit in the format `[x_1 ... x_n y_1 ...y_n]`, where [x_i,y_i]` is the target position at the i-th BPM. cm_ords: List of CM ordinates to be used for correction (SC.ORD.CM) @@ -281,12 +281,12 @@ def SCfeedbackRun(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=No trackMode='ORB', Z0=np.zeros(6), nTurns=1, rfStep=1E3, useIdealRing=True) - MinvCO = SCgetPinv(np.column_stack((MCO, 1E8 * eta)), alpha=10) - SC = SCfeedbackRun(SC, MinvCO, target=0, maxsteps=50, scaleDisp=1E8) + SC = correct(SC, np.column_stack((MCO, 1E8 * eta)), alpha=10, target=0, maxsteps=50, scaleDisp=1E8) """ - LOGGER.debug('SCfeedbackRun: Start') - bpm_ords, cm_ords, reference = _check_ords(SC, response_matrix, reference, bpm_ords, cm_ords) + LOGGER.debug('Orbit/trajectory correction: Start') + 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) @@ -307,17 +307,17 @@ def SCfeedbackRun(SC, response_matrix, reference=None, cm_ords=None, bpm_ords=No # Check stopping criteria if np.any(np.isnan(bpm_readings[0, :])): - raise RuntimeError('SCfeedbackRun: FAIL (lost transmission)') + raise RuntimeError('Orbit/trajectory correction: FAIL (lost transmission)') if max(rms_orbit_history[-1]) < target and _is_stable_or_converged(min(NREPRO, maxsteps), eps, rms_orbit_history): - LOGGER.debug(f"SCfeedbackRun: Success (target reached after {steps:d} steps)") + LOGGER.debug(f"Orbit/trajectory correction: Success (target reached after {steps:d} steps)") return SC if _is_stable_or_converged(NREPRO, eps, rms_orbit_history): - LOGGER.debug(f"SCfeedbackRun: Success (converged after {steps:d} steps)") + LOGGER.debug(f"Orbit/trajectory correction: Success (converged after {steps:d} steps)") return SC if _is_stable_or_converged(min(NREPRO, maxsteps), eps, rms_orbit_history) or maxsteps == 1: - LOGGER.debug("SCfeedbackRun: Success (maxsteps reached)") + LOGGER.debug("Orbit/trajectory correction: Success (maxsteps reached)") return SC - raise RuntimeError("SCfeedbackRun: FAIL (maxsteps reached, unstable)") + raise RuntimeError("Orbit/trajectory correction: FAIL (maxsteps reached, unstable)") def _check_ords(SC, response_matrix, reference, bpm_ords, cm_ords): diff --git a/pySC/correction/ramp_errors.py b/pySC/correction/ramp_errors.py index e9d91c2..e1c5f21 100644 --- a/pySC/correction/ramp_errors.py +++ b/pySC/correction/ramp_errors.py @@ -1,8 +1,8 @@ import numpy as np from pySC.core.constants import SUPPORT_TYPES, RF_PROPERTIES -from pySC.correction.orbit_trajectory import SCfeedbackRun +from pySC.correction import orbit_trajectory from pySC.lattice_properties.response_model import SCgetModelRM -from pySC.utils.sc_tools import SCgetPinv, SCscaleCircumference +from pySC.utils.sc_tools import SCscaleCircumference from pySC.utils import logging_tools LOGGER = logging_tools.get_logger(__name__) @@ -15,7 +15,6 @@ def SCrampUpErrors(SC, nStepsRamp=10, eps=1e-5, target=0, alpha=10, maxsteps=30) errFieldsRF = ['Offset', 'CalError'] SC0 = SC M = SCgetModelRM(SC, SC.ORD.BPM, SC.ORD.CM, nTurns=SC.INJ.nTurns, trackMode=SC.INJ.trackMode) - Mplus = SCgetPinv(M, alpha=alpha) for scale in np.linspace(1 / nStepsRamp, 1, nStepsRamp): LOGGER.debug(f'Ramping up errors with scaling factor {scale:.2f}.') SC = scaleSupport(SC, SC0, errFieldsSup, scale) @@ -25,7 +24,7 @@ def SCrampUpErrors(SC, nStepsRamp=10, eps=1e-5, target=0, alpha=10, maxsteps=30) SC = scaleInjection(SC, SC0, scale) SC = scaleCircumference(SC, SC0, scale) try: - SC = SCfeedbackRun(SC, Mplus, target=target, maxsteps=maxsteps, eps=eps) + SC = orbit_trajectory.correct(SC, M, alpha=alpha, target=target, maxsteps=maxsteps, eps=eps) except RuntimeError: if 2 * nStepsRamp > 100: raise Exception(f'Ramping up failed at scaling {scale:.2f} with {nStepsRamp} ramping steps. ' diff --git a/pySC/example.py b/pySC/example.py index 3dc820d..508e509 100644 --- a/pySC/example.py +++ b/pySC/example.py @@ -5,7 +5,7 @@ from at import Lattice from pySC.utils.at_wrapper import atloco from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.correction.orbit_trajectory import SCfeedbackFirstTurn, SCfeedbackStitch, SCfeedbackRun, SCfeedbackBalance +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 @@ -112,20 +112,18 @@ def _marker(name): 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) - Minv2 = SCgetPinv(RM2, alpha=50) SC.INJ.nParticles = 1 SC.INJ.nTurns = 1 SC.INJ.nShots = 1 SC.INJ.trackMode = 'TBT' eps = 5E-4 # Noise level bpm_reading(SC) - SC = SCfeedbackFirstTurn(SC, Minv1) + SC = orbit_trajectory.first_turn(SC, RM1, alpha=50) SC.INJ.nTurns = 2 - SC = SCfeedbackStitch(SC, Minv2, nBPMs=3, maxsteps=20) - # SC = SCfeedbackRun(SC, Minv2, target=300E-6, maxsteps=30, eps=eps) - SC = SCfeedbackBalance(SC, Minv2, maxsteps=32, eps=eps) + SC = orbit_trajectory.stitch(SC, RM2, n_bpms=3, maxsteps=20, alpha=50) + # SC = orbit_trajectory.correct(SC, RM2, target=300E-6, maxsteps=30, eps=eps, alpha=50) + SC = orbit_trajectory.balance(SC, RM2, maxsteps=32, eps=eps, alpha=50) # plot_cm_strengths(SC) # Performing trajectory BBA @@ -141,7 +139,7 @@ def _marker(name): for rel_setting in np.linspace(0.1, 1, 5): SC.set_magnet_setpoints(sextOrds, rel_setting, False, 2, method='rel') try: - SC = SCfeedbackBalance(SC, Minv2, maxsteps=32, eps=eps) + SC = orbit_trajectory.balance(SC, RM2, maxsteps=32, eps=eps, alpha=50) except RuntimeError: pass @@ -171,11 +169,10 @@ def _marker(name): SC.INJ.trackMode = 'ORB' MCO = SCgetModelRM(SC, SC.ORD.BPM, SC.ORD.CM, trackMode='ORB') eta = SCgetModelDispersion(SC, SC.ORD.BPM, SC.ORD.RF) - + resp_with_disp = np.column_stack((MCO, 1E8 * eta)) for alpha in range(10, 0, -1): - MinvCO = SCgetPinv(np.column_stack((MCO, 1E8 * eta)), alpha=alpha) try: - CUR = SCfeedbackRun(SC, MinvCO, target=0, maxsteps=50, scaleDisp=1E8) + CUR = orbit_trajectory.correct(SC, resp_with_disp, target=0, maxsteps=50, scaleDisp=1E8, alpha=alpha) except RuntimeError: break B0rms = np.sqrt(np.mean(np.square(bpm_reading(SC)[0]), axis=1)) @@ -201,12 +198,11 @@ def _marker(name): # {Ords, normal/skew, ind/fam, deltaK} [SCgetOrds(SC.RING, 'QD'), False, 'individual', 1E-4]) - Morbinv = SCgetPinv(MCO, alpha=50) for n in range(6): _, bpm_data, cm_data, fit_parameters, loco_flags, ring_data = atloco(loco_meas_data, bpm_data, cm_data, fit_parameters, loco_flags, ring_data) SC = apply_lattice_correction(SC, fit_parameters) - SC = SCfeedbackRun(SC, Morbinv, target=0, maxsteps=30) + 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, diff --git a/pySC/matlab_index.py b/pySC/matlab_index.py index 96849fe..c7dd4ac 100644 --- a/pySC/matlab_index.py +++ b/pySC/matlab_index.py @@ -21,8 +21,7 @@ from pySC.core.lattice_setting import switch_cavity_and_radiation from pySC.correction.bba import trajectory_bba, orbit_bba, fake_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, \ - SCfeedbackRun as frun, SCfeedbackBalance as fbalance +from pySC.correction import orbit_trajectory from pySC.correction.ramp_errors import SCrampUpErrors as ramp_up_errors from pySC.correction.rf import correct_rf_phase, correct_rf_frequency from pySC.correction.tune import tune_scan @@ -72,28 +71,28 @@ def SCdynamicAperture(RING, dE, /, *, bounds=np.array([0, 1e-3]), nturns=1000, t def SCfeedbackBalance(SC, Mplus, /, *, R0=None, CMords=None, BPMords=None, eps=1e-4, maxsteps=10, verbose=False): - return fbalance(SC, Mplus, reference=R0, cm_ords=CMords, bpm_ords=BPMords, eps=eps, maxsteps=maxsteps) + return orbit_trajectory.balance(SC, np.linalg.pinv(Mplus), reference=R0, cm_ords=CMords, bpm_ords=BPMords, + eps=eps, maxsteps=maxsteps) def SCfeedbackFirstTurn(SC, Mplus, /, *, R0=None, CMords=None, BPMords=None, maxsteps=100, wiggleAfter=20, wiggleSteps=32, wiggleRange=np.array([500E-6, 1000E-6]), verbose=False): - return first_turn(SC, Mplus, reference=R0, cm_ords=CMords, bpm_ords=BPMords, maxsteps=maxsteps, - wiggle_after=wiggleAfter, - wiggle_steps=wiggleSteps, wiggle_range=wiggleRange) + # wiggle parameters are constants of pySC.correction.orbit_trajectory module + return orbit_trajectory.first_turn(SC, np.linalg.pinv(Mplus), reference=R0, cm_ords=CMords, bpm_ords=BPMords, + maxsteps=maxsteps) def SCfeedbackRun(SC, Mplus, /, *, R0=None, CMords=None, BPMords=None, eps=1e-4, target=0, maxsteps=30, scaleDisp=0, weight=None, verbose=False): - return frun(SC, Mplus, reference=R0, cm_ords=CMords, bpm_ords=BPMords, eps=eps, target=target, maxsteps=maxsteps, - scaleDisp=scaleDisp) + return orbit_trajectory.correct(SC, np.linalg.pinv(Mplus), reference=R0, cm_ords=CMords, bpm_ords=BPMords, + eps=eps, target=target, maxsteps=maxsteps, scaleDisp=scaleDisp) def SCfeedbackStitch(SC, Mplus, /, *, R0=None, CMords=None, BPMords=None, nBPMs=4, maxsteps=30, nRepro=3, wiggle_steps=32, wiggle_range=np.array([500E-6, 1000E-6])): - return stitch(SC, Mplus, reference=R0, cm_ords=CMords, bpm_ords=BPMords, nBPMs=nBPMs, maxsteps=maxsteps, - nRepro=nRepro, - wiggle_steps=wiggle_steps, wiggle_range=wiggle_range) + # wiggle parameters are constants of pySC.correction.orbit_trajectory module + return orbit_trajectory.stitch(SC, np.linalg.pinv(Mplus), reference=R0, cm_ords=CMords, bpm_ords=BPMords, n_bpms=nBPMs, maxsteps=maxsteps) def SCfitInjectionZ(SC, mode, /, *, nDims=np.array([0, 1]), nBPMs=np.array([0, 1, 2]), nShots=None, verbose=0, diff --git a/tests/test_example.py b/tests/test_example.py index ef42833..51e6a7c 100644 --- a/tests/test_example.py +++ b/tests/test_example.py @@ -2,12 +2,12 @@ from tests.test_at_wrapper import at_lattice import numpy as np from pySC.core.simulated_commissioning import SimulatedCommissioning -from pySC.correction.orbit_trajectory import SCfeedbackFirstTurn, SCfeedbackStitch, SCfeedbackRun, SCfeedbackBalance +from pySC.correction import orbit_trajectory from pySC.core.beam import bpm_reading, beam_transmission 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, SCgetPinv +from pySC.utils.sc_tools import SCgetOrds from pySC.core.lattice_setting import switch_cavity_and_radiation from pySC.correction.rf import correct_rf_phase, correct_rf_frequency @@ -69,18 +69,16 @@ def test_example(at_lattice): 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) - minv2 = SCgetPinv(rm2, alpha=50) sc.INJ.nParticles = 1 sc.INJ.nTurns = 1 sc.INJ.nShots = 1 sc.INJ.trackMode = 'TBT' eps = 5E-4 # Noise level - sc = SCfeedbackFirstTurn(sc, minv1) + sc = orbit_trajectory.first_turn(sc, rm1, alpha=50) sc.INJ.nTurns = 2 - sc = SCfeedbackStitch(sc, minv2, nBPMs=3, maxsteps=20) - sc = SCfeedbackBalance(sc, minv2, maxsteps=32, eps=eps) + sc = orbit_trajectory.stitch(sc, rm2, alpha=50, n_bpms=3, maxsteps=20) + sc = orbit_trajectory.balance(sc, rm2, alpha=50, maxsteps=32, eps=eps) # Performing BBA sc.INJ.nParticles = 1 @@ -99,7 +97,7 @@ def test_example(at_lattice): # Turning on the sextupoles for rel_setting in np.linspace(0.1, 1, 5): sc.set_magnet_setpoints(sext_ords, rel_setting, False, 2, method='rel') - sc = SCfeedbackBalance(sc, minv2, maxsteps=32, eps=eps) + sc = orbit_trajectory.balance(sc, rm2, alpha=50, maxsteps=32, eps=eps) sc.RING = switch_cavity_and_radiation(sc.RING, 'cavityon') @@ -116,11 +114,10 @@ def test_example(at_lattice): sc.INJ.trackMode = 'ORB' mco = SCgetModelRM(sc, sc.ORD.BPM, sc.ORD.CM, trackMode='ORB') eta = SCgetModelDispersion(sc, sc.ORD.BPM, sc.ORD.RF) - + resp_with_disp = np.column_stack((mco, 1E8 * eta)) for alpha in range(10, 0, -1): - minv_co = SCgetPinv(np.column_stack((mco, 1E8 * eta)), alpha=alpha) try: - cur = SCfeedbackRun(sc, minv_co, target=0, maxsteps=50, scaleDisp=1E8) + cur = orbit_trajectory.correct(sc, resp_with_disp, alpha=50, target=0, maxsteps=50, scaleDisp=1E8) except RuntimeError: break B0rms = np.sqrt(np.mean(np.square(bpm_reading(sc)[0]), axis=1))