Merge branch 'master' into tune_scan

doc/modules/correction.rst

@@ -12,3 +12,6 @@ Correction algorithms
 .. automodule:: pySC.correction.tune
     :members:
 
+.. automodule:: pySC.correction.chroma
+    :members:
+

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))

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