Add dtype handling to scope interface.

4 files changed, 14 insertions, 12 deletions

diff --git a/pyecsca/sca/scope/base.py b/pyecsca/sca/scope/base.py
index 5b2cde2..2663cb2 100644
--- a/pyecsca/sca/scope/base.py
+++ b/pyecsca/sca/scope/base.py
@@ -90,12 +90,13 @@ class Scope(object):
         """
         raise NotImplementedError
 
-    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
+    def retrieve(self, channel: str, type: SampleType, dtype = None) -> Optional[Trace]:
         """
         Retrieve a captured trace of a channel.
 
         :param channel: The channel to retrieve the trace from.
         :param type: The type of returned samples.
+        :param dtype: The data type of the returned samples, should be numpy dtype-like.
         :return: The captured trace (if any).
         """
         raise NotImplementedError
diff --git a/pyecsca/sca/scope/chipwhisperer.py b/pyecsca/sca/scope/chipwhisperer.py
index cb12c91..49fab5d 100644
--- a/pyecsca/sca/scope/chipwhisperer.py
+++ b/pyecsca/sca/scope/chipwhisperer.py
@@ -52,11 +52,11 @@ class ChipWhispererScope(Scope):  # pragma: no cover
     def capture(self, timeout: Optional[int] = None) -> bool:
         return not self.scope.capture()
 
-    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
+    def retrieve(self, channel: str, type: SampleType, dtype = np.float16) -> Optional[Trace]:
         data = self.scope.get_last_trace()
         if data is None:
             return None
-        return Trace(data, {"sampling_frequency": self.scope.clock.clkgen_freq, "channel": channel})
+        return Trace(np.array(data, dtype=dtype), {"sampling_frequency": self.scope.clock.clkgen_freq, "channel": channel})
 
     def stop(self) -> None:
         pass
diff --git a/pyecsca/sca/scope/picoscope_alt.py b/pyecsca/sca/scope/picoscope_alt.py
index 5a1b74b..10f7959 100644
--- a/pyecsca/sca/scope/picoscope_alt.py
+++ b/pyecsca/sca/scope/picoscope_alt.py
@@ -1,4 +1,5 @@
 from time import time_ns, sleep
+import numpy as np
 from typing import Optional, Tuple, Sequence, Union
 
 from picoscope.ps3000 import PS3000
@@ -57,11 +58,11 @@ class PicoScopeAlt(Scope):  # pragma: no cover
                 return False
         return True
 
-    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
+    def retrieve(self, channel: str, type: SampleType, dtype = np.float32) -> Optional[Trace]:
         if type == SampleType.Raw:
-            data = self.ps.getDataRaw(channel)
+            data = self.ps.getDataRaw(channel).astype(dtype=dtype, copy=False)
         else:
-            data = self.ps.getDataV(channel)
+            data = self.ps.getDataV(channel, dtype=dtype)
         if data is None:
             return None
         return Trace(data, {"sampling_frequency": self.frequency, "channel": channel, "sample_type": type})
diff --git a/pyecsca/sca/scope/picoscope_sdk.py b/pyecsca/sca/scope/picoscope_sdk.py
index fe23e5e..46c1da2 100644
--- a/pyecsca/sca/scope/picoscope_sdk.py
+++ b/pyecsca/sca/scope/picoscope_sdk.py
@@ -30,20 +30,20 @@ from ..trace import Trace
 
 
 def adc2volt(adc: Union[np.ndarray, ctypes.c_int16],
-             volt_range: float, adc_minmax: int) -> Union[np.ndarray, float]:  # pragma: no cover
+             volt_range: float, adc_minmax: int, dtype = np.float32) -> Union[np.ndarray, float]:  # pragma: no cover
     if isinstance(adc, ctypes.c_int16):
         return (adc.value / adc_minmax) * volt_range
     if isinstance(adc, np.ndarray):
-        return ((adc / adc_minmax) * volt_range).astype(np.float16, copy=False)
+        return ((adc / adc_minmax) * volt_range).astype(dtype=dtype, copy=False)
     raise ValueError
 
 
 def volt2adc(volt: Union[np.ndarray, float],
-             volt_range: float, adc_minmax: int) -> Union[np.ndarray, ctypes.c_int16]:  # pragma: no cover
+             volt_range: float, adc_minmax: int, dtype = np.float32) -> Union[np.ndarray, ctypes.c_int16]:  # pragma: no cover
     if isinstance(volt, float):
         return ctypes.c_int16(int((volt / volt_range) * adc_minmax))
     if isinstance(volt, np.ndarray):
-        return ((volt / volt_range) * adc_minmax).astype(np.int16, copy=False)
+        return ((volt / volt_range) * adc_minmax).astype(dtype=dtype, copy=False)
     raise ValueError
 
 
@@ -192,7 +192,7 @@ class PicoScopeSdk(Scope):  # pragma: no cover
                 return False
         return True
 
-    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
+    def retrieve(self, channel: str, type: SampleType, dtype = np.float32) -> Optional[Trace]:
         if self.samples is None:
             raise ValueError
         actual_samples = ctypes.c_int32(self.samples)
@@ -200,7 +200,7 @@ class PicoScopeSdk(Scope):  # pragma: no cover
         assert_pico_ok(
                 self.__dispatch_call("GetValues", self.handle, 0, ctypes.byref(actual_samples), 1,
                                      0, 0, ctypes.byref(overflow)))
-        arr = np.array(self.buffers[channel], dtype=np.int16)
+        arr = np.array(self.buffers[channel], dtype=dtype)
         if type == SampleType.Raw:
             data = arr
         else: