add fft sage impls and rename halo dir to zk

script/research/zk/fft/fft2.sage

@@ -0,0 +1,132 @@
+# sage: w.multiplicative_order()
+# 1330
+# sage: 11^3 - 1
+# 1330
+# sage: K.<w> = GF(11^3, repr="int")
+# sage: w
+# 11
+
+p = 199
+# sage: factor(199^3 - 1)
+# 2 * 3^3 * 11 * 13267
+n = 3^3
+n = 10
+n = 5
+
+assert p.is_prime()
+
+def find_ext_order(p, n):
+    N = 1
+    while True:
+        pNx_order = p^N - 1
+
+        # Does n divide the group order 𝔽_{p^N}^×?
+        if pNx_order % n == 0:
+            return N
+
+        N += 1
+
+def find_nth_root_unity(K, n):
+    # It cannot be a quadratic residue if n is odd
+    #assert n % 2 == 1
+
+    # So there is an nth root of unity in p^N. Now we have to find it.
+    pNx_order = p^N - 1
+
+    ω = K.gens()[0]
+    ω = ω^(pNx_order/n)
+    assert ω^n == 1
+    assert ω^(n - 1) != 1
+
+    return ω
+
+N = find_ext_order(p, n)
+print(f"N = {N}")
+print()
+K.<a> = GF(p^N, repr="int")
+ω = find_nth_root_unity(K, n)
+
+L.<X> = K[]
+
+f = 3*X^4 + 7*X^3 + X^2 + 4
+g = 2*X^4 + 2*X^2 + 110
+f = X^2 + 2*X + 4
+g = 2*X^2 + 110
+assert f.degree() < n/2
+assert g.degree() < n/2
+assert f.degree() + g.degree() < n
+print(f"f = {f}")
+print(f"g = {g}")
+print(f"fg = {f*g}")
+print()
+
+def vectorify(f):
+    assert f.degree() < n
+    fT = vector([f[i] for i in range(f.degree() + 1)] +
+                # Zero padding
+                [0 for _ in range(n - f.degree() - 1)])
+    assert len(fT) == n
+    # Just check decomposed polynomial is in the correct order
+    assert sum([fT[i]*X^i for i in range(n)]) == f
+    return fT
+
+fT = vectorify(f)
+gT = vectorify(g)
+
+def nXn_vandermonde(n, ω):
+    # We hardcode this one so you know what is looks like
+    if n == 5:
+        Vω = matrix([
+            [1,   1,   1,   1,   1],
+            [1, ω^1, ω^2, ω^3, ω^4],
+            [1, ω^2, ω^4, ω^1, ω^3],
+            [1, ω^3, ω^1, ω^4, ω^2],
+            [1, ω^4, ω^3, ω^2, ω^1],
+        ])
+        return Vω
+
+    # This is the code to generate it
+    Vω = matrix([[ω^(i * j) for j in range(n)] for i in range(n)])
+    return Vω
+
+Vω = nXn_vandermonde(n, ω)
+Vω_inv = nXn_vandermonde(n, ω^-1)/n
+# Lemma: V_ω^{-1} = 1/n V_{ω^-1}
+assert Vω^-1 == Vω_inv
+
+DFT_ω_f = Vω * fT
+f_evals = [f(X=ω^i) for i in range(n)]
+print(f"DFT_ω(f) = {DFT_ω_f}")
+print(f"f(ω^i) = {f_evals}")
+print()
+
+DFT_ω_g = Vω * gT
+g_evals = [g(X=ω^i) for i in range(n)]
+print(f"DFT_ω(g) = {DFT_ω_g}")
+print(f"g(ω^i) = {g_evals}")
+print()
+
+def convolution(f, g):
+    return f*g % (X^n - 1)
+def pointwise_prod(fT, gT):
+    return [a_i*b_i for a_i, b_i in zip(fT, gT)]
+
+print(f"deg(f) + deg(g) = {f.degree() + g.degree()}")
+fжg = convolution(f, g)
+print(f"f☼g = {fжg}")
+assert fжg == f*g
+fжgT = vectorify(fжg)
+
+DFT_ω_fжg = Vω * fжgT
+for i in range(n):
+    assert fжg(X=ω^i) == f(ω^i)*g(ω^i)
+print(f"DFT_ω(f☼g) = {DFT_ω_fжg}")
+DFT_fg_prod = vector(pointwise_prod(DFT_ω_f, DFT_ω_g))
+print(f"DFT_ω(f)·DFT_ω(g) = {DFT_fg_prod}")
+assert DFT_ω_fжg == DFT_fg_prod
+
+inv_DFT_fg = Vω_inv * DFT_fg_prod
+fgT = vectorify(f*g)
+assert inv_DFT_fg == fgT
+print(f"DFT^-1(DFT_ω(f)·DFT_ω(g)) = {inv_DFT_fg}")
+

script/research/zk/fft/fft3.sage

@@ -0,0 +1,116 @@
+p = 199
+n = 4
+
+assert p.is_prime()
+
+def find_ext_order(p, n):
+    N = 1
+    while True:
+        pNx_order = p^N - 1
+
+        # Does n divide the group order 𝔽_{p^N}^×?
+        if pNx_order % n == 0:
+            return N
+
+        N += 1
+
+def find_nth_root_unity(K, n):
+    # It cannot be a quadratic residue if n is odd
+    #assert n % 2 == 1
+
+    # So there is an nth root of unity in p^N. Now we have to find it.
+    pNx_order = p^N - 1
+
+    ω = K.gens()[0]
+    ω = ω^(pNx_order/n)
+    assert ω^n == 1
+    assert ω^(n - 1) != 1
+
+    return ω
+
+N = find_ext_order(p, n)
+print(f"p = {p}")
+print(f"n = {n}")
+print(f"N = {N}")
+print(f"p^N = {p^N}")
+K.<a> = GF(p^N, repr="int")
+ω = find_nth_root_unity(K, n)
+print(f"ω = {ω}")
+print()
+
+L.<X> = K[]
+
+f = 10*X + 110
+assert f.degree() < n/2
+print(f"f = {f}")
+print()
+
+def vectorify(f):
+    assert f.degree() < n
+    fT = vector([f[i] for i in range(f.degree() + 1)] +
+                # Zero padding
+                [0 for _ in range(n - f.degree() - 1)])
+    assert len(fT) == n
+    # Just check decomposed polynomial is in the correct order
+    assert sum([fT[i]*X^i for i in range(n)]) == f
+    return fT
+
+def dot(a, b):
+    assert len(a) == len(b)
+    return [a_i*b_i for a_i, b_i in zip(a, b)]
+
+assert n == 2^2
+m = 4
+print(f"m = {m}")
+ω_powers = vector(ω^i for i in range(m/2))
+print(f"ω^i = {ω_powers}")
+
+fT = vectorify(f)
+print(f"fT  = {fT}")
+
+# Rewrite f(X) = g(X) + X^(n/2) h(X)
+f_g, f_h = vector(fT[:m/2]), vector(fT[m/2:])
+print(f"    = {f_g}, {f_h}")
+
+r8 =  f_g + f_h
+s8 = dot((f_g - f_h), ω_powers)
+assert len(r8) == len(s8) == m/2
+print(f"r8 = {r8}")
+print(f"s8 = {s8}")
+print()
+
+m = 2
+ω_powers = vector(ω_i for ω_i in ω_powers[::2])
+assert len(ω_powers) == m/2
+print(f"m = {m}")
+
+# Corresponds to r_4
+r4_g, r4_h = vector(r8[:m/2]), vector(r8[m/2:])
+print(f"r4_g, r4_h = {r4_g}, {r4_h}")
+
+r4_r2 = r4_g + r4_h
+r4_s2 = dot((r4_g - r4_h), ω_powers)
+print(f"r4_r2 = {r4_r2}")
+print(f"r4_s2 = {r4_s2}")
+print()
+
+# Corresponds to s_4
+s4_g, s4_h = vector(s8[:m/2]), vector(s8[m/2:])
+print(f"s4_g, s4_h = {s4_g}, {s4_h}")
+
+s4_r2 = s4_g + s4_h
+s4_s2 = dot((s4_g - s4_h), ω_powers)
+print(f"s4_r2 = {s4_r2}")
+print(f"s4_s2 = {s4_s2}")
+print()
+
+# Final step
+m = 1
+print(f"m = {m}")
+print("STOP")
+# Just return the values directly
+print()
+
+f_evals = [f(X=ω^i) for i in range(n)]
+print(f"f(ω^i) = {f_evals}")
+

script/research/zk/fft/fft4.sage

@@ -0,0 +1,101 @@
+import itertools
+
+p = 199
+n = 8
+
+assert p.is_prime()
+
+def find_ext_order(p, n):
+    N = 1
+    while True:
+        pNx_order = p^N - 1
+
+        # Does n divide the group order 𝔽_{p^N}^×?
+        if pNx_order % n == 0:
+            return N
+
+        N += 1
+
+def find_nth_root_unity(K, n):
+    # It cannot be a quadratic residue if n is odd
+    #assert n % 2 == 1
+
+    # So there is an nth root of unity in p^N. Now we have to find it.
+    pNx_order = p^N - 1
+
+    ω = K.gens()[0]
+    ω = ω^(pNx_order/n)
+    assert ω^n == 1
+    assert ω^(n - 1) != 1
+
+    return ω
+
+N = find_ext_order(p, n)
+print(f"p = {p}")
+print(f"n = {n}")
+print(f"N = {N}")
+print(f"p^N = {p^N}")
+K.<a> = GF(p^N, repr="int")
+ω = find_nth_root_unity(K, n)
+print(f"ω = {ω}")
+print()
+
+L.<X> = K[]
+
+f = 7*X^3 + X^2 + 110*X + 4
+assert f.degree() < n/2
+print(f"f = {f}")
+print()
+
+def vectorify(f):
+    assert f.degree() < n
+    fT = vector([f[i] for i in range(f.degree() + 1)] +
+                # Zero padding
+                [0 for _ in range(n - f.degree() - 1)])
+    assert len(fT) == n
+    # Just check decomposed polynomial is in the correct order
+    assert sum([fT[i]*X^i for i in range(n)]) == f
+    return fT
+
+def dot(a, b):
+    assert len(a) == len(b)
+    return [a_i*b_i for a_i, b_i in zip(a, b)]
+
+# ABC, DEF -> ADBECF
+def alternate(list1, list2):
+    return itertools.chain(*zip(list1, list2))
+
+def calc_dft(ω_powers, f):
+    m = len(f)
+    indent = " " * (n - m)
+    print(f"{indent}calc_dft({ω_powers}, {f})")
+    print(f"{indent}  m = {m}")
+    if m == 1:
+        print(f"{indent}  m = 1 so return f")
+        return f
+    g, h = vector(f[:m/2]), vector(f[m/2:])
+    print(f"{indent}  g = {g}")
+    print(f"{indent}  h = {h}")
+
+    r = g + h
+    s = dot(g - h, ω_powers)
+    print(f"{indent}  r = {r}")
+    print(f"{indent}  s = {s}")
+    print()
+
+    ω_powers = vector(ω_i for ω_i in ω_powers[::2])
+    rT = calc_dft(ω_powers, r)
+    sT = calc_dft(ω_powers, s)
+
+    result = list(alternate(rT, sT))
+    print(f"{indent}return {result}")
+    return result
+
+ω_powers = vector(ω^i for i in range(n/2))
+fT = vectorify(f)
+dft = calc_dft(ω_powers, fT)
+print()
+print(f"DFT(f) = {dft}")
+f_evals = [f(X=ω^i) for i in range(n)]
+print(f"f(ω^i) = {f_evals}")
+