3 files changed, 41 insertions, 6 deletions

diff --git a/pyecsca/sca/__init__.py b/pyecsca/sca/__init__.py
index 7e44c62..159d402 100644
--- a/pyecsca/sca/__init__.py
+++ b/pyecsca/sca/__init__.py
@@ -4,7 +4,7 @@ from .edit import *
 from .filter import *
 from .process import *
 from .sampling import *
-from .ttest import *
+from .test import *
 from .trace import *
 from .trace_set.base import *
 from .trace_set.inspector import *
diff --git a/pyecsca/sca/ttest.py b/pyecsca/sca/test.py
index 329f37d..e192048 100644
--- a/pyecsca/sca/ttest.py
+++ b/pyecsca/sca/test.py
@@ -1,7 +1,8 @@
+from typing import Sequence, Optional
+
 import numpy as np
 from public import public
-from scipy.stats import ttest_ind
-from typing import Sequence, Optional
+from scipy.stats import ttest_ind, ks_2samp
 
 from .trace import Trace, CombinedTrace
 
@@ -30,7 +31,8 @@ def welch_ttest(first_set: Sequence[Trace], second_set: Sequence[Trace]) -> Opti
 
 
 @public
-def student_ttest(first_set: Sequence[Trace], second_set: Sequence[Trace]) -> Optional[CombinedTrace]:
+def student_ttest(first_set: Sequence[Trace], second_set: Sequence[Trace]) -> Optional[
+    CombinedTrace]:
     """
     Perform the Students's t-test sample wise on two sets of traces `first_set` and `second_set`.
     Useful for Test Vector Leakage Analysis (TVLA).
@@ -40,3 +42,22 @@ def student_ttest(first_set: Sequence[Trace], second_set: Sequence[Trace]) -> Op
     :return: Student's t-values (samplewise)
     """
     return ttest_func(first_set, second_set, True)
+
+
+@public
+def ks_test(first_set: Sequence[Trace], second_set: Sequence[Trace]) -> Optional[CombinedTrace]:
+    """
+    Perform the Kolmogorov-Smirnov two sample test on equality of distributions sample wise on
+    two sets of traces `first_set` and `second_set`.
+    :param first_set:
+    :param second_set:
+    :return: Kolmogorov-Smirnov test statistic values (samplewise)
+    """
+    if not first_set or not second_set or len(first_set) == 0 or len(second_set) == 0:
+        return None
+    first_stack = np.stack([first.samples for first in first_set])
+    second_stack = np.stack([second.samples for second in second_set])
+    results = np.empty(len(first_set[0].samples), dtype=first_set[0].samples.dtype)
+    for i in range(len(first_set[0].samples)):
+        results[i] = ks_2samp(first_stack[..., i], second_stack[..., i])[0]
+    return CombinedTrace(None, None, results, parents=[*first_set, *second_set])
diff --git a/test/sca/test_ttest.py b/test/sca/test_test.py
index 02f78c5..898c6d6 100644
--- a/test/sca/test_ttest.py
+++ b/test/sca/test_test.py
@@ -1,7 +1,8 @@
 from unittest import TestCase
 
 import numpy as np
-from pyecsca.sca import Trace, welch_ttest, student_ttest
+
+from pyecsca.sca import Trace, welch_ttest, student_ttest, ks_test
 
 
 class TTestTests(TestCase):
@@ -17,9 +18,22 @@ class TTestTests(TestCase):
         a = Trace(None, None, np.array([19.8, 20.4, 19.6, 17.8, 18.5, 18.9, 18.3, 18.9, 19.5, 22.0]))
         b = Trace(None, None, np.array([28.2, 26.6, 20.1, 23.3, 25.2, 22.1, 17.7, 27.6, 20.6, 13.7]))
         c = Trace(None, None, np.array([20.2, 21.6, 27.1, 13.3, 24.2, 20.1, 11.7, 25.6, 26.6, 21.4]))
-        
+
         result = welch_ttest([a, b], [b, c])
         self.assertIsNotNone(result)
 
     def test_students_ttest(self):
+        self.assertIsNone(student_ttest([], []))
         self.assertIsNotNone(student_ttest([self.a, self.b], [self.c, self.d]))
+
+
+class KolmogorovSmirnovTests(TestCase):
+
+    def test_ks_test(self):
+        self.assertIsNone(ks_test([], []))
+
+        a = Trace(None, b"\xff", np.array([20, 80], dtype=np.dtype("i1")))
+        b = Trace(None, b"\xff", np.array([30, 42], dtype=np.dtype("i1")))
+        c = Trace(None, b"\x00", np.array([78, 56], dtype=np.dtype("i1")))
+        d = Trace(None, b"\x00", np.array([98, 36], dtype=np.dtype("i1")))
+        self.assertIsNotNone(ks_test([a, b], [c, d]))