f2dc2027f6
in the original GitHub fork: https://github.com/ntoronto/racket Some things about this are known to be broken (most egregious is that the array tests DO NOT RUN because of a problem in typed/rackunit), about half has no coverage in the tests, and half has no documentation. Fixes and docs are coming. This is committed now to allow others to find errors and inconsistency in the things that appear to be working, and to give the author a (rather incomplete) sense of closure.
145 lines
6.3 KiB
Racket
145 lines
6.3 KiB
Racket
#lang typed/racket/base
|
|
|
|
(require racket/flonum
|
|
racket/fixnum
|
|
racket/list
|
|
racket/future
|
|
"../../base.rkt"
|
|
"../../parameters.rkt"
|
|
"../unsafe.rkt"
|
|
"flvector.rkt"
|
|
"vector.rkt")
|
|
|
|
(provide vector-fft flvector-fft!
|
|
vector-inverse-fft flvector-inverse-fft!)
|
|
|
|
;; Fast Fourier Transform
|
|
|
|
(: init-d (-> Integer))
|
|
(define (init-d)
|
|
(define t (max 1 (max-math-threads)))
|
|
(exact-ceiling (/ (log t) (log 2))))
|
|
|
|
(: vector-fft (case-> ((Vectorof Float-Complex) -> (Vectorof Float-Complex))
|
|
((Vectorof Float-Complex) Integer -> (Vectorof Float-Complex))
|
|
((Vectorof Float-Complex) Integer Integer -> (Vectorof Float-Complex))))
|
|
(define vector-fft
|
|
(case-lambda
|
|
[(as) (vector-fft as 0 (vector-length as))]
|
|
[(as n) (vector-fft as 0 n)]
|
|
[(as start end)
|
|
(define n (vector-length as))
|
|
(define as-r (for/flvector: #:length n ([a (in-vector as)]) (real-part a)))
|
|
(define as-i (for/flvector: #:length n ([a (in-vector as)]) (imag-part a)))
|
|
(define bs-r (make-flvector n))
|
|
(define bs-i (make-flvector n))
|
|
(flvector-fft! as-r as-i start end bs-r bs-i 0)
|
|
(for/vector: #:length n ([x (in-flvector bs-r)]
|
|
[y (in-flvector bs-i)]) : Float-Complex
|
|
(make-rectangular x y))]))
|
|
|
|
(: flvector-fft! (FlVector FlVector Integer Integer FlVector FlVector Integer -> Void))
|
|
(define (flvector-fft! as-r as-i start end bs-r bs-i b-start)
|
|
(define len (min (flvector-length as-r) (flvector-length as-i)))
|
|
(cond
|
|
[(not (power-of-two? (- end start)))
|
|
(error 'vector-fft "expected power-of-two length; given length ~e" (- end start))]
|
|
[else
|
|
(let ([as-r (flvector-copy as-r start end)]
|
|
[as-i (flvector-copy as-i start end)])
|
|
(define n (flvector-length as-r))
|
|
(define xs-r (make-flvector n 0.0))
|
|
(define xs-i (make-flvector n 0.0))
|
|
(define a (real->double-flonum (first (dft-convention))))
|
|
(define b (real->double-flonum (second (dft-convention))))
|
|
(vector-fft/depth as-r as-i xs-r xs-i b n 0 (init-d))
|
|
(define c (flexpt (->fl n) (* 0.5 (- 1.0 a))))
|
|
(let loop ([#{j : Nonnegative-Fixnum} 0])
|
|
(when (j . < . n)
|
|
(define r (/ (unsafe-flvector-ref as-r j) c))
|
|
(define i (/ (unsafe-flvector-ref as-i j) c))
|
|
(let ([j (+ j b-start)])
|
|
(unsafe-flvector-set! bs-r j r)
|
|
(unsafe-flvector-set! bs-i j i))
|
|
(loop (+ j 1)))))]))
|
|
|
|
(: decimate-in-time! (FlVector FlVector FlVector FlVector Index Index -> Void))
|
|
(define (decimate-in-time! as-r as-i xs-r xs-i n/2 start)
|
|
(let loop ([#{i : Nonnegative-Fixnum} 0])
|
|
(when (i . < . n/2)
|
|
(define si (+ start i))
|
|
(define si2 (unsafe-fx+ si i))
|
|
(define si21 (unsafe-fx+ si2 1))
|
|
(define sin2 (unsafe-fx+ si n/2))
|
|
(unsafe-flvector-set! xs-r si (unsafe-flvector-ref as-r si2))
|
|
(unsafe-flvector-set! xs-i si (unsafe-flvector-ref as-i si2))
|
|
(unsafe-flvector-set! xs-r sin2 (unsafe-flvector-ref as-r si21))
|
|
(unsafe-flvector-set! xs-i sin2 (unsafe-flvector-ref as-i si21))
|
|
(loop (+ i 1)))))
|
|
|
|
(: twiddle-factor! (FlVector FlVector Float Index Index -> Void))
|
|
(define (twiddle-factor! xs-r xs-i b n/2 start)
|
|
(define c (/ (* b pi 0.0+1.0i) (->fl n/2)))
|
|
(let loop ([#{k : Nonnegative-Fixnum} 0])
|
|
(when (k . < . n/2)
|
|
(define k-start (+ k start))
|
|
(define res (* (make-rectangular (unsafe-flvector-ref xs-r k-start)
|
|
(unsafe-flvector-ref xs-i k-start))
|
|
(exp (* c (->fl k)))))
|
|
(unsafe-flvector-set! xs-r k-start (real-part res))
|
|
(unsafe-flvector-set! xs-i k-start (imag-part res))
|
|
(loop (+ k 1)))))
|
|
|
|
(: combine! (FlVector FlVector FlVector FlVector Index Index -> Void))
|
|
(define (combine! as-r as-i xs-r xs-i n/2 start)
|
|
(let loop ([#{k : Nonnegative-Fixnum} 0])
|
|
(when (k . < . n/2)
|
|
(define sk (+ start k))
|
|
(define sk2 (unsafe-fx+ sk n/2))
|
|
(define br (unsafe-flvector-ref xs-r sk))
|
|
(define bi (unsafe-flvector-ref xs-i sk))
|
|
(define cr (unsafe-flvector-ref xs-r sk2))
|
|
(define ci (unsafe-flvector-ref xs-i sk2))
|
|
(unsafe-flvector-set! as-r sk2 (- br cr))
|
|
(unsafe-flvector-set! as-i sk2 (- bi ci))
|
|
(unsafe-flvector-set! as-r sk (+ br cr))
|
|
(unsafe-flvector-set! as-i sk (+ bi ci))
|
|
(loop (+ k 1)))))
|
|
|
|
(: vector-fft/depth (FlVector FlVector FlVector FlVector Float Index Index Integer -> Void))
|
|
(define (vector-fft/depth as-r as-i xs-r xs-i b n start d)
|
|
(unless (= n 1)
|
|
(define n/2 (quotient n 2))
|
|
(decimate-in-time! as-r as-i xs-r xs-i n/2 start)
|
|
(cond [(= d 0) (define start+n/2 (assert (+ start n/2) index?))
|
|
(vector-fft/depth xs-r xs-i as-r as-i b n/2 start 0)
|
|
(vector-fft/depth xs-r xs-i as-r as-i b n/2 start+n/2 0)
|
|
(twiddle-factor! xs-r xs-i b n/2 start+n/2)]
|
|
[else
|
|
(define bs (future (λ () (vector-fft/depth xs-r xs-i as-r as-i b n/2 start (- d 1)))))
|
|
(define cs (future (λ ()
|
|
(define start+n/2 (assert (+ start n/2) index?))
|
|
(vector-fft/depth xs-r xs-i as-r as-i b n/2 start+n/2 (- d 1))
|
|
(twiddle-factor! xs-r xs-i b n/2 start+n/2))))
|
|
(touch bs)
|
|
(touch cs)])
|
|
(combine! as-r as-i xs-r xs-i n/2 start)))
|
|
|
|
;; ---------------------------------------------------------------------------------------------------
|
|
;; Inverse Fast Fourier Transform
|
|
|
|
(: vector-inverse-fft (case-> ((Vectorof Float-Complex) -> (Vectorof Float-Complex))
|
|
((Vectorof Float-Complex) Integer -> (Vectorof Float-Complex))
|
|
((Vectorof Float-Complex) Integer Integer -> (Vectorof Float-Complex))))
|
|
(define vector-inverse-fft
|
|
(case-lambda
|
|
[(as) (vector-fft as 0 (vector-length as))]
|
|
[(as n) (vector-fft as 0 n)]
|
|
[(as start end) (parameterize ([dft-convention (dft-inverse-convention)])
|
|
(vector-fft as start end))]))
|
|
|
|
(: flvector-inverse-fft! (FlVector FlVector Integer Integer FlVector FlVector Integer -> Void))
|
|
(define (flvector-inverse-fft! as-r as-i start end bs-r bs-i b-start)
|
|
(parameterize ([dft-convention (dft-inverse-convention)])
|
|
(flvector-fft! as-r as-i start end bs-r bs-i b-start)))
|