Skip to content

Instantly share code, notes, and snippets.

@johnhw
Created February 8, 2024 11:57
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save johnhw/b3a6931f5de7f17abf96ece13ce4d0d5 to your computer and use it in GitHub Desktop.
Save johnhw/b3a6931f5de7f17abf96ece13ce4d0d5 to your computer and use it in GitHub Desktop.
FFT based multiplication example in NumPy
from numpy.fft import rfft, irfft
import numpy as np
def int_to_vec(x, base=10):
"""Convert an integer to a list of digits in a given base."""
return [int(s, base) for s in str(x)]
def vec_to_int(v, base=10):
"""Convert a list of digits to an integer in a given base;
values of the list can be any integer, carries will
be propagated to fit the base."""
n, m, carry = 0, 1, 0
for x in reversed(v):
y = x + carry
n += m * (y % base)
carry = y // base
m *= base
return n + m * carry
def fft_mul(a, b):
"""Perform multplication by convolution in the frequency domain."""
n = len(a) + len(b)
n = n + (n % 2) # make it even length
a = [0] * (n - len(a)) + a
b = [0] * (n - len(b)) + b
return [int(round(v)) for v in irfft(rfft(a) * rfft(b))[:-1]]
def mul(a,b,base=10):
"""Multiply two integers in a given base using FFT-based convolution."""
return vec_to_int(fft_mul(int_to_vec(a, base), int_to_vec(b, base)), base)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment