The Fast Fourier Transform in Petalisp

fft.lisp

(use-package :petalisp)

(defun fft (x)
  (let ((n (shape-size (shape x))))
    (assert (= (shape-rank (shape x)) 1))
    (assert (zerop (logand n (1- n))) ()
            "The length is not a power of 2")
    (if (<= n 1)
        x
        (let* ((even (fft (reshape x (~ 0 2 (1- n))
                                     (τ (x) ((/ x 2))))))
               (odd  (fft (reshape x (~ 1 2 (1- n))
                                     (τ (x) ((/ (1- x) 2))))))
               (odd* (α #'* odd
                        (α #'cis
                           (α #'*
                              (/ pi (shape-size (shape x)) -0.5)
                              (indices odd))))))
          (fuse (α #'+ even odd*)
                (reshape (α #'- even odd*)
                         (τ (i) ((+ i (shape-size (shape odd)))))))))))

That's pretty much the way to go in Petalisp. Here is how I would have written it:

(defun fft (x)
  (trivia:ematch (shape x)
    ((~ _) x)
    ((~ start end)
     (let* ((n (the power-of-two (- end start -1)))
            (even (fft (collapse (reshape x (~ 0 2 (1- n))))))
            (odd  (fft (collapse (reshape x (~ 1 2 (1- n))))))
            (odd* (α #'* odd
                     (α #'cis
                        (α #'*
                           (/ pi n -0.5)
                           (indices odd))))))
       (fuse (α #'+ even odd*)
             (reshape (α #'- even odd*)
                      (τ (i) ((+ i (shape-size (shape odd)))))))))))

The only thing I changed is that I use trivia:ematch to destructure shapes, the collapse function instead of τ, and a custom type for powers of two.

The last expression bothers me. That is a pattern that I have used quite often already. Maybe I should introduce a stack function to place arrays next to each other for a supplied axis. Then one could replace the last three lines with

(stack 0 (α #'+ even odd*) (α #'- even odd*))

@nodefunallowed What do you think?

Sure, using pattern matching, collapse and the type is definitely nicer.

The code I had based this off used an append (and was pretty silly for using lists), and I was thinking about making a lazy array concatenate to make the last expression's purpose more obvious. stack is probably a better name if it can go in any dimension.