Add PicoScope variant override. Return offsets from align functions.

3 files changed, 27 insertions, 20 deletions

diff --git a/pyecsca/sca/scope/picoscope_sdk.py b/pyecsca/sca/scope/picoscope_sdk.py
index c78958d..0962af1 100644
--- a/pyecsca/sca/scope/picoscope_sdk.py
+++ b/pyecsca/sca/scope/picoscope_sdk.py
@@ -62,8 +62,9 @@ class PicoScopeSdk(Scope):  # pragma: no cover
         "rising": 2,
         "falling": 3
     }
+    _variant: Optional[str]
 
-    def __init__(self):
+    def __init__(self, variant: Optional[str] = None):
         super().__init__()
         self.handle: ctypes.c_int16 = ctypes.c_int16()
         self.frequency: Optional[int] = None
@@ -73,6 +74,7 @@ class PicoScopeSdk(Scope):  # pragma: no cover
         self.timebase: Optional[int] = None
         self.buffers: MutableMapping = {}
         self.ranges: MutableMapping = {}
+        self._variant = variant
 
     def open(self) -> None:
         assert_pico_ok(self.__dispatch_call("OpenUnit", ctypes.byref(self.handle)))
@@ -82,11 +84,14 @@ class PicoScopeSdk(Scope):  # pragma: no cover
         return list(self.CHANNELS.keys())
 
     def get_variant(self):
+        if self._variant is not None:
+            return self._variant
         info = (ctypes.c_int8 * 6)()
         size = ctypes.c_int16()
         assert_pico_ok(self.__dispatch_call("GetUnitInfo", self.handle, ctypes.byref(info), 6,
                                             ctypes.byref(size), 3))
-        return "".join(chr(i) for i in info[:size.value])
+        self._variant = "".join(chr(i) for i in info[:size.value])
+        return self._variant
 
     def setup_frequency(self, frequency: int, pretrig: int, posttrig: int) -> Tuple[int, int]:
         return self.set_frequency(frequency, pretrig, posttrig)
@@ -216,8 +221,8 @@ class PicoScopeSdk(Scope):  # pragma: no cover
 if isinstance(ps3000, CannotFindPicoSDKError):
     @public
     class PS3000Scope(PicoScopeSdk):  # pragma: no cover
-        def __init__(self):
-            super().__init__()
+        def __init__(self, variant: Optional[str] = None):
+            super().__init__(variant)
             raise ps3000
 else:
     @public
@@ -270,8 +275,8 @@ else:
 if isinstance(ps4000, CannotFindPicoSDKError):
     @public
     class PS4000Scope(PicoScopeSdk):  # pragma: no cover
-        def __init__(self):
-            super().__init__()
+        def __init__(self, variant: Optional[str] = None):
+            super().__init__(variant)
             raise ps4000
 else:
     @public
@@ -323,8 +328,8 @@ else:
 if isinstance(ps5000, CannotFindPicoSDKError):
     @public
     class PS5000Scope(PicoScopeSdk):  # pragma: no cover
-        def __init__(self):
-            super().__init__()
+        def __init__(self, variant: Optional[str] = None):
+            super().__init__(variant)
             raise ps5000
 else:
     @public
@@ -368,8 +373,8 @@ else:
 if isinstance(ps6000, CannotFindPicoSDKError):
     @public
     class PS6000Scope(PicoScopeSdk):  # pragma: no cover
-        def __init__(self):
-            super().__init__()
+        def __init__(self, variant: Optional[str] = None):
+            super().__init__(variant)
             raise ps6000
 else:
     @public
diff --git a/pyecsca/sca/trace/align.py b/pyecsca/sca/trace/align.py
index cd26aad..02d1762 100644
--- a/pyecsca/sca/trace/align.py
+++ b/pyecsca/sca/trace/align.py
@@ -12,8 +12,9 @@ from .trace import Trace
 
 
 def align_reference(reference: Trace, *traces: Trace,
-                    align_func: Callable[[Trace], Tuple[bool, int]]) -> List[Trace]:
+                    align_func: Callable[[Trace], Tuple[bool, int]]) -> Tuple[List[Trace], List[int]]:
     result = [deepcopy(reference)]
+    offsets = [0]
     for trace in traces:
         length = len(trace.samples)
         include, offset = align_func(trace)
@@ -28,13 +29,14 @@ def align_reference(reference: Trace, *traces: Trace,
             else:
                 result_samples[-offset:] = trace.samples[:length + offset]
         result.append(trace.with_samples(result_samples))
-    return result
+        offsets.append(offset)
+    return result, offsets
 
 
 @public
 def align_correlation(reference: Trace, *traces: Trace,
                       reference_offset: int, reference_length: int,
-                      max_offset: int, min_correlation: float = 0.5) -> List[Trace]:
+                      max_offset: int, min_correlation: float = 0.5) -> Tuple[List[Trace], List[int]]:
     """
     Align `traces` to the reference `trace`. Using the cross-correlation of a part of the reference
     trace starting at `reference_offset` with `reference_length` and try to match it to a part of
@@ -75,7 +77,7 @@ def align_correlation(reference: Trace, *traces: Trace,
 
 @public
 def align_peaks(reference: Trace, *traces: Trace,
-                reference_offset: int, reference_length: int, max_offset: int) -> List[Trace]:
+                reference_offset: int, reference_length: int, max_offset: int) -> Tuple[List[Trace], List[int]]:
     """
     Align `traces` to the reference `trace` so that the maximum value within the reference trace
     window from `reference_offset` of `reference_length` aligns with the maximum
@@ -106,7 +108,7 @@ def align_peaks(reference: Trace, *traces: Trace,
 @public
 def align_offset(reference: Trace, *traces: Trace,
                  reference_offset: int, reference_length: int, max_offset: int,
-                 dist_func: Callable[[np.ndarray, np.ndarray], float], max_dist: float = float("inf")) -> List[Trace]:
+                 dist_func: Callable[[np.ndarray, np.ndarray], float], max_dist: float = float("inf")) -> Tuple[List[Trace], List[int]]:
     """
     Align `traces` to the reference `trace` so that the value of the `dist_func` is minimized
     between the reference trace window from `reference_offset` of `reference_length` and the trace
@@ -145,7 +147,7 @@ def align_offset(reference: Trace, *traces: Trace,
 
 @public
 def align_sad(reference: Trace, *traces: Trace,
-              reference_offset: int, reference_length: int, max_offset: int) -> List[Trace]:
+              reference_offset: int, reference_length: int, max_offset: int) -> Tuple[List[Trace], List[int]]:
     """
     Align `traces` to the reference `trace` so that the Sum Of Absolute Differences between the
     reference trace window from `reference_offset` of `reference_length` and the trace being aligned
diff --git a/test/sca/test_align.py b/test/sca/test_align.py
index 15d452a..006d6b9 100644
--- a/test/sca/test_align.py
+++ b/test/sca/test_align.py
@@ -13,7 +13,7 @@ class AlignTests(Plottable):
         a = Trace(first_arr)
         b = Trace(second_arr)
         c = Trace(third_arr)
-        result = align_correlation(a, b, c, reference_offset=1, reference_length=3, max_offset=4, min_correlation=0.65)
+        result, offsets = align_correlation(a, b, c, reference_offset=1, reference_length=3, max_offset=4, min_correlation=0.65)
         self.assertIsNotNone(result)
         self.assertEqual(len(result), 2)
         np.testing.assert_equal(result[0].samples, first_arr)
@@ -22,7 +22,7 @@ class AlignTests(Plottable):
     @slow
     def test_large_align(self):
         example = InspectorTraceSet.read("test/data/example.trs")
-        result = align_correlation(*example, reference_offset=100000, reference_length=20000, max_offset=15000)
+        result, offsets  = align_correlation(*example, reference_offset=100000, reference_length=20000, max_offset=15000)
         self.assertIsNotNone(result)
 
     @slow
@@ -36,7 +36,7 @@ class AlignTests(Plottable):
         second_arr = np.array([10, 10, 10, 10, 90, 40, 50, 20, 10, 17, 16, 10], dtype=np.dtype("i1"))
         a = Trace(first_arr)
         b = Trace(second_arr)
-        result = align_peaks(a, b, reference_offset=2, reference_length=5, max_offset=3)
+        result, offsets  = align_peaks(a, b, reference_offset=2, reference_length=5, max_offset=3)
         self.assertEqual(np.argmax(result[0].samples), np.argmax(result[1].samples))
 
     def test_sad_align(self):
@@ -44,7 +44,7 @@ class AlignTests(Plottable):
         second_arr = np.array([10, 10, 90, 40, 50, 20, 10, 17, 16, 10, 10], dtype=np.dtype("i1"))
         a = Trace(first_arr)
         b = Trace(second_arr)
-        result = align_sad(a, b, reference_offset=2, reference_length=5, max_offset=3)
+        result, offsets  = align_sad(a, b, reference_offset=2, reference_length=5, max_offset=3)
         self.assertEqual(len(result), 2)
 
     def test_dtw_align_scale(self):