Tune scan

pySC/correction/tune.py

@@ -129,32 +129,43 @@ 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.001, 0.001]),xtol=1E-4, 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: [1,1]
+            xtol(float, optional): Step tolerance for solver. Default: 1e-4
+            ftol(float, optional): Merit tolerance for solver. Default: 1e-4
+            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
+        target_tune = SC.IDEALRING.get_tune(get_integer=fit_integer)
+    LOGGER.debug(f'Fitting tunes from [{SC.RING.get_tune(get_integer=fit_integer)}] to [{target_tune}].')
+    SP0 = []
+    for n in range(len(q_ords)):
+        SP0.append(np.zeros_like(q_ords[n])) #working with a list of two arrays
     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}]')
+    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)
-    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
+    for nFam in range(len(q_ords)):
+        SC.set_magnet_setpoints(q_ords[nFam], setpoints[nFam] + init_setpoints[nFam], False, 1, method='abs', dipole_compensation=True)
+    nu = SC.RING.get_tune(get_integer=fit_integer)
+    nu = nu[0:2]
+    return np.sqrt(np.mean((nu - target) ** 2))

tests/test_tune_chroma.py

@@ -0,0 +1,29 @@
+import copy
+import at
+import pytest
+import numpy as np
+from numpy.testing import assert_equal, assert_allclose
+from tests.test_at_wrapper import at_lattice
+from pySC.correction.tune import tune_scan, fit_tune
+from pySC.correction.chroma import fit_chroma
+from pySC.core.simulated_commissioning import SimulatedCommissioning
+from pySC.utils.sc_tools import SCgetOrds
+
+
+def test_fit_tune(at_lattice):
+    sc = SimulatedCommissioning(at_lattice)
+    sc.register_cavities(SCgetOrds(sc.RING, 'RFCav'),
+                         FrequencyOffset=5E3,
+                         VoltageOffset=5E3,
+                         TimeLagOffset=0.5)
+    sc.register_magnets(SCgetOrds(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,
+                        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 = fit_tune(sc, q_ords=[SCgetOrds(sc.RING, 'QF'), SCgetOrds(sc.RING, 'QD')], target_tune=np.array([0.16, 0.21]))
+    return sc