1load("fft.sage")
 2
 3#####
 4# Roots of Unity test:
 5q = 17
 6F = GF(q)
 7n = 4
 8g, primitive_w = get_primitive_root_of_unity(F, n)
 9print("generator:", g)
10print("primitive_w:", primitive_w)
11
12w = get_nth_roots_of_unity(n, primitive_w)
13print(f"{n}th roots of unity: {w}")
14assert w == [1, 13, 16, 4]
15
16
17#####
18# FFT test:
19
20def isprime(num):
21    for n in range(2,int(num^1/2)+1):
22        if num%n==0:
23            return False
24    return True
25
26# list valid values for q
27for i in range(20):
28    if isprime(8*i+1):
29        print("q =", 8*i+1)
30
31q = 41
32F = GF(q)
33n = 4
34# q needs to be a prime, s.t. q-1 is divisible by n
35assert (q-1)%n==0
36print("q =", q, "n = ", n)
37
38# ft: Vandermonde matrix for the nth roots of unity
39w, ft, ft_inv = fft(F, n)
40print("nth roots of unity:", w)
41print("Vandermonde matrix:")
42print(ft)
43
44a = vector([3,4,5,9])
45print("a:", a)
46
47# interpolate f_a(x)
48fa_coef = ft_inv * a
49print("fa_coef:", fa_coef)
50
51P.<x> = PolynomialRing(F)
52fa = P(list(fa_coef))
53print("f_a(x):", fa)
54
55# check that evaluating fa(x) at the roots of unity returns the expected values of a
56for i in range(len(a)):
57    assert fa(w[i]) == a[i]
58
59