Naive negacyclic FFT

fft.py

import numpy as np


def poly_mul(a, b):
  a_preprocessed = np.concatenate([a, -a], dtype=np.complex64)
  b_preprocessed = np.concatenate([b, -b], dtype=np.complex64)
  a_ft = np.fft.fft(a_preprocessed)
  b_ft = np.fft.fft(b_preprocessed)
  prod = a_ft * b_ft
  np_fft_mul = np.round(0.5 * np.real(np.fft.ifft(prod))).astype(np.int64)
  return np_fft_mul[:a.shape[0]]


def _np_polymul(poly1, poly2):
  # poly_mod represents the polynomial to divide by: x^N + 1, N = len(a)
  poly_mod = np.zeros(len(poly1) + 1, np.uint32)
  poly_mod[0] = 1
  poly_mod[len(poly1)] = 1

  # Reversing the list order because numpy polymul interprets the polynomial
  # with higher-order coefficients first, whereas our code does the opposite
  np_mul = np.polymul(list(reversed(poly1)), list(reversed(poly2)))
  (_, np_poly_mod) = np.polydiv(np_mul, poly_mod)
  np_pad = np.pad(
      np_poly_mod, (len(poly1) - len(np_poly_mod), 0),
      "constant",
      constant_values=(0, 0))
  return np.array(list(reversed(np_pad)), dtype=int)


if __name__ == "__main__":
  a = np.random.randint(low=0, high=2**16 - 1, size=(512,))
  b = np.random.randint(low=0, high=2**16 - 1, size=(512,))

  output = poly_mul(a, b)
  expected = _np_polymul(a, b)

  abs_diff = np.abs(output - expected)

  # print(f"output=\t\t{output}")
  # print(f"expected=\t{expected}")
  print(f"max_abs_diff=\t{np.max(abs_diff)}")