diff --git a/pySC/correction/tune.py b/pySC/correction/tune.py index cb38b9f..e2443f7 100644 --- a/pySC/correction/tune.py +++ b/pySC/correction/tune.py @@ -11,7 +11,7 @@ from pySC.core.beam import beam_transmission, plot_transmission from pySC.utils import logging_tools -from pySC.utils.at_wrapper import atlinopt + LOGGER = logging_tools.get_logger(__name__) @@ -129,32 +129,37 @@ def plot_scan(fin_trans, max_turns, first_quads, rel_quad_changes): return f, [ax1, ax2, ax3] -def fit_tune(SC, q_ords, target_tune=None, xtol=1E-4, ftol=1E-3, fit_integer=True): - # TODO check if experimantally feasible +def fit_tune(SC, q_ords, target_tune=None, init_step_size=np.array([0.0001, 0.0001]),xtol=1E-3, ftol=1E-3, fit_integer=True): + """ + Applies a tune correction using two quadrupole families. + Note: this is not beam based but assumes the tunes can be measured reasonably well. + + Args: + SC: SimulatedCommissioning instance + q_ords: [2xN] array or list [[1 x NQ1],[1 x NQ2], [1 x NQ3], ...] of quadrupole ordinates + target_tune (optional, [1x2] array): Target tunes for correction. Default: tunes of 'SC.IDEALRING' + init_step_size ([1x2] array, optional): Initial step size for the solver. Default: [0.0001,0.0001] + xtol(float, optional): Step tolerance for solver. Default: 1e-3 + ftol(float, optional): Merit tolerance for solver. Default: 1e-3 + fit_integer(bool, optional): Flag specifying if the integer part should be fitted as well. Default: True. + + Returns: + SC: SimulatedCommissioning instance with corrected tunes. + Example: + SC = fit_tune(SC, q_ords=[SCgetOrds(sc.RING, 'QF'), SCgetOrds(sc.RING, 'QD')], target_tune=numpy.array([0.16,0.21])) + """ if target_tune is None: - target_tune = tune(SC, fit_integer, ideal=True) - LOGGER.debug(f'Fitting tunes from [{tune(SC, fit_integer)}] to [{target_tune}].') - SP0 = np.zeros((len(q_ords), len(q_ords[0]))) # TODO can the lengts vary - for nFam in range(len(q_ords)): - for n in range(len(q_ords[nFam])): - SP0[nFam][n] = SC.RING[q_ords[nFam][n]].SetPointB[1] - fun = lambda x: _fit_tune_fun(SC, q_ords, x, SP0, target_tune, fit_integer) - sol = fmin(fun, xtol=xtol, ftol=ftol) - SC.set_magnet_setpoints(q_ords, sol + SP0, False, 1, method='abs', dipole_compensation=True) - LOGGER.debug(f' Final tune: [{tune(SC, fit_integer)}]\n Setpoints change: [{sol}]') + target_tune = SC.IDEALRING.get_tune(get_integer=fit_integer)[:2] + LOGGER.debug(f'Fitting tunes from [{SC.RING.get_tune(get_integer=fit_integer)}] to [{target_tune}].') + fun = lambda x: _fit_tune_fun(SC, q_ords, x, target_tune, fit_integer) + sol = fmin(fun, init_step_size, xtol=xtol, ftol=ftol) + LOGGER.debug(f' Final tune: [{SC.RING.get_tune(get_integer=fit_integer)}]\n Setpoints change: [{sol}]') return SC -def _fit_tune_fun(SC, q_ords, setpoints, init_setpoints, target, fit_integer): - SC.set_magnet_setpoints(q_ords, setpoints + init_setpoints, False, 1, method='abs', dipole_compensation=True) - nu = tune(SC, fit_integer) +def _fit_tune_fun(SC, q_ords, setpoints, target, fit_integer): + for nFam in range(len(q_ords)): + # TODO check the add method here + SC.set_magnet_setpoints(q_ords[nFam], setpoints[nFam], False, 1, method='add', dipole_compensation=True) + nu = SC.RING.get_tune(get_integer=fit_integer)[:2] return np.sqrt(np.mean((nu - target) ** 2)) - - -def tune(SC, fit_integer: bool = False, ideal: bool = False): - ring = SC.IDEALRING if ideal else SC.RING - if fit_integer: - ld, _, _ = atlinopt(ring, 0, range(len(ring) + 1)) - return ld[-1].mu / 2 / np.pi - _, nu, _ = atlinopt(ring, 0) - return nu diff --git a/tests/test_tune.py b/tests/test_tune.py new file mode 100644 index 0000000..a936e2c --- /dev/null +++ b/tests/test_tune.py @@ -0,0 +1,24 @@ +import pytest +import numpy as np +from numpy.testing import assert_allclose +from tests.test_at_wrapper import at_lattice +from pySC.core.simulated_commissioning import SimulatedCommissioning +from pySC.correction.tune import fit_tune +from pySC.utils import sc_tools + + +def test_fit_tune(sc): + sc = fit_tune(sc, q_ords=[sc_tools.ords_from_regex(sc.RING, 'QF'), + sc_tools.ords_from_regex(sc.RING, 'QD')], + target_tune=sc.RING.get_tune(get_integer=True)[:2] + [0.005, -0.005], fit_integer=True) + # TODO with this tolerance it is not testing much + assert_allclose(actual=sc.RING.get_tune()[:2], desired=sc.RING.get_tune()[:2] + [0.005, -0.005], rtol=1e-2) + return sc + + +@pytest.fixture +def sc(at_lattice): + SC = SimulatedCommissioning(at_lattice) + SC.register_magnets(sc_tools.ords_from_regex(SC.RING, 'QF')) + SC.register_magnets(sc_tools.ords_from_regex(SC.RING, 'QD')) + return SC