This script provides a function to compute the Cooley-Tukey Radix-2 Decimation in Time (DIT) Fast Fourier Transform (FFT) on a given input vector.

fft.py

# !/usr/bin/env python3
# -*- coding: utf-8 -*-

# Author: Etienne Collin
# Date: 2023/12/09
# Email: [email protected]

"""
This script provides a function to compute the Cooley-Tukey Radix-2 Decimation in Time (DIT)
Fast Fourier Transform (FFT) on a given input vector. Additionally, the script may be run
through the terminal by giving it some arguments.

Functions:
- fft(x, inverse=False, omega=None, modulo=None, debug=False): Perform FFT or IFFT
  on an input vector.
- cli_parse(): Runs the fft function on a user-provided input vector via command-line arguments

Usage:
- Use the fft function to compute FFT or IFFT on a list:
    ```python
    result = fft(input_vector, inverse=False, omega=None, modulo=None, debug=False)
    ```
- Use the main function to run the script from the command line:
    ```bash
    python3 fft.py input [-h] [-d] [-i] [-m MODULO] [-o OMEGA]
    ```
"""

import math


def fft(x, inverse=False, omega=None, modulo=None, debug=False):
    """
    Perform the Cooley-Tukey Radix-2 Decimation in Time (DIT) Fast Fourier Transform (FFT)
    on a given input vector.

    Parameters:
    - x (list): Input vector for which FFT is computed. The length of the vector must
                be a power of 2.
    - inverse (bool, optional): If True, compute the Inverse FFT (IFFT) instead of the FFT.
                                Default is False.
    - omega (complex, optional): The primary N-th root of unity. If None, it will be
                                calculated based on the length of the input vector and
                                the inverse flag. Default is None.
    - modulo (int, optional): If provided, compute the FFT in a modular ring with the
                                specified modulus. Both omega and modulo must be either
                                None or not None. Default is None.
    - debug (bool, optional): If True, print intermediate results during FFT computation.
                                Default is False.

    Returns:
    - list: The result of the FFT or IFFT on the input vector.

    Raises:
    - ValueError: If the number of elements in the input vector is not a power of 2.
    - TypeError:
        - If omega and modulo are not both specified or both None when computing in
        a modular ring.
        - If omega and modulo are not both integers when computing in a modular ring.
        - If the input vector contains complex numbers when computing in a modular ring.

    Note:
    - The input vector length must be a power of 2.
    - If computing in a modular ring, both omega and modulo must be specified or None.
    - If computing the IFFT, the output is normalized by dividing each element by the
      length of the input vector.
    - If a modulus is provided, the result is taken modulo the specified value.

    Examples:
    >>> x = [1, 2, 3, 4, 4, 3, 2, 1]
    >>> omega = None
    >>> modulo = None
    >>> debug = False

    # Forward FFT
    >>> forward = fft(x, False, omega, modulo, debug)
    >>> print("Forward: ", forward)

    # Inverse FFT
    >>> inverse = fft(forward, True, omega, modulo, debug)
    >>> print("Inverse: ", inverse)
    """

    def fft_run(x, omega):
        # Base case of the recursion
        n = len(x)
        if n == 1:
            return x

        # If we compute the FFT in a modular ring, then the recursive calls
        # must be done with omega squared. Otherwise, the
        # recursive calls must be done with omega set to None
        next_omega = omega**2 if omega is not None else None
        y_even = fft_run(x[::2], next_omega)
        y_odd = fft_run(x[1::2], next_omega)

        # If we do not compute the FFT in a modular ring, then the omega must
        # be set to the correct value.
        if omega is None:
            omega = math.e ** ((-1) ** inverse * 2 * math.pi * 1j / n)

        # Compute the FFT
        y = [0j] * n
        for k in range(n // 2):
            y[k] = y_even[k] + omega**k * y_odd[k]
            y[k + n // 2] = y_even[k] - omega**k * y_odd[k]

        # Print the FFT if in debug mode
        if debug:
            print(y)

        # Return the FFT
        return y

    # Check if the number of elements in the input vector is a power of 2
    n = len(x)
    if not ((n & (n - 1) == 0) and n != 0):
        raise ValueError("The number of elements in the input vector must be a power of 2")

    # If we compute the FFT in a modular ring, then the omega and modulo
    # parameters must be specified and integers
    if bool(omega is None) ^ bool(modulo is None):
        raise TypeError("Omega and modulo must both be either None or not None")

    # If we compute the FFT in a modular ring, then the omega and modulo
    # parameters must be integers
    if bool(omega is not None) and bool(modulo is not None):
        if not (isinstance(omega, int) or not isinstance(modulo, int)):
            raise TypeError("Omega and modulo must both be integers")

    # If we compute the FFT in a modular ring, then the input vector must not
    # contain complex numbers
    if modulo is not None:
        if any(isinstance(i, complex) for i in x):
            raise TypeError("Input vector must not contain complex numbers when computing in a modular ring")

    # Compute the inverse omega if needed
    if inverse:
        omega = pow(omega, -1, modulo) if omega is not None else None

    # Run the FFT
    fft_result = fft_run(x, omega)

    # If we are computing the inverse FFT, then we need to divide
    # each element by the length of the input vector
    if inverse:
        fft_result = [(i * pow(len(x), -1, modulo)) for i in fft_result]

    # If the FFT is computed in a modular ring, then we need to apply the
    # modulo operation to each element
    if modulo is not None:
        fft_result = [i % modulo for i in fft_result]

    return fft_result


def cli_parse():
    """
    Runs the fft function on a user-provided input vector via command-line arguments.
    """

    import argparse
    import sys
    from argparse import ArgumentError

    # Disable traceback
    sys.tracebacklimit = 0

    # Create the parser
    parser = argparse.ArgumentParser(
        prog="Fast Fourier Transform (FFT)",
        description="Compute the FFT on a given input vector",
    )

    # Required positional argument
    parser.add_argument(
        "input",
        action="store",
        type=str,
        help="A string containing the input vector on which to compute the FFT. \
        The values are comma or space separated ints, floats and complex numbers)",
    )

    # Optional argument
    parser.add_argument("-d", "--debug", action="store_true", help="Print intermediate FFT results")
    parser.add_argument("-i", "--inverse", action="store_true", help="Compute the Inverse FFT (IFFT)")
    parser.add_argument(
        "-m",
        "--modulo",
        type=int,
        help="The modulo of the modular ring (int). Must be used with omega",
    )
    parser.add_argument(
        "-o",
        "--omega",
        type=int,
        help="The primary N-th root of unity (int). Must be used with modulo",
    )

    # Parse the arguments
    try:
        args = parser.parse_args()
    except ArgumentError:
        print("> Invalid arguments. Use -h or --help for help.")
        return

    # Prepare the input string and parse it into a list of ints/floats/complex numbers
    try:
        x = args.input.strip().replace(",", " ").replace("[", "").replace("]", "").split()
        x = [complex(elem).real if complex(elem).imag == 0 else complex(elem) for elem in x]
    except ValueError:
        print("> The input vector can only contain numbers")
        return

    # Get the parsed arguments
    inverse = args.inverse
    omega = args.omega
    modulo = args.modulo
    debug = args.debug

    # Run the FFT
    try:
        fft_result = fft(x, inverse, omega, modulo, debug)
    except (ValueError, TypeError) as e:
        print(">", type(e).__name__ + ":", e)
        return

    # Print the result
    prefix = "> FFTI result:" if inverse else "> FFT result:"
    print(prefix, fft_result)


# Only run the prompt if the file is executed directly
# Otherwise, the functions can be imported and used in another file
if __name__ == "__main__":
    cli_parse()