Negacyclic polymul based on tangent fft

tangent_fft.py

import numpy as np
import math


def primitive_nth_root(n):
  """Return a primitive nth root of unity."""
  return math.cos(2 * math.pi / n) + 1.0j * math.sin(2 * math.pi / n)


def poly_mul(a, b):
  """Computes a poly multiplication mod (X^N + 1) where N = len(a).

  Uses the Tangent FFT idea described in
  https://nufhe.readthedocs.io/en/latest/implementation_details.html

  Args:
    a: a polynomial in degree-increasing order of coefficients
    b: a polynomial in degree-increasing order of coefficients

  Returns:
    The product polynomial
  """
  n = a.shape[0]
  primitive_root = primitive_nth_root(2 * n)
  root_powers = primitive_root**np.arange(n // 2)

  # change - to +
  # a_preprocessed = (a[:n // 2] - 1j * a[n // 2:]) * root_powers
  # b_preprocessed = (b[:n // 2] - 1j * b[n // 2:]) * root_powers
  a_preprocessed = (a[:n // 2] + 1j * a[n // 2:]) * root_powers
  b_preprocessed = (b[:n // 2] + 1j * b[n // 2:]) * root_powers

  a_ft = np.fft.fft(a_preprocessed)
  b_ft = np.fft.fft(b_preprocessed)
  prod = a_ft * b_ft

  # remove conjugation
  # ifft_prod = np.conj(np.fft.ifft(prod))
  ifft_prod = np.fft.ifft(prod)

  # invert twist instead of applying a second twist
  # ifft_rotated = ifft_prod * root_powers
  ifft_rotated = ifft_prod * primitive_root**np.arange(0, -n // 2, -1)

  first_half = np.real(ifft_rotated)
  second_half = np.imag(ifft_rotated)
  return np.round(np.concatenate([first_half, second_half])).astype(a.dtype)


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.array([1, 2, 3, 4, 5, 6, 7, 8])
  # b = np.array([2, 3, 4, 5, 6, 7, 8, 9])
  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)}")