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.
93 lines
3.5 KiB
Racket
93 lines
3.5 KiB
Racket
#lang typed/racket/base
|
|
|
|
(require "../../parameters.rkt"
|
|
"../../base.rkt"
|
|
"../../flonum.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
|
|
fcarray-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)))
|