From c382cbcbbb18823bdf0a9bb7efcea87eab1ae1fc Mon Sep 17 00:00:00 2001 From: Lina Hoummi Date: Wed, 24 Jan 2024 18:46:19 +0100 Subject: [PATCH] Fixed Chromaticity fit (#43) --------- Co-authored-by: Lina Hoummi Co-authored-by: lmalina --- doc/modules/correction.rst | 3 +++ pySC/correction/chroma.py | 51 +++++++++++++++++++++++++++----------- tests/test_chroma.py | 36 +++++++++++++++++++++++++++ 3 files changed, 75 insertions(+), 15 deletions(-) create mode 100644 tests/test_chroma.py diff --git a/doc/modules/correction.rst b/doc/modules/correction.rst index f7e29b9..b46a868 100644 --- a/doc/modules/correction.rst +++ b/doc/modules/correction.rst @@ -12,3 +12,6 @@ Correction algorithms .. automodule:: pySC.correction.tune :members: +.. automodule:: pySC.correction.chroma + :members: + diff --git a/pySC/correction/chroma.py b/pySC/correction/chroma.py index 15f3c0f..27b08ee 100644 --- a/pySC/correction/chroma.py +++ b/pySC/correction/chroma.py @@ -1,36 +1,57 @@ +""" +Chromaticity +------------- + +This module contains functions to fit the chromaticity of 'SC.RING'. +""" + import numpy as np from scipy.optimize import fmin -from pySC.utils.at_wrapper import atlinopt from pySC.utils import logging_tools LOGGER = logging_tools.get_logger(__name__) -def fit_chroma(SC, s_ords, target_chroma=None, init_step_size=np.array([2, 2]), xtol=1E-4, ftol=1E-3, - tune_knobs_ords=None, tune_knobs_delta_k=None): +def fit_chroma(SC, s_ords, target_chroma=None, init_step_size=np.array([2, 2]), xtol=1E-4, ftol=1E-3): + """ + Applies a chromaticity correction using two sextupole families. + + Args: + SC: SimulatedCommissioning instance + s_ords: [2xN] array or list [[1 x NS1],[1 x NS2], [1 x NS3], ...] of sextupole ordinates + target_chroma ([1x2] array, optional): Target chromaticity for correction. Default: chromaticity of 'SC.IDEALRING' + init_step_size ([1x2] array, optional): Initial step size for the solver. Default: [2,2] + xtol(float, optional): Step tolerance for solver. Default: 1e-4 + ftol(float, optional): Merit tolerance for solver. Default: 1e-3 + + Returns: + SC: SimulatedCommissioning instance with corrected chromaticity. + Example: + SC = fit_chroma(SC, s_ords=[SCgetOrds(sc.RING, 'SF'), SCgetOrds(sc.RING, 'SD')], target_chroma=numpy.array([1,1])) + """ if target_chroma is None: - _, _, target_chroma = atlinopt(SC.IDEALRING, 0, []) + target_chroma = SC.IDEALRING.get_chrom()[0:2] if np.sum(np.isnan(target_chroma)): LOGGER.error('Target chromaticity must not contain NaN. Aborting.') return SC - if tune_knobs_ords is not None and tune_knobs_delta_k is not None: - for nFam in range(len(tune_knobs_ords)): - SC.set_magnet_setpoints(tune_knobs_ords[nFam], tune_knobs_delta_k[nFam], False, 1, - method='add') # TODO quads here? - LOGGER.debug(f'Fitting chromaticities from {atlinopt(SC.RING, 0, [])[2]} to {target_chroma}.') # first two elements - SP0 = np.zeros((len(s_ords), len(s_ords[0]))) # TODO can the lengts vary + + LOGGER.debug(f'Fitting chromaticities from {SC.RING.get_chrom()} to {target_chroma}.') # first two elements + SP0 = [] + for n in range(len(s_ords)): + SP0.append(np.zeros_like(s_ords[n])) for nFam in range(len(s_ords)): for n in range(len(s_ords[nFam])): SP0[nFam][n] = SC.RING[s_ords[nFam][n]].SetPointB[2] fun = lambda x: _fit_chroma_fun(SC, s_ords, x, SP0, target_chroma) sol = fmin(fun, init_step_size, xtol=xtol, ftol=ftol) - SC.set_magnet_setpoints(s_ords, sol + SP0, False, 1, method='abs', dipole_compensation=True) - LOGGER.debug(f' Final chromaticity: {atlinopt(SC.RING, 0, [])[2]}\n Setpoints change: {sol}.') # first two elements + # TODO needs to set the solution to SC + LOGGER.debug(f' Final chromaticity: {SC.RING.get_chrom()}\n Setpoints change: {sol}.') # first two elements return SC -def _fit_chroma_fun(SC, q_ords, setpoints, init_setpoints, target): - SC.set_magnet_setpoints(q_ords, setpoints + init_setpoints, False, 2, method='abs', dipole_compensation=True) - _, _, nu = atlinopt(SC.RING, 0, []) +def _fit_chroma_fun(SC, s_ords, setpoints, init_setpoints, target): + for n in range(len(s_ords)): + SC.set_magnet_setpoints(s_ords[n], setpoints[n] + init_setpoints[n], False, 2, method='abs', dipole_compensation=True) + nu = SC.RING.get_chrom()[0:2] return np.sqrt(np.mean((nu - target) ** 2)) diff --git a/tests/test_chroma.py b/tests/test_chroma.py new file mode 100644 index 0000000..d8aac8f --- /dev/null +++ b/tests/test_chroma.py @@ -0,0 +1,36 @@ +import at +import numpy as np +from numpy.testing import assert_allclose +import pytest +from pathlib import Path +from pySC.core.simulated_commissioning import SimulatedCommissioning +from pySC.utils import logging_tools, sc_tools +from pySC.correction.chroma import fit_chroma + +LOGGER = logging_tools.get_logger(__name__) +INPUTS = Path(__file__).parent / "inputs" + + +def test_chroma_hmba(at_ring): + np.random.seed(12345678) + sc = SimulatedCommissioning(at_ring) + sc.register_bpms(sc_tools.ords_from_regex(sc.RING, 'BPM'), Roll=0.0) + sc.register_magnets(sc_tools.ords_from_regex(sc.RING, 'SF|SD'), CalErrorB=np.array([0, 0, 0.05])) # [1/m] + sc.register_cavities(sc_tools.ords_from_regex(sc.RING, 'RFC')) + sc.apply_errors() + s_ords = [sc_tools.ords_from_regex(sc.RING, '^SF'), sc_tools.ords_from_regex(sc.RING, '^SD')] + target = np.array([2.0,2.0]) + sc = fit_chroma(sc, s_ords, target_chroma=target) + assert_allclose(sc.RING.get_chrom()[0:2], target, atol=2e-6, rtol=1e-6) + + +@pytest.fixture +def at_ring(): + ring = at.load_mat(f'{INPUTS}/hmba.mat') + ring.enable_6d() + at.set_cavity_phase(ring) + at.set_rf_frequency(ring) + + ring.tapering(niter=3, quadrupole=True, sextupole=True) + ring = at.Lattice(ring) + return ring