diff --git a/algebra/ff/babybear.py b/algebra/ff/babybear.py index e0950f6..b2ee162 100644 --- a/algebra/ff/babybear.py +++ b/algebra/ff/babybear.py @@ -4,3 +4,4 @@ class BabyBear(PrimeField): # BabyBear prime: 2^31 - 2^27 + 1 P = 2013265921 + w = 31 diff --git a/algebra/ff/m31.py b/algebra/ff/m31.py index 1e61c5a..9675479 100644 --- a/algebra/ff/m31.py +++ b/algebra/ff/m31.py @@ -3,4 +3,5 @@ class M31(PrimeField): # Mersenne31 prime: 2^31 - 1 - P = 2147483647 + P = 2**31 - 1 + w = 7 diff --git a/algebra/ff/prime_field.py b/algebra/ff/prime_field.py index 504593f..f8fbded 100644 --- a/algebra/ff/prime_field.py +++ b/algebra/ff/prime_field.py @@ -4,6 +4,7 @@ class PrimeField: P: int = None + w: int = None def __init__(self, x): if isinstance(x, (int, float, list, Tensor)): diff --git a/algebra/ff/test.py b/algebra/ff/test.py deleted file mode 100644 index e4760fd..0000000 --- a/algebra/ff/test.py +++ /dev/null @@ -1,6 +0,0 @@ -from tinygrad import Tensor - -a = Tensor([1, 2, 3]) -b = Tensor([4, 5, 6]) - -print((a % 3).numpy()) diff --git a/algebra/poly/ntt.py b/algebra/poly/ntt.py new file mode 100644 index 0000000..e2d7542 --- /dev/null +++ b/algebra/poly/ntt.py @@ -0,0 +1,122 @@ +from tinygrad.tensor import Tensor +from tinygrad import dtypes +import numpy as np + + +def _next_valid_n(l, p): + """Find next n >= l that divides p-1 efficiently""" + p_minus_1 = p - 1 + for n in range(l, min(p, 2 * l)): + if p_minus_1 % n == 0: + return n + + bit_length = l.bit_length() + power_of_two = 1 << bit_length + if power_of_two < l: + power_of_two = power_of_two << 1 + + while power_of_two < p: + if p_minus_1 % power_of_two == 0: + return power_of_two + power_of_two = power_of_two << 1 + + for divisor in range(l, int(p_minus_1**0.5) + 1): + if p_minus_1 % divisor == 0: + quotient = p_minus_1 // divisor + if quotient >= l: + return quotient + return divisor + + return p - 1 + + +def ntt(polynomial: Tensor, prime: int, primitive_root: int) -> Tensor: + """ + Compute the Number Theoretic Transform of a polynomial. + + Args: + polynomial: Coefficient vector of the polynomial (Tensor or list/array) + prime: The prime modulus + primitive_root: A primitive root modulo prime + + Returns: + The NTT of the polynomial as a Tensor + """ + return _transform(polynomial, prime, primitive_root, inverse=False) + + +def intt(transformed: Tensor, prime: int, primitive_root: int) -> Tensor: + """ + Compute the Inverse Number Theoretic Transform. + + Args: + transformed: The transformed polynomial (Tensor or list/array) + prime: The prime modulus + primitive_root: A primitive root modulo prime + + Returns: + The original polynomial coefficients as a Tensor + """ + return _transform(transformed, prime, primitive_root, inverse=True) + + +def _transform(x, prime: int, primitive_root: int, inverse: bool = False) -> Tensor: + """Internal function to perform NTT or INTT transformation""" + if not isinstance(x, Tensor): + x = Tensor(x, dtype=dtypes.uint64) + + dtype = x.dtype + n = _next_valid_n(len(x), prime) + print(f"n: {n}, prime: {prime}, primitive_root: {primitive_root}") + + if len(x) < n: + padded_np = Tensor.zeros(n, dtype=dtype).contiguous() + padded_np[: len(x)] = x + x = padded_np + + # Create powers matrix (i*j) % n + i = Tensor.arange(n, dtype=dtypes.uint64).reshape(n, 1) + j = Tensor.arange(n, dtype=dtypes.uint64).reshape(1, n) + powers = (i * j) % n + + # Get omega (or inverse omega for INTT) + omega = pow(primitive_root, (prime - 1) // n, prime) + if inverse: + omega = pow(omega, prime - 2, prime) # Modular inverse + + # Compute all powers of omega + omega_powers = np.ones(n, dtype=np.uint64) + current = 1 + for k in range(n): + omega_powers[k] = current + current = (current * omega) % prime + + # Create transformation matrix + matrix = Tensor(omega_powers, dtype=dtypes.uint64)[powers] + result = (matrix @ (x % prime)) % prime + + # For inverse, apply scaling factor (1/n mod prime) + if inverse: + n_inv = pow(n, prime - 2, prime) + result = (result * n_inv) % prime + + # Perform the transformation + return result + + +if __name__ == "__main__": + from algebra.ff.m31 import M31 + from random import randint + + n, prime, primitive_root = 10, M31.P, M31.w + polynomial = Tensor([randint(0, prime - 1) for _ in range(n)], dtype=dtypes.uint64) + print(f"polynomial: {polynomial.numpy()}") + + transformed = ntt(polynomial, prime, primitive_root) + print("NTT result:", transformed.numpy()) + + original = intt(transformed, prime, primitive_root) + print("INTT result:", original.numpy()) + + assert (polynomial.numpy() == original.numpy()[: len(polynomial)]).all(), "INTT(NTT(polynomial)) != polynomial" + print("Verification successful") diff --git a/algebra/poly/univariate.py b/algebra/poly/univariate.py index ca5c83c..6867ca4 100644 --- a/algebra/poly/univariate.py +++ b/algebra/poly/univariate.py @@ -1,6 +1,8 @@ from tinygrad.tensor import Tensor from tinygrad import dtypes from algebra.ff.prime_field import PrimeField as PF +from algebra.poly.ntt import ntt, intt +import numpy as np class Polynomial: @@ -14,17 +16,15 @@ class Polynomial: """ def __init__(self, coeffs: list[int] | Tensor, prime_field: PF = None): - """ - Initialize the polynomial. - - coeffs: a list of field elements (instances of a PrimeField subclass) - prime_field: optional prime field class to use for modular arithmetic - """ + """Initialize polynomial with coefficients and optional prime field.""" self.PrimeField = prime_field if isinstance(coeffs, list): + coeffs = np.trim_zeros(coeffs, "b") self.coeffs = Tensor(coeffs, dtype=dtypes.int32) - elif isinstance(coeffs, Tensor): - self.coeffs = coeffs + else: + coeffs_np = coeffs.numpy() + coeffs_np = np.trim_zeros(coeffs_np, "b") + self.coeffs = Tensor(coeffs_np, dtype=coeffs.dtype) def degree(self) -> int: """ @@ -49,6 +49,20 @@ def evaluate(self, x: int | Tensor): result = result * x + coeff return result + def ntt(self): + """ + Compute the Number Theoretic Transform of the polynomial. + """ + p_ntt = ntt(self.coeffs.cast(dtypes.uint64), self.PrimeField.P, self.PrimeField.w).cast(self.coeffs.dtype) + return Polynomial(p_ntt, self.PrimeField) + + def intt(self): + """ + Compute the Inverse Number Theoretic Transform of the polynomial. + """ + p_intt = intt(self.coeffs.cast(dtypes.uint64), self.PrimeField.P, self.PrimeField.w).cast(self.coeffs.dtype) + return Polynomial(p_intt, self.PrimeField) + def __evaluate_all(self, xs: Tensor): """ Evaluate the polynomial at all elements in xs using Horner's method. diff --git a/pyproject.toml b/pyproject.toml index 9f6e1a0..67ae1dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ authors = [{ name = "Pia", email = "gayeongparkk@gmail.com" }] requires-python = ">=3.10" dependencies = [ "numpy>=2.2.3", + "sympy>=1.13.3", "tinygrad>=0.10.2", ] diff --git a/tests/test_poly.py b/tests/test_poly.py index aa8753c..b9b87dd 100644 --- a/tests/test_poly.py +++ b/tests/test_poly.py @@ -2,6 +2,7 @@ from algebra.ff.m31 import M31 from algebra.ff.babybear import BabyBear from tinygrad import Tensor +from random import randint def test_polynomial_operations_m31(): @@ -70,3 +71,8 @@ def test_polynomial_operations_babybear(): # Test evaluate_all result = p1(Tensor([1, 2, 3])).numpy() assert (result == [6, 17, 34]).all() + + p7 = Polynomial([randint(0, 100) for _ in range(8)], M31) + p7_ntt = p7.ntt() + p7_intt = p7_ntt.intt() + assert (p7_intt.coeffs.numpy() == p7.coeffs.numpy()).all() diff --git a/uv.lock b/uv.lock index dc6fdd9..a490d78 100644 --- a/uv.lock +++ b/uv.lock @@ -8,6 +8,7 @@ version = "0.1.0" source = { virtual = "." } dependencies = [ { name = "numpy" }, + { name = "sympy" }, { name = "tinygrad" }, ] @@ -20,6 +21,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "numpy", specifier = ">=2.2.3" }, + { name = "sympy", specifier = ">=1.13.3" }, { name = "tinygrad", specifier = ">=0.10.2" }, ] @@ -56,6 +58,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374", size = 5892 }, ] +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198 }, +] + [[package]] name = "numpy" version = "2.2.3" @@ -178,6 +189,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/63/6a/aca01554949f3a401991dc32fe22837baeaccb8a0d868256cbb26a029778/ruff-0.9.7-py3-none-win_arm64.whl", hash = "sha256:b075a700b2533feb7a01130ff656a4ec0d5f340bb540ad98759b8401c32c2037", size = 10177763 }, ] +[[package]] +name = "sympy" +version = "1.13.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/11/8a/5a7fd6284fa8caac23a26c9ddf9c30485a48169344b4bd3b0f02fef1890f/sympy-1.13.3.tar.gz", hash = "sha256:b27fd2c6530e0ab39e275fc9b683895367e51d5da91baa8d3d64db2565fec4d9", size = 7533196 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/ff/c87e0622b1dadea79d2fb0b25ade9ed98954c9033722eb707053d310d4f3/sympy-1.13.3-py3-none-any.whl", hash = "sha256:54612cf55a62755ee71824ce692986f23c88ffa77207b30c1368eda4a7060f73", size = 6189483 }, +] + [[package]] name = "tinygrad" version = "0.10.2"