1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
|
#lang typed/racket/base
(require "../../base.rkt"
racket/flonum
"../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
(parameterize ([array-strictness #f])
(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)))
|
