Add secure np.cum{ulative_}sum().

Author: lschoe

mpyc/runtime.py

@@ -3143,6 +3143,48 @@ async def np_sum(self, a, axis=None, keepdims=False, initial=0):
         return np.sum(a, axis=axis, keepdims=keepdims, initial=initial.value)
         # TODO: handle switch from initial (field elt) to initial.value inside finfields.py
 
+    @asyncoro.mpc_coro_no_pc
+    async def np_cumsum(self, a, axis=None):
+        """Secure cumulative sum of array a along given axis.
+
+        If axis is None, array a is flattened first.
+        """
+        shape = (a.size,) if a.shape == () or axis is None else a.shape
+        if isinstance(a, self.SecureFixedPointArray):
+            rettype = (type(a), a.integral, shape)
+        else:
+            rettype = (type(a), shape)
+        await self.returnType(rettype)
+        a = await self.gather(a)
+        return np.cumsum(a, axis=axis)
+
+    @asyncoro.mpc_coro_no_pc
+    async def np_cumulative_sum(self, a, axis=None, include_initial=False):
+        """Secure cumulative sum of array a along given axis.
+
+        Only for 0D and 1D arrays, axis is allowed to be None.
+        If include_initial holds, the initial zero(s) are included in the output.
+        """
+        if axis is None:
+            if a.ndim < 2:
+                axis = 0
+            else:
+                raise ValueError('For arrays which have more than one dimension '
+                                 '``axis`` argument is required.')
+
+        shape = (a.size,) if a.shape == () else a.shape
+        if include_initial:
+            shape = list(shape)
+            shape[axis] += 1
+            shape = tuple(shape)
+        if isinstance(a, self.SecureFixedPointArray):
+            rettype = (type(a), a.integral, shape)
+        else:
+            rettype = (type(a), shape)
+        await self.returnType(rettype)
+        a = await self.gather(a)
+        return np.cumulative_sum(a, axis=axis, include_initial=include_initial)
+
     @asyncoro.mpc_coro_no_pc
     async def np_negative(self, a):
         """Secure elementwise negation -a (additive inverse) of a."""

tests/test_runtime.py

@@ -246,6 +246,9 @@ def __lt__(self, other):
         np.assertEqual(mpc.run(mpc.output(c.sum(keepdims=True))), a.sum(keepdims=True))
         np.assertEqual(mpc.run(mpc.output(c.sum(axis=(0, 2), keepdims=True))),
                        a.sum(axis=(0, 2), keepdims=True))
+        np.assertEqual(mpc.run(mpc.output(np.cumsum(c))), np.cumsum(a))
+        np.assertEqual(mpc.run(mpc.output(np.cumulative_sum(c, axis=1, include_initial=True))),
+                       np.cumulative_sum(a, axis=1, include_initial=True))
         np.assertEqual(mpc.run(mpc.output(c**254 * c**0 * c**-253)), a)
 
         # TODO: c //= 2 secure int __floordiv__() etc.
@@ -317,6 +320,9 @@ def test_secfxp_array(self):
         np.assertEqual(mpc.run(mpc.output(np.equal(c, 0))), False)
         np.assertEqual(mpc.run(mpc.output(np.sum(c, axis=(-2, 1)))), np.sum(a, axis=(-2, 1)))
         np.assertEqual(mpc.run(mpc.output(c.sum(axis=1, initial=1.5))), a.sum(axis=1, initial=1.5))
+        np.assertEqual(mpc.run(mpc.output(np.cumsum(c, axis=0))), np.cumsum(a, axis=0))
+        np.assertEqual(mpc.run(mpc.output(np.cumulative_sum(c, axis=1, include_initial=True))),
+                       np.cumulative_sum(a, axis=1, include_initial=True))
         self.assertEqual(np.prod(c, axis=(-2, 1)).integral, False)
         np.assertEqual(mpc.run(mpc.output(np.prod(c, axis=(-2, 1)))), np.prod(a, axis=(-2, 1)))
         a = a.flatten()[:3]
@@ -393,6 +399,16 @@ def test_secfxp_array(self):
         self.assertEqual(np.sum(c1).integral, False)
         self.assertEqual(np.sum(c1, axis=0).integral, False)
 
+        self.assertEqual(np.cumsum(c2).integral, True)
+        self.assertEqual(np.cumsum(c2, axis=0).integral, True)
+        self.assertEqual(np.cumsum(c1).integral, False)
+        self.assertEqual(np.cumsum(c1, axis=0).integral, False)
+
+        self.assertEqual(np.cumulative_sum(c2, axis=1).integral, True)
+        self.assertEqual(np.cumulative_sum(c2, axis=0).integral, True)
+        self.assertEqual(np.cumulative_sum(c1, axis=1).integral, False)
+        self.assertEqual(np.cumulative_sum(c1, axis=0).integral, False)
+
         self.assertEqual(mpc.np_sgn(c1).integral, True)
         self.assertEqual(mpc.np_sgn(c2).integral, True)
         self.assertEqual(np.absolute(c1).integral, False)
@@ -505,6 +521,10 @@ def test_secfld_array(self):
         np.assertEqual(mpc.run(mpc.output(np.outer(a, c))), np.outer(a, a))
         np.assertEqual(mpc.run(mpc.output(np.convolve(a[0][0], c[0][0]))),
                        np.convolve(a[0][0], a[0][0]))
+        np.assertEqual(mpc.run(mpc.output(np.sum(c))), np.sum(a))
+        np.assertEqual(mpc.run(mpc.output(np.cumsum(c))), np.cumsum(a))
+        np.assertEqual(mpc.run(mpc.output(np.cumulative_sum(c, axis=0, include_initial=True))),
+                       np.cumulative_sum(a, axis=0, include_initial=True))
         np.assertEqual(mpc.run(mpc.output(np.roll(c, 1))), np.roll(a, 1))
         self.assertEqual(len(c), 1)
         self.assertEqual(len(c.T), 2)
@@ -525,6 +545,7 @@ def test_np_errors(self):
         self.assertRaises(ValueError, np.convolve, c, c)
         self.assertRaises(ValueError, np.convolve, c[0], c[0][:0])
         self.assertRaises(ValueError, np.convolve, c[0][:0], c[0])
+        self.assertRaises(ValueError, np.cumulative_sum, c)
 
     def test_async(self):
         mpc.options.no_async = False