Skip to content

Instantly share code, notes, and snippets.

@asanso
Last active Nov 26, 2021
Embed
What would you like to do?
# p and values taken from the example https://vitalik.ca/general/2019/05/12/fft.html
# based on the work https://solvable.group/posts/ecfft/
p =337
F = GF(p)
log_n = 3
n = 2^log_n
S = [1, 148, 336, 189]
S_prime = [85, 111,252, 226]
cosets = {}
cosets[0] = [1, 85, 148, 111, 336, 252, 189, 226]
cosets[1] = [1,148,336,189]
cosets[2] = [1,336]
def precompute(log_n, S, S_prime):
Ss = {}
Ss_prime = {}
matrices = {}
inverse_matrices = {}
for i in range(log_n, -1, -1):
n = 1 << i
nn = n // 2
matrices[n] = []
inverse_matrices[n] = []
R2.<X> = F[]
for j in range(nn):
s0, s1 = S[j], S[j + nn]
M = Matrix(F, [[1,s0],[1, s1]])
inverse_matrices[n].append(M.inverse())
s0, s1 = S_prime[j], S_prime[j + nn]
M = Matrix(F, [[1,s0],[1, s1]])
matrices[n].append(M)
S = [(x^2 %p) for x in S[:nn]]
S_prime = [(x^2 % p) for x in S_prime[:nn]]
return matrices, inverse_matrices
# Precompute the data needed to compute EXTEND_S,S'
matrices, inverse_matrices, = precompute(log_n-1, S, S_prime)
def extend(P_evals):
n = len(P_evals)
nn = n // 2
if n == 1:
return P_evals
P0_evals = []
P1_evals = []
for j in range(nn):
s0, s1 = S[j], S[j + nn]
y0, y1 = P_evals[j], P_evals[j + nn]
Mi = inverse_matrices[n][j]
p0, p1 = Mi * vector([y0, y1])
P0_evals.append(p0)
P1_evals.append(p1)
P0_evals_prime = extend(P0_evals)
P1_evals_prime = extend(P1_evals)
ansL = []
ansR = []
for M, p0, p1 in zip(matrices[n], P0_evals_prime, P1_evals_prime):
v = M * vector([p0, p1])
ansL.append(v[0])
ansR.append(v[1])
return ansL + ansR
def enter(P_coeffs):
n = len(P_coeffs)
nn = n // 2
if len(P_coeffs) == 1:
return P_coeffs
low = P_coeffs[:nn]
high = P_coeffs[nn:]
low_evals = enter(low)
high_evals = enter(high)
low_evals_prime = extend(low_evals)
high_evals_prime = extend(high_evals)
res = []
coset = cosets[log_n -log(n,2)]
for i in range(nn):
res.append((low_evals[i] + coset[2 * i]**nn * high_evals[i])%p)
res.append((low_evals_prime[i] + coset[2 * i+ 1]**nn * high_evals_prime[i]) %p)
return res
# Generate the same polynomial as the example in https://vitalik.ca/general/2019/05/12/fft.html
R1.<X> = F[]
P = 6*X^7 + 2*X^6 + 9*X^5 + 5*X^4 + X^3+4*X^2+X+3
result = enter(P.coefficients())
assert result == [31, 70, 109, 74, 334, 181, 232, 4]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment