You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

147 lines
3.3 KiB

# Primitive Root of Unity
def get_primitive_root_of_unity(F, n):
# using the method described by Thomas Pornin in
# https://crypto.stackexchange.com/a/63616
q = F.order()
for k in range(q):
if k==0:
continue
g = F(k)
# g = F.random_element()
if g==0:
continue
w = g ^ ((q-1)/n)
if w^(n/2) != 1:
return g, w
# Roots of Unity
def get_nth_roots_of_unity(n, primitive_w):
w = [0]*n
for i in range(n):
w[i] = primitive_w^i
return w
# fft (Fast Fourier Transform) returns:
# - nth roots of unity
# - Vandermonde matrix for the nth roots of unity
# - Inverse Vandermonde matrix
def fft(F, n):
g, primitive_w = get_primitive_root_of_unity(F, n)
w = get_nth_roots_of_unity(n, primitive_w)
ft = matrix(F, n)
for j in range(n):
row = []
for k in range(n):
row.append(primitive_w^(j*k))
ft.set_row(j, row)
ft_inv = ft^-1
return w, ft, ft_inv
# Fast polynomial multiplicaton using FFT
def poly_mul(fa, fb, F, n):
w, ft, ft_inv = fft(F, n)
# compute evaluation points from polynomials fa & fb at the roots of unity
a_evals = []
b_evals = []
for i in range(n):
a_evals.append(fa(w[i]))
b_evals.append(fb(w[i]))
# multiply elements in a_evals by b_evals
c_evals = map(operator.mul, a_evals, b_evals)
c_evals = vector(c_evals)
# using FFT, convert the c_evals into fc(x)
fc_coef = c_evals*ft_inv
fc2=P(fc_coef.list())
return fc2, c_evals
# Tests
#####
# Roots of Unity test:
q = 17
F = GF(q)
n = 4
g, primitive_w = get_primitive_root_of_unity(F, n)
print("generator:", g)
print("primitive_w:", primitive_w)
w = get_nth_roots_of_unity(n, primitive_w)
print(f"{n}th roots of unity: {w}")
assert w == [1, 13, 16, 4]
#####
# FFT test:
def isprime(num):
for n in range(2,int(num^1/2)+1):
if num%n==0:
return False
return True
# list valid values for q
for i in range(20):
if isprime(8*i+1):
print("q =", 8*i+1)
q = 41
F = GF(q)
n = 4
# q needs to be a prime, s.t. q-1 is divisible by n
assert (q-1)%n==0
print("q =", q, "n = ", n)
# ft: Vandermonde matrix for the nth roots of unity
w, ft, ft_inv = fft(F, n)
print("nth roots of unity:", w)
print("Vandermonde matrix:")
print(ft)
fa_eval = vector([3,4,5,9])
print("fa_eval:", fa_eval)
# interpolate f_a(x)
fa_coef = ft_inv * fa_eval
print("fa_coef:", fa_coef)
P.<x> = PolynomialRing(F)
fa = P(list(fa_coef))
print("f_a(x):", fa)
# check that evaluating fa(x) at the roots of unity returns the expected values of fa_eval
for i in range(len(fa_eval)):
assert fa(w[i]) == fa_eval[i]
# go from coefficient form to evaluation form
fa_eval2 = ft * fa_coef
print("fa_eval'", fa_eval)
assert fa_eval2 == fa_eval
# Fast polynomial multiplicaton using FFT
print("\n---------")
print("---Fast polynomial multiplication using FFT")
n = 8
# q needs to be a prime, s.t. q-1 is divisible by n
assert (q-1)%n==0
print("q =", q, "n = ", n)
fa=P([1,2,3,4])
fb=P([1,2,3,4])
fc_expected = fa*fb
print("fc expected result:", fc_expected) # expected result
print("fc expected coef", fc_expected.coefficients())
fc, c_evals = poly_mul(fa, fb, F, n)
print("c_evals=(a_evals*b_evals)=", c_evals)
print("fc:", fc)
assert fc_expected == fc