From 962b6ba25c92b4ba579d5189d952425d2632d522 Mon Sep 17 00:00:00 2001
From: malina <lukas.malina@pm.me>
Date: Wed, 24 Jan 2024 18:56:51 +0100
Subject: [PATCH] Little fixes

---
 pySC/correction/tune.py   |  9 +++++----
 tests/test_tune.py        | 24 ++++++++++++++++++++++++
 tests/test_tune_chroma.py | 19 -------------------
 3 files changed, 29 insertions(+), 23 deletions(-)
 create mode 100644 tests/test_tune.py
 delete mode 100644 tests/test_tune_chroma.py

diff --git a/pySC/correction/tune.py b/pySC/correction/tune.py
index fe937d4..e2443f7 100644
--- a/pySC/correction/tune.py
+++ b/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,9 +149,9 @@ 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
@@ -159,6 +159,7 @@ def fit_tune(SC, q_ords, target_tune=None, init_step_size=np.array([0.0001, 0.00
 
 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))
\ No newline at end of file
+    return np.sqrt(np.mean((nu - target) ** 2))
diff --git a/tests/test_tune.py b/tests/test_tune.py
new file mode 100644
index 0000000..a936e2c
--- /dev/null
+++ b/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
diff --git a/tests/test_tune_chroma.py b/tests/test_tune_chroma.py
deleted file mode 100644
index 1629693..0000000
--- a/tests/test_tune_chroma.py
+++ /dev/null
@@ -1,19 +0,0 @@
-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.sc_tools import SCgetOrds
-
-
-def test_fit_tune(at_lattice):
-    SC = sc(at_lattice)
-    SC = fit_tune(SC, q_ords=[SCgetOrds(SC.RING, 'QF'), SCgetOrds(SC.RING, 'QD')], target_tune=SC.RING.get_tune(get_integer=True)[:2]+[0.005,-0.005], fit_integer=True)
-    assert_allclose(actual=SC.RING.get_tune()[:2], desired=SC.RING.get_tune()[:2]+[0.005,-0.005], rtol=1e-2)
-    return sc
-
-def sc(at_lattice):
-    SC = SimulatedCommissioning(at_lattice)
-    SC.register_magnets(SCgetOrds(SC.RING, 'QF'))
-    SC.register_magnets(SCgetOrds(SC.RING, 'QD'))
-    return SC
\ No newline at end of file