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.)
92 lines
3.4 KiB
Racket
92 lines
3.4 KiB
Racket
#lang typed/racket/base
|
|
|
|
(require "../../base.rkt"
|
|
"../../flonum.rkt"
|
|
"../parameters.rkt"
|
|
"../unsafe.rkt"
|
|
"../vector/vector-fft.rkt"
|
|
"fcarray-struct.rkt"
|
|
"array-struct.rkt"
|
|
"array-transform.rkt"
|
|
"for-each.rkt"
|
|
"utils.rkt")
|
|
|
|
(provide array-axis-fft
|
|
array-fft
|
|
array-axis-inverse-fft
|
|
array-inverse-fft)
|
|
|
|
;; Fast Fourier Transform
|
|
|
|
(: fcarray-last-axis-fft (FCArray -> FCArray))
|
|
(define (fcarray-last-axis-fft arr)
|
|
(define ds (array-shape arr))
|
|
(define dims (vector-length ds))
|
|
(define k (- dims 1))
|
|
(cond
|
|
[(not (index? k))
|
|
(raise-argument-error 'fcarray-last-axis-fft "FCArray with at least one axis" arr)]
|
|
[else
|
|
(define xs (fcarray-real-data arr))
|
|
(define ys (fcarray-imag-data arr))
|
|
(define dk (unsafe-vector-ref ds k))
|
|
(define n (array-size arr))
|
|
(define new-xs (make-flvector n))
|
|
(define new-ys (make-flvector n))
|
|
(for-each-array+data-index
|
|
(unsafe-vector-remove ds k)
|
|
(λ (js j)
|
|
(define old-j (unsafe-fx* j dk))
|
|
(flvector-fft! xs ys old-j (unsafe-fx+ old-j dk) new-xs new-ys old-j)))
|
|
(unsafe-fcarray ds new-xs new-ys)]))
|
|
|
|
(: array-axis-fft ((Array Number) Integer -> (Array Float-Complex)))
|
|
(define (array-axis-fft arr k)
|
|
(define ds (array-shape arr))
|
|
(define dims (vector-length ds))
|
|
(cond [(= dims 0)
|
|
(raise-argument-error 'array-axis-fft "Array with at least one axis" 0 arr k)]
|
|
[(or (0 . > . k) (k . >= . dims))
|
|
(raise-argument-error 'array-axis-fft (format "Index less than ~a" dims) 1 arr k)]
|
|
[(= k (- dims 1))
|
|
(fcarray-last-axis-fft (array->fcarray arr))]
|
|
[else
|
|
(array-axis-swap (fcarray-last-axis-fft (array->fcarray (array-axis-swap arr k (- dims 1))))
|
|
k (- dims 1))]))
|
|
|
|
(: fcarray-fft (FCArray -> FCArray))
|
|
(define (fcarray-fft arr)
|
|
(define dims (array-dims arr))
|
|
(cond [(zero? dims) (raise-argument-error 'fcarray-fft "FCArray with at least one axis" arr)]
|
|
[(not (andmap power-of-two? (vector->list (array-shape arr))))
|
|
(raise-argument-error 'fcarray-fft "FCArray with power-of-two shape" arr)]
|
|
[else
|
|
(define dims-1 (- dims 1))
|
|
(cond [(zero? dims-1) (fcarray-last-axis-fft arr)]
|
|
[else
|
|
(let loop ([#{k : Positive-Fixnum} 1] [arr (array-axis-fft arr 0)])
|
|
(cond [(k . < . dims-1) (loop (+ k 1) (array-axis-fft arr k))]
|
|
[else (fcarray-last-axis-fft (array->fcarray arr))]))])]))
|
|
|
|
(: array-fft ((Array Number) -> FCArray))
|
|
(define (array-fft arr)
|
|
(define dims (array-dims arr))
|
|
(cond [(= dims 0) (raise-argument-error 'array-fft "Array with at least one axis" arr)]
|
|
[(not (andmap power-of-two? (vector->list (array-shape arr))))
|
|
(raise-argument-error 'array-fft "Array with power-of-two shape" arr)]
|
|
[else
|
|
(fcarray-fft (array->fcarray arr))]))
|
|
|
|
;; ---------------------------------------------------------------------------------------------------
|
|
;; Inverse Fast Fourier Transform
|
|
|
|
(: array-axis-inverse-fft ((Array Number) Integer -> (Array Float-Complex)))
|
|
(define (array-axis-inverse-fft arr k)
|
|
(parameterize ([dft-convention (dft-inverse-convention)])
|
|
(array-axis-fft arr k)))
|
|
|
|
(: array-inverse-fft ((Array Number) -> (Array Float-Complex)))
|
|
(define (array-inverse-fft arr)
|
|
(parameterize ([dft-convention (dft-inverse-convention)])
|
|
(array-fft arr)))
|