5a43f2c6bc
Cleaned up other docs in preparation for alpha-testing announcement Created `math/utils' module for stuff that doesn't go anywhere else (e.g. FFT scaling convention, max-math-threads parameters) Reduced the number of macros that expand to applications of `array-map' Added `flvector-sum', defined `flsum' in terms of it Reduced the number of pointwise `flvector', `flarray' and `fcarray' operations Reworked `inline-build-flvector' and `inline-flvector-map' to be faster and expand to less code in both typed and untyped Racket Redefined conversions like `list->flvector' in terms of for loops (can do it now that TR has working `for/flvector:', etc.)
144 lines
6.3 KiB
Racket
144 lines
6.3 KiB
Racket
#lang typed/racket/base
|
|
|
|
(require racket/flonum
|
|
racket/fixnum
|
|
racket/list
|
|
racket/future
|
|
"../../base.rkt"
|
|
"../../flonum.rkt"
|
|
"../parameters.rkt"
|
|
"../unsafe.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)))
|