1# Primitive Root of Unity
  2def get_primitive_root_of_unity(F, n):
  3    # using the method described by Thomas Pornin in
  4    # https://crypto.stackexchange.com/a/63616
  5    q = F.order()
  6    for k in range(q):
  7        if k==0:
  8            continue
  9        g = F(k)
 10        #  g = F.random_element()
 11        if g==0:
 12            continue
 13        w = g ^ ((q-1)/n)
 14        if w^(n/2) != 1:
 15            return g, w
 16
 17# Roots of Unity
 18def get_nth_roots_of_unity(n, primitive_w):
 19    w = [0]*n
 20    for i in range(n):
 21        w[i] = primitive_w^i
 22    return w
 23
 24
 25# fft (Fast Fourier Transform) returns:
 26# - nth roots of unity
 27# - Vandermonde matrix for the nth roots of unity
 28# - Inverse Vandermonde matrix
 29def fft(F, n):
 30    g, primitive_w = get_primitive_root_of_unity(F, n)
 31
 32    w = get_nth_roots_of_unity(n, primitive_w)
 33
 34    ft = matrix(F, n)
 35    for j in range(n):
 36        row = []
 37        for k in range(n):
 38            row.append(primitive_w^(j*k))
 39        ft.set_row(j, row)
 40    ft_inv = ft^-1
 41    return w, ft, ft_inv
 42
 43
 44# Fast polynomial multiplication using FFT
 45def poly_mul(fa, fb, F, n):
 46    w, ft, ft_inv = fft(F, n)
 47
 48    # compute evaluation points from polynomials fa & fb at the roots of unity
 49    a_evals = []
 50    b_evals = []
 51    for i in range(n):
 52        a_evals.append(fa(w[i]))
 53        b_evals.append(fb(w[i]))
 54
 55    # multiply elements in a_evals by b_evals
 56    c_evals = map(operator.mul, a_evals, b_evals)
 57    c_evals = vector(c_evals)
 58
 59    # using FFT, convert the c_evals into fc(x)
 60    fc_coef = c_evals*ft_inv
 61    fc2=P(fc_coef.list())
 62    return fc2, c_evals
 63
 64
 65# Tests
 66
 67#####
 68# Roots of Unity test:
 69q = 17
 70F = GF(q)
 71n = 4
 72g, primitive_w = get_primitive_root_of_unity(F, n)
 73print("generator:", g)
 74print("primitive_w:", primitive_w)
 75
 76w = get_nth_roots_of_unity(n, primitive_w)
 77print(f"{n}th roots of unity: {w}")
 78assert w == [1, 13, 16, 4]
 79
 80
 81#####
 82# FFT test:
 83
 84def isprime(num):
 85    for n in range(2,int(num^1/2)+1):
 86        if num%n==0:
 87            return False
 88    return True
 89
 90# list valid values for q
 91for i in range(20):
 92    if isprime(8*i+1):
 93        print("q =", 8*i+1)
 94
 95q = 41
 96F = GF(q)
 97n = 4
 98# q needs to be a prime, s.t. q-1 is divisible by n
 99assert (q-1)%n==0
100print("q =", q, "n = ", n)
101
102# ft: Vandermonde matrix for the nth roots of unity
103w, ft, ft_inv = fft(F, n)
104print("nth roots of unity:", w)
105print("Vandermonde matrix:")
106print(ft)
107
108fa_eval = vector([3,4,5,9])
109print("fa_eval:", fa_eval)
110
111# interpolate f_a(x)
112fa_coef = ft_inv * fa_eval
113print("fa_coef:", fa_coef)
114
115P.<x> = PolynomialRing(F)
116fa = P(list(fa_coef))
117print("f_a(x):", fa)
118
119# check that evaluating fa(x) at the roots of unity returns the expected values of fa_eval
120for i in range(len(fa_eval)):
121    assert fa(w[i]) == fa_eval[i]
122
123# go from coefficient form to evaluation form
124fa_eval2 = ft * fa_coef
125print("fa_eval'", fa_eval)
126assert fa_eval2 == fa_eval
127
128
129# Fast polynomial multiplication using FFT
130print("\n---------")
131print("---Fast polynomial multiplication using FFT")
132
133n = 8
134# q needs to be a prime, s.t. q-1 is divisible by n
135assert (q-1)%n==0
136print("q =", q, "n = ", n)
137
138fa=P([1,2,3,4])
139fb=P([1,2,3,4])
140fc_expected = fa*fb
141print("fc expected result:", fc_expected) # expected result
142print("fc expected coef", fc_expected.coefficients())
143
144fc, c_evals = poly_mul(fa, fb, F, n)
145print("c_evals=(a_evals*b_evals)=", c_evals)
146print("fc:", fc)
147assert fc_expected == fc