Feature/array statistics (#24)

Add arithmetic mean operation

Author: mathias-nillion

nada_algebra/array.py

@@ -22,6 +22,7 @@
 from nada_algebra.types import (
     Rational,
     SecretRational,
+    rational,
     public_rational,
     secret_rational,
     get_log_scale,
@@ -326,6 +327,34 @@ def apply(self, func: Callable[[Any], Any]) -> "NadaArray":
         """
         return NadaArray(np.frompyfunc(func, 1, 1)(self.inner))
 
+    @copy_metadata(np.ndarray.mean)
+    def mean(self, axis=None, dtype=None, out=None, keepdims=False) -> Any:
+        sum_arr = self.inner.sum(axis=axis, dtype=dtype, keepdims=keepdims)
+
+        if self.dtype in (Rational, SecretRational):
+            nada_type = rational
+        else:
+            nada_type = Integer
+
+        if axis is None:
+            count = nada_type(self.size)
+        else:
+            if keepdims:
+                count = np.expand_dims(count, axis=axis)
+                count = np.frompyfunc(nada_type, 1, 1)(count)
+            else:
+                count = nada_type(self.shape[axis])
+
+        mean_arr = sum_arr / count
+
+        if out is not None:
+            out[...] = mean_arr
+            return out
+
+        if isinstance(mean_arr, np.ndarray):
+            return NadaArray(mean_arr)
+        return mean_arr
+
     @staticmethod
     def output_array(array: np.ndarray, party: Party, prefix: str) -> list:
         """

nada_algebra/funcs.py

@@ -419,6 +419,11 @@ def diagonal(a: NadaArray, *args, **kwargs):
     return a.diagonal(*args, **kwargs)
 
 
+@copy_metadata(np.diagonal)
+def mean(a: NadaArray, *args, **kwargs):
+    return a.mean(*args, **kwargs)
+
+
 @copy_metadata(np.prod)
 def prod(a: NadaArray, *args, **kwargs):
     return a.prod(*args, **kwargs)

tests/nada-tests/nada-project.toml

@@ -141,3 +141,7 @@ prime_size = 128
 [[programs]]
 path = "src/functional_operations.py"
 prime_size = 128
+
+[[programs]]
+path = "src/array_statistics.py"
+prime_size = 128

tests/nada-tests/src/array_statistics.py

@@ -0,0 +1,36 @@
+from nada_dsl import *
+import nada_algebra as na
+
+
+def nada_main():
+    parties = na.parties(2)
+
+    a = na.array([3, 2], parties[0], "A", SecretInteger)
+    b = na.array([3, 2], parties[0], "B", na.SecretRational)
+
+    a_sum = a.sum()
+    b_sum = b.sum()
+
+    a_sum_arr = a.sum(axis=0)
+    b_sum_arr = b.sum(axis=0)
+
+    a_mean = a.mean()
+    b_mean = b.mean()
+
+    a_mean_arr = a.mean(axis=0)
+    b_mean_arr = b.mean(axis=0)
+
+    output_1 = [
+        Output(a_sum, "a_sum", parties[1]),
+        Output(a_mean, "a_mean", parties[1]),
+        Output(b_sum.value, "b_sum", parties[1]),
+        Output(b_mean.value, "b_mean", parties[1]),
+    ]
+    output_2 = (
+        a_sum_arr.output(parties[1], "a_sum_arr")
+        + b_sum_arr.output(parties[1], "b_sum_arr")
+        + a_mean_arr.output(parties[1], "a_mean_arr")
+        + b_mean_arr.output(parties[1], "b_mean_arr")
+    )
+
+    return output_1 + output_2

tests/nada-tests/tests/array_statistics.yaml

@@ -0,0 +1,54 @@
+---
+program: array_statistics
+inputs:
+  secrets:
+    A_0_0:
+      SecretInteger: "1"
+    A_0_1:
+      SecretInteger: "2"
+    A_1_0:
+      SecretInteger: "4"
+    A_1_1:
+      SecretInteger: "5"
+    A_2_0:
+      SecretInteger: "4"
+    A_2_1:
+      SecretInteger: "5"
+    B_0_0:
+      SecretInteger: "1"
+    B_0_1:
+      SecretInteger: "2"
+    B_1_0:
+      SecretInteger: "4"
+    B_1_1:
+      SecretInteger: "5"
+    B_2_0:
+      SecretInteger: "4"
+    B_2_1:
+      SecretInteger: "5"
+  public_variables: {}
+expected_outputs:
+  a_mean_arr_1:
+    SecretInteger: "4"
+  b_mean_arr_0:
+    SecretInteger: "3"
+  b_sum:
+    SecretInteger: "21"
+  b_mean:
+    SecretInteger: "3"
+  a_mean_arr_0:
+    SecretInteger: "3"
+  b_mean_arr_1:
+    SecretInteger: "4"
+  b_sum_arr_0:
+    SecretInteger: "9"
+  a_sum:
+    SecretInteger: "21"
+  a_sum_arr_0:
+    SecretInteger: "9"
+  b_sum_arr_1:
+    SecretInteger: "12"
+  a_mean:
+    SecretInteger: "3"
+  a_sum_arr_1:
+    SecretInteger: "12"

tests/test_all.py

@@ -36,6 +36,7 @@
     "rational_advanced",
     "array_attributes",
     "functional_operations",
+    "array_statistics",
     # Not supported yet
     # "unsigned_matrix_inverse",
     # "private_inverse"