Tune scan

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

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