From 85688142c87228b8fb62e2a72a6963b6acc5a36c Mon Sep 17 00:00:00 2001
From: Rastislav Turanyi <rastislav.turanyi@stfc.ac.uk>
Date: Fri, 22 Nov 2024 14:01:22 +0000
Subject: [PATCH] Use only first and last chopper in non-Fermi chopper width

---
 src/resolution_functions/models/pychop.py | 18 ++++++++++--------
 tests/unit_tests/test_pychop.py           |  2 +-
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/resolution_functions/models/pychop.py b/src/resolution_functions/models/pychop.py
index 892bf5f..debac64 100644
--- a/src/resolution_functions/models/pychop.py
+++ b/src/resolution_functions/models/pychop.py
@@ -535,7 +535,7 @@ def get_long_frequency(frequency: list[int],
     def _get_chop_times(cls,
                         model_data: PyChopModelDataNonFermi,
                         e_init: float,
-                        chopper_frequency: list[int]) -> list[list[np.ndarray]]:
+                        chopper_frequency: list[int]) -> list[list[Float[np.ndarray, 'times']]]:
         frequencies = cls.get_long_frequency(chopper_frequency, model_data)
         choppers = model_data.choppers
 
@@ -549,8 +549,8 @@ def _get_chop_times(cls,
         # if there's only one disk we prepend a dummy disk with full opening at zero distance
         # so that the distance calculations (which needs the difference between disk pos) works
         if len(choppers) == 1:
-            choppers = {
-                'prepended': {
+            choppers = [
+                {
                     'distance': 0,
                     'nslot': 1,
                     'slot_ang_pos': None,
@@ -559,15 +559,17 @@ def _get_chop_times(cls,
                     'radius': 500,
                     'num_disk': 1
                 },
-                **choppers
-            }
+                *list(choppers.values())
+            ]
             frequencies = np.array([model_data.source_frequency, frequencies[0]])
+        else:
+            choppers = list(choppers.values())
 
         chop_times = []
 
         # first we optimise on the main Ei
         # TODO: Calculate only the first and last choppers (only those used upstream)
-        for i, (frequency, chopper) in enumerate(zip(frequencies, choppers.values())):
+        for frequency, chopper in zip([frequencies[0], frequencies[-1]], [choppers[0], choppers[-1]]):
             chopper: DiskChopper
             this_phase, phase_independence = chopper['default_phase'], chopper['is_phase_independent']
 
@@ -605,7 +607,7 @@ def _get_chop_times(cls,
                 next_win_t = uSec / model_data.source_frequency + (uSec / frequency)
 
                 while realTimeOp[0] < next_win_t:
-                    chop_times[i].append(deepcopy(realTimeOp))
+                    chop_times[-1].append(deepcopy(realTimeOp))
                     slt0 = islt % slot
                     slt1 = (islt + 1) % slot
                     angdiff = angles[slt1] - angles[slt0]
@@ -619,7 +621,7 @@ def _get_chop_times(cls,
                 realTimeOp -= next_win_t * np.ceil(realTimeOp[0] / next_win_t)
 
                 while realTimeOp[0] < (uSec / p_frames + next_win_t):
-                    chop_times[i].append(deepcopy(realTimeOp))
+                    chop_times[-1].append(deepcopy(realTimeOp))
                     realTimeOp += next_win_t
 
         return chop_times
diff --git a/tests/unit_tests/test_pychop.py b/tests/unit_tests/test_pychop.py
index bcd0056..9f638a1 100644
--- a/tests/unit_tests/test_pychop.py
+++ b/tests/unit_tests/test_pychop.py
@@ -342,7 +342,7 @@ def test_chop_times(matrix, pychop_nonfermi_data):
 
     actual = PyChopModelNonFermi._get_chop_times(data, e_init, chopper_frequencies)
 
-    for aa, ee in zip(actual, expected):
+    for aa, ee in zip(actual, [expected[0], expected[-1]]):
         for a, e in zip(aa, ee):
             try:
                 assert_allclose(a, e, rtol=0, atol=1e-8)