feat: Added get_item, get_attr and shape (#5)

Author: jcabrero

nada_algebra/__init__.py

@@ -1,3 +1,4 @@
 """This is the __init__.py module"""
 
+from nada_algebra.array import NadaArray
 from nada_algebra.funcs import *

nada_algebra/array.py

@@ -31,6 +31,46 @@ class NadaArray:
 
     inner: np.ndarray
 
+    def __getitem__(self, item):
+        """
+        Get an item from the array.
+
+        Args:
+            item: The item to retrieve.
+
+        Returns:
+            NadaArray: A new NadaArray representing the retrieved item.
+        """
+        if len(self.inner.shape) == 1:
+            return self.inner[item]
+        return NadaArray(self.inner[item])
+
+    def __setitem__(self, key, value):
+        """
+        Set an item in the array.
+
+        Args:
+            key: The key to set.
+            value: The value to set.
+        """
+        if isinstance(value, NadaArray):
+            # print("NadaArray")
+            self.inner[key] = value.inner
+        else:
+            self.inner[key] = value
+
+    def __getattr__(self, name: str):
+        """
+        Get an attribute from the array.
+
+        Args:
+            name (str): The attribute name.
+
+        Returns:
+            NadaArray: A new NadaArray representing the retrieved attribute.
+        """
+        return getattr(self.inner, name)
+
     def __add__(
         self,
         other: Union[
@@ -268,7 +308,17 @@ def output_array(array: np.ndarray, party: Party, prefix: str) -> list:
         Returns:
             list: A list of Output objects.
         """
-        if isinstance(array, (SecretInteger, SecretUnsignedInteger)):
+        if isinstance(
+            array,
+            (
+                SecretInteger,
+                SecretUnsignedInteger,
+                PublicInteger,
+                PublicUnsignedInteger,
+                Integer,
+                UnsignedInteger,
+            ),
+        ):
             return [Output(array, f"{prefix}_0", party)]
 
         if len(array.shape) == 1:

nada_algebra/client.py

@@ -3,6 +3,7 @@
 secret and public variable integers and generating named party objects and input dictionaries.
 """
 
+from typing import Union
 from py_nillion_client import (
     SecretInteger,
     SecretUnsignedInteger,
@@ -29,12 +30,12 @@ def parties(num: int, prefix: str = "Party") -> list:
 def array(
     arr: np.ndarray,
     prefix: str,
-    nada_type: (
-        type(SecretInteger)
-        | type(SecretUnsignedInteger)
-        | type(PublicVariableInteger)
-        | type(PublicVariableUnsignedInteger)
-    ) = SecretInteger,
+    nada_type: Union[
+        SecretInteger,
+        SecretUnsignedInteger,
+        PublicVariableInteger,
+        PublicVariableUnsignedInteger,
+    ] = SecretInteger,
 ) -> dict:
     """
     Recursively generates a dictionary of Nillion input objects for each element

nada_algebra/funcs.py

@@ -141,4 +141,4 @@ def output(arr: NadaArray, party: Party, prefix: str):
     Returns:
         list: A list of Output objects.
     """
-    return arr.output(party, prefix)
+    return NadaArray.output_array(arr, party, prefix)

poetry.lock

@@ -0,0 +1,7 @@
+name = "matrix_inverse"
+version = "0.1.0"
+authors = [""]
+
+[[programs]]
+path = "src/main.py"
+prime_size = 64

tests/gauss_jordan/nada-project.toml

@@ -0,0 +1,93 @@
+from nada_dsl import *
+import numpy as np
+
+import nada_algebra as na
+from nada_algebra.array import NadaArray
+
+# from nada_crypto import random_lu_matrix, public_modular_inverse
+
+LOG_SCALE = 16
+SCALE = 1 << LOG_SCALE
+PRIME_64 = 18446744072637906947
+PRIME_128 = 340282366920938463463374607429104828419
+PRIME_256 = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF98C00003
+PRIME = PRIME_64
+
+
+def public_modular_inverse(
+    value: Integer | UnsignedInteger, modulo: int
+) -> PublicUnsignedInteger | UnsignedInteger:
+    """
+    Calculates the modular inverse of a value with respect to a prime modulus.
+
+    Args:
+        `value`: The value for which the modular inverse is to be calculated.
+        `modulo`: The prime modulo with respect to which the modular inverse is to be calculated.
+
+    Returns:
+        The modular inverse of the value with respect to the modulo.
+
+    Raises:
+        Exception: If the input type is not a `PublicUnsignedInteger` or `UnsignedInteger`.
+    """
+    return value ** UnsignedInteger(modulo - 2)
+
+
+def gauss_jordan_zn(mat: na.NadaArray, modulo: int):
+    """
+    Perform Gauss-Jordan elimination on Z_n on a given matrix.
+
+    Parameters:
+    - `matrix` (numpy.ndarray): The input matrix to perform Gauss-Jordan elimination on.
+    - `modulo` (int): The modulo representing the field `Z_n`
+
+    Returns:
+    numpy.ndarray: The reduced row echelon form of the input matrix.
+    """
+
+    # Make a copy of the matrix to avoid modifying the original
+    rows = mat.inner.shape[0]
+    cols = mat.inner.shape[1]
+
+    # Forward elimination
+    for i in range(rows):
+        # Find pivot row
+        pivot_row = i
+        while pivot_row < rows and (mat[pivot_row][i] == UnsignedInteger(0)) is Boolean(
+            True
+        ):
+            pivot_row += 1
+
+        # Swap pivot row with current row
+        mat[[i, pivot_row]] = mat[[pivot_row, i]]
+
+        # Scale pivot row to have leading 1
+        diagonal_element = mat[i][i]
+        pivot_inv = public_modular_inverse(diagonal_element, modulo)
+
+        mat[i] = mat[i] * pivot_inv
+
+        # Perform row operations to eliminate entries below pivot
+        for j in range(i + 1, rows):
+            factor = mat[j][i]
+            mat[j] -= mat[i] * factor
+
+    # Backward elimination
+    for i in range(rows - 1, -1, -1):
+        for j in range(i - 1, -1, -1):
+            factor = mat[j][i]
+            mat[j] -= mat[i] * factor
+
+    return mat
+
+
+def nada_main():
+    parties = na.parties(3)
+
+    A = na.array([3, 3], parties[0], "A", nada_type=SecretUnsignedInteger)
+
+    A = A.reveal()
+    A_inv = gauss_jordan_zn(A, PRIME)
+    outputs = na.output(A_inv, parties[2], "my_output")
+
+    return outputs

tests/gauss_jordan/src/main.py

@@ -0,0 +1,5 @@
+# This directory is kept purposely, so that no compilation errors arise.
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore

tests/gauss_jordan/target/.gitignore

@@ -0,0 +1,44 @@
+---
+program: main
+inputs:
+  secrets:
+    A_0_0:
+      SecretUnsignedInteger: "2"
+    A_0_1:
+      SecretUnsignedInteger: "4"
+    A_0_2:
+      SecretUnsignedInteger: "6"
+    A_1_0:
+      SecretUnsignedInteger: "1"
+    A_1_1:
+      SecretUnsignedInteger: "3"
+    A_1_2:
+      SecretUnsignedInteger: "5"
+    A_2_0:
+      SecretUnsignedInteger: "3"
+    A_2_1:
+      SecretUnsignedInteger: "1"
+    A_2_2:
+      SecretUnsignedInteger: "2"
+  public_variables: 
+    B: 
+      UnsignedInteger: "456"
+expected_outputs:
+  my_output_0_0:
+    UnsignedInteger: "1"
+  my_output_0_1:
+    UnsignedInteger: "0"
+  my_output_0_2:
+    UnsignedInteger: "0"
+  my_output_1_0:
+    UnsignedInteger: "0"
+  my_output_1_1:
+    UnsignedInteger: "1"
+  my_output_1_2:
+    UnsignedInteger: "0"
+  my_output_2_0:
+    UnsignedInteger: "0"
+  my_output_2_1:
+    UnsignedInteger: "0"
+  my_output_2_2:
+    UnsignedInteger: "1"

tests/gauss_jordan/tests/gauss_jordan.yaml

@@ -0,0 +1,7 @@
+name = "get_item"
+version = "0.1.0"
+authors = [""]
+
+[[programs]]
+path = "src/main.py"
+prime_size = 128