[Python] Radix-4 NTT Python Code Implementation Debugging

I am currently implementing a radix-4 NTT algorithm based on the pseudocode I have written Python code for it.For n = 4, with modulus q = 3329, the 2n-th root of unity is ψ = 40, and the 4th root of unity is 1600.My implementation gives correct results compared to the direct NTT definition for n=4.

PseudoCode For Radix-4:https://postimg.cc/yDC25VR1 PseudoCode For definition:https://postimg.cc/3WPq4KnX

However, I am running into issues when scaling to larger sizes: For n = 16, modulus q = 3329, I used ψ = 193 as the 2n-th root of unity (this follows Kyber’s specification). The 4th root of unity I used is ω = 1600.

I tested the code by running the NTT on the input [1, 2, 3, …, 16].The result from the radix-4 pseudocode is not matching the output obtained from the definition of NTT.

Here is the input and the output I got for reference: Input: [1, 2, 3, …, 16] Output (radix-4 algo):Output: [1618, 3118, 2098, 2565, 1159, 1073, 246, 2062, 1288, 2094, 654, 1204, 145, 535, 2752, 708]

Code is given as follows:



# Implementation of radix-4:
import math


psi = 193
q = 3329
n = 16
w4 = 1600
def mod_pow(base, exp, mod=q):
return pow(base, exp, mod)

def bit_reverse(val, bits):
"""rev bits"""
reversed_val = 0
for i in range(bits):
if (val >> i) & 1:
reversed_val |= 1 << (bits - 1 - i)
return reversed_val

#ROMs
wa1_ROM = []
wa2_ROM = []
wa3_ROM = []

log4_n = int(math.log(n, 4))

#Precompute
for p in range(log4_n - 1, -1, -1): # Stage loop
J = 4 ** p # Spacing
omega_m = mod_pow(psi, J,q) # Base twiddle factor for stage p

for k in range(n // (4 * J)): # Group loop
i = 2 * bit_reverse(k, int(math.log(n // (4 * J), 2))) + 1

val1 = mod_pow(omega_m, i, q)
val2 = mod_pow(omega_m, 2 * i, q)
val3 = mod_pow(omega_m, 3 * i, q)

wa1_ROM.append(val1)
wa2_ROM.append(val2)
wa3_ROM.append(val3)

print(f"Stage p={p}, Spacing J={J}, Group k={k}, i={i}")
print(f" ω1 = {val1} (ψ^{J * i} mod {q})")
print(f" ω2 = {val2} (ψ^{J * 2 * i} mod {q})")
print(f" ω3 = {val3} (ψ^{J * 3 * i} mod {q})")
print("-" * 50)

print("wa1_ROM:", wa1_ROM)
print("wa2_ROM:", wa2_ROM)
print("wa3_ROM:", wa3_ROM)


#Step 2: Radix-4 NTT

def radix4_ntt(a):
t=0
for stage_num, p in enumerate(range(log4_n - 1, -1, -1), start=1):

J = 4 ** p
print(f"\n===== Stage {stage_num} (p={p}, J={J}) =====")
for k in range(n // (4 * J)):
w1 = wa1_ROM[t]
w2 = wa2_ROM[t]
w3 = wa3_ROM[t]
print(f"\n Group k={k}, ROM index t={t}: w1={w1}, w2={w2}, w3={w3}")
t += 1
for j in range(J):
T0 = (a[4*k*J + j] + a[4*k*J + j + 2*J] * w2) % q
T1 = (a[4*k*J + j] - a[4*k*J + j + 2*J] * w2) % q
T2 = (a[4*k*J + j + J] * w1 + a[4*k*J + j + 3*J] * w3) % q
T3 = (a[4*k*J + j + J] * w1 - a[4*k*J + j + 3*J] * w3) % q

print(f" j={j}:")
print(f" T0 = (a[{4*k*J + j}] + a[{4*k*J + j + 2*J}] * {w2}) % {q} = {T0}")
print(f" T1 = (a[{4*k*J + j}] - a[{4*k*J + j + 2*J}] * {w2}) % {q} = {T1}")
print(f" T2 = (a[{4*k*J + j + J}] * {w1} + a[{4*k*J + j + 3*J}] * {w3}) % {q} = {T2}")
print(f" T3 = (a[{4*k*J + j + J}] * {w1} - a[{4*k*J + j + 3*J}] * {w3}) % {q} = {T3}")

a[4*k*J + j] = (T0 + T2) % q
a[4*k*J + j + J] = (T1 + T3 * w4) % q
a[4*k*J + j + 2*J] = (T0 - T2) % q
a[4*k*J + j + 3*J] = (T1 - T3 * w4) % q

print(f" a[{4*k*J + j}] = (T0 + T2) % {q} = {a[4*k*J + j]}")
print(f" a[{4*k*J + j + J}] = (T1 + T3 * {w4}) % {q} = {a[4*k*J + j + J]}")
print(f" a[{4*k*J + j + 2*J}] = (T0 - T2) % {q} = {a[4*k*J + j + 2*J]}")
print(f" a[{4*k*J + j + 3*J}] = (T1 - T3 * {w4}) % {q} = {a[4*k*J + j + 3*J]}")

return a

a = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16]
print("Input:", a)
result = radix4_ntt(a[:])
print("\nOutput:", result)


Whereas using the definition For the same input I am getting output as NTT: [1618, 2559, 2893, 1509, 2565, 2836, 2084, 3317, 2098, 1098, 1242, 2031, 3118, 976, 2350, 1012] By taking 2nth root of unity as 193 and nth root of unity as 630.

NTT Defination Code:

def mod_pow(base, exp, mod):
"""Modular exponentiation"""
return pow(base, exp, mod)

def ntt_transform(a, N, q, phi, omega):
"""
Computes A_i = sum_{j=0}^{N-1} a_j * phi^(j) * omega^(i*j) mod q
for i = 0,1,...,N-1

Parameters:
a : input array of length N
N : transform length
q : modulus
phi : primitive 2N-th root of unity modulo q
omega : primitive N-th root of unity modulo q
Returns:
List of transformed values A of length N
"""
A = []
for i in range(N):
total = 0
for j in range(N):
term = (a[j] * mod_pow(phi, j, q) * mod_pow(omega, i * j, q)) % q
total = (total + term) % q
A.append(total)
return A


# Example usage:
if __name__ == "__main__":
# Example parameters (you can change these)
N = 16
q = 3329
phi = 193 # example value for φ_{2N}
omega = 630 # example value for ω_N
a = list(range(1,N+1)) # input array

result = ntt_transform(a, N, q, phi, omega)
print("NTT:", result)





My questions are:

1. My computation of the roots of unity (2nth root of unity = 193, 4th root of unity = 1600 and nth root of unity = 630) is correct for Kyber NTT with n = 16 and q = 3329.
2. Is my python code correct for radix-4 NTT is correct, or if I’m misinterpreting the pseudocode?

Any guidance or verification would be greatly appreciated! Thanks in advance.

Continue reading...