feat: Added host Pearson coefficient runner function

1 files changed, 38 insertions, 9 deletions

diff --git a/pyecsca/sca/stacked_traces/correlate.py b/pyecsca/sca/stacked_traces/correlate.py
index 10705ca..862e664 100644
--- a/pyecsca/sca/stacked_traces/correlate.py
+++ b/pyecsca/sca/stacked_traces/correlate.py
@@ -1,13 +1,46 @@
 import numpy as np
 import numpy.typing as npt
 from numba import cuda
+from numba.cuda.cudadrv.devicearray import DeviceNDArray
 from math import sqrt
+from typing import List, Optional, Union
+from .combine import GPUTraceManager
+from .stacked_traces import StackedTraces
+from ..trace.trace import CombinedTrace
+
+
+def gpu_pearson_corr(intermediate_values: npt.NDArray[np.number],
+                     stacked_traces: Optional[StackedTraces] = None,
+                     trace_manager: Optional[GPUTraceManager] = None,
+                     **tm_kwargs) -> Union[CombinedTrace, List[CombinedTrace]]:
+    if (stacked_traces is None) == (trace_manager is None):
+        raise ValueError("Either samples or trace manager must be given.")
+
+    if trace_manager is None:
+        assert stacked_traces is not None
+        trace_manager = GPUTraceManager(stacked_traces, **tm_kwargs)
+
+    if (len(intermediate_values.shape) != 1
+        or (intermediate_values.shape[0]
+            != trace_manager.get_traces_shape()[0])):
+        raise ValueError("Intermediate values have to be a vector "
+                         "as long as trace_count")
+
+    intermed_sum: np.number = np.sum(intermediate_values)
+    intermed_sq_sum: np.number = np.sum(np.square(intermediate_values))
+
+    return trace_manager.run(
+        _gpu_pearson_corr,
+        [intermediate_values, [intermed_sum], [intermed_sq_sum]]
+    )
 
 
 @cuda.jit(device=True, cache=True)
-def gpu_pearson_corr(samples: npt.NDArray[np.number],
-                     intermediate_values: npt.NDArray[np.number],
-                     result: cuda.devicearray.DeviceNDArray):
+def _gpu_pearson_corr(samples: DeviceNDArray,
+                      intermediate_values: DeviceNDArray,
+                      intermed_sum: DeviceNDArray,
+                      intermed_sq_sum: DeviceNDArray,
+                      result: DeviceNDArray):
     """
     Calculates the Pearson correlation coefficient between the given samples and intermediate values using GPU acceleration.
 
@@ -26,19 +59,15 @@ def gpu_pearson_corr(samples: npt.NDArray[np.number],
     n = samples.shape[0]
     samples_sum = 0.
     samples_sq_sum = 0.
-    intermed_sum = 0.
-    intermed_sq_sum = 0.
     product_sum = 0.
 
     for row in range(n):
         samples_sum += samples[row, col]
         samples_sq_sum += samples[row, col] ** 2
-        intermed_sum += intermediate_values[row]
-        intermed_sq_sum += intermediate_values[row] ** 2
         product_sum += samples[row, col] * intermediate_values[row]
 
     numerator = n * product_sum - samples_sum * intermed_sum
-    denominator = (sqrt(n * samples_sq_sum - samples_sum * samples_sum)
-                   * sqrt(n * intermed_sq_sum - intermed_sum * intermed_sum))
+    denominator = (sqrt(n * samples_sq_sum - samples_sum ** 2)
+                   * sqrt(n * intermed_sq_sum[0] - intermed_sum[0] ** 2))
 
     result[col] = numerator / denominator