FastFourierTransforms.sage

# 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]