Little fixes

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__)
 
@@ -149,16 +149,17 @@ def fit_tune(SC, q_ords, target_tune=None, init_step_size=np.array([0.0001, 0.00
             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 = SC.IDEALRING.get_tune(get_integer=fit_integer)
+        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)
+    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, 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))
+    return np.sqrt(np.mean((nu - target) ** 2))

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