github.com/gopherd/gonum@v0.0.4/dsp/fourier/fourier.go (about)

     1  // Copyright ©2018 The Gonum Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package fourier
     6  
     7  import "github.com/gopherd/gonum/dsp/fourier/internal/fftpack"
     8  
     9  // FFT implements Fast Fourier Transform and its inverse for real sequences.
    10  type FFT struct {
    11  	work []float64
    12  	ifac [15]int
    13  
    14  	// real temporarily store complex data as
    15  	// pairs of real values to allow passing to
    16  	// the backing code. The length of real
    17  	// must always be half the length of work.
    18  	real []float64
    19  }
    20  
    21  // NewFFT returns an FFT initialized for work on sequences of length n.
    22  func NewFFT(n int) *FFT {
    23  	var t FFT
    24  	t.Reset(n)
    25  	return &t
    26  }
    27  
    28  // Len returns the length of the acceptable input.
    29  func (t *FFT) Len() int { return len(t.real) }
    30  
    31  // Reset reinitializes the FFT for work on sequences of length n.
    32  func (t *FFT) Reset(n int) {
    33  	if 2*n <= cap(t.work) {
    34  		t.work = t.work[:2*n]
    35  		t.real = t.real[:n]
    36  	} else {
    37  		t.work = make([]float64, 2*n)
    38  		t.real = make([]float64, n)
    39  	}
    40  	fftpack.Rffti(n, t.work, t.ifac[:])
    41  }
    42  
    43  // Coefficients computes the Fourier coefficients of the input sequence,
    44  // converting the time series in seq into the frequency spectrum, placing
    45  // the result in dst and returning it. This transform is unnormalized; a
    46  // call to Coefficients followed by a call of Sequence will multiply the
    47  // input sequence by the length of the sequence.
    48  //
    49  // If the length of seq is not t.Len(), Coefficients will panic.
    50  // If dst is nil, a new slice is allocated and returned. If dst is not nil and
    51  // the length of dst does not equal t.Len()/2+1, Coefficients will panic.
    52  func (t *FFT) Coefficients(dst []complex128, seq []float64) []complex128 {
    53  	if len(seq) != t.Len() {
    54  		panic("fourier: sequence length mismatch")
    55  	}
    56  	if dst == nil {
    57  		dst = make([]complex128, t.Len()/2+1)
    58  	} else if len(dst) != t.Len()/2+1 {
    59  		panic("fourier: destination length mismatch")
    60  	}
    61  	copy(t.real, seq)
    62  	fftpack.Rfftf(len(t.real), t.real, t.work, t.ifac[:])
    63  	dst[0] = complex(t.real[0], 0)
    64  	if len(seq) < 2 {
    65  		return dst
    66  	}
    67  	if len(seq)%2 == 1 {
    68  		dst[len(dst)-1] = complex(t.real[len(t.real)-2], t.real[len(t.real)-1])
    69  	} else {
    70  		dst[len(dst)-1] = complex(t.real[len(t.real)-1], 0)
    71  	}
    72  	for i := 1; i < len(dst)-1; i++ {
    73  		dst[i] = complex(t.real[2*i-1], t.real[2*i])
    74  	}
    75  	return dst
    76  }
    77  
    78  // Sequence computes the real periodic sequence from the Fourier coefficients,
    79  // converting the frequency spectrum in coeff into a time series, placing the
    80  // result in dst and returning it. This transform is unnormalized; a call to
    81  // Coefficients followed by a call of Sequence will multiply the input sequence
    82  // by the length of the sequence.
    83  //
    84  // If the length of coeff is not t.Len()/2+1, Sequence will panic.
    85  // If dst is nil, a new slice is allocated and returned. If dst is not nil and
    86  // the length of dst does not equal the length of coeff, Sequence will panic.
    87  func (t *FFT) Sequence(dst []float64, coeff []complex128) []float64 {
    88  	if len(coeff) != t.Len()/2+1 {
    89  		panic("fourier: coefficients length mismatch")
    90  	}
    91  	if dst == nil {
    92  		dst = make([]float64, t.Len())
    93  	} else if len(dst) != t.Len() {
    94  		panic("fourier: destination length mismatch")
    95  	}
    96  	dst[0] = real(coeff[0])
    97  	if len(dst) < 2 {
    98  		return dst
    99  	}
   100  	nf := coeff[len(coeff)-1]
   101  	if len(dst)%2 == 1 {
   102  		dst[len(dst)-2] = real(nf)
   103  		dst[len(dst)-1] = imag(nf)
   104  	} else {
   105  		dst[len(dst)-1] = real(nf)
   106  	}
   107  
   108  	for i, cv := range coeff[1 : len(coeff)-1] {
   109  		dst[2*i+1] = real(cv)
   110  		dst[2*i+2] = imag(cv)
   111  	}
   112  	fftpack.Rfftb(len(dst), dst, t.work, t.ifac[:])
   113  	return dst
   114  }
   115  
   116  // Freq returns the relative frequency center for coefficient i.
   117  // Freq will panic if i is negative or greater than or equal to t.Len().
   118  func (t *FFT) Freq(i int) float64 {
   119  	if i < 0 || t.Len() <= i {
   120  		panic("fourier: index out of range")
   121  	}
   122  	step := 1 / float64(t.Len())
   123  	return step * float64(i)
   124  }
   125  
   126  // CmplxFFT implements Fast Fourier Transform and its inverse for complex sequences.
   127  type CmplxFFT struct {
   128  	work []float64
   129  	ifac [15]int
   130  
   131  	// real temporarily store complex data as
   132  	// pairs of real values to allow passing to
   133  	// the backing code. The length of real
   134  	// must always be half the length of work.
   135  	real []float64
   136  }
   137  
   138  // NewCmplxFFT returns an CmplxFFT initialized for work on sequences of length n.
   139  func NewCmplxFFT(n int) *CmplxFFT {
   140  	var t CmplxFFT
   141  	t.Reset(n)
   142  	return &t
   143  }
   144  
   145  // Len returns the length of the acceptable input.
   146  func (t *CmplxFFT) Len() int { return len(t.work) / 4 }
   147  
   148  // Reset reinitializes the FFT for work on sequences of length n.
   149  func (t *CmplxFFT) Reset(n int) {
   150  	if 4*n <= cap(t.work) {
   151  		t.work = t.work[:4*n]
   152  		t.real = t.real[:2*n]
   153  	} else {
   154  		t.work = make([]float64, 4*n)
   155  		t.real = make([]float64, 2*n)
   156  	}
   157  	fftpack.Cffti(n, t.work, t.ifac[:])
   158  }
   159  
   160  // Coefficients computes the Fourier coefficients of a complex input sequence,
   161  // converting the time series in seq into the frequency spectrum, placing
   162  // the result in dst and returning it. This transform is unnormalized; a call
   163  // to Coefficients followed by a call of Sequence will multiply the input
   164  // sequence by the length of the sequence.
   165  //
   166  // If the length of seq is not t.Len(), Coefficients will panic.
   167  // If dst is nil, a new slice is allocated and returned. If dst is not nil and
   168  // the length of dst does not equal the length of seq, Coefficients will panic.
   169  // It is safe to use the same slice for dst and seq.
   170  func (t *CmplxFFT) Coefficients(dst, seq []complex128) []complex128 {
   171  	if len(seq) != t.Len() {
   172  		panic("fourier: sequence length mismatch")
   173  	}
   174  	if dst == nil {
   175  		dst = make([]complex128, len(seq))
   176  	} else if len(dst) != len(seq) {
   177  		panic("fourier: destination length mismatch")
   178  	}
   179  	for i, cv := range seq {
   180  		t.real[2*i] = real(cv)
   181  		t.real[2*i+1] = imag(cv)
   182  	}
   183  	fftpack.Cfftf(len(dst), t.real, t.work, t.ifac[:])
   184  	for i := range dst {
   185  		dst[i] = complex(t.real[2*i], t.real[2*i+1])
   186  	}
   187  	return dst
   188  }
   189  
   190  // Sequence computes the complex periodic sequence from the Fourier coefficients,
   191  // converting the frequency spectrum in coeff into a time series, placing the
   192  // result in dst and returning it. This transform is unnormalized; a call to
   193  // Coefficients followed by a call of Sequence will multiply the input sequence
   194  // by the length of the sequence.
   195  //
   196  // If the length of coeff is not t.Len(), Sequence will panic.
   197  // If dst is nil, a new slice is allocated and returned. If dst is not nil and
   198  // the length of dst does not equal the length of coeff, Sequence will panic.
   199  // It is safe to use the same slice for dst and coeff.
   200  func (t *CmplxFFT) Sequence(dst, coeff []complex128) []complex128 {
   201  	if len(coeff) != t.Len() {
   202  		panic("fourier: coefficients length mismatch")
   203  	}
   204  	if dst == nil {
   205  		dst = make([]complex128, len(coeff))
   206  	} else if len(dst) != len(coeff) {
   207  		panic("fourier: destination length mismatch")
   208  	}
   209  	for i, cv := range coeff {
   210  		t.real[2*i] = real(cv)
   211  		t.real[2*i+1] = imag(cv)
   212  	}
   213  	fftpack.Cfftb(len(dst), t.real, t.work, t.ifac[:])
   214  	for i := range dst {
   215  		dst[i] = complex(t.real[2*i], t.real[2*i+1])
   216  	}
   217  	return dst
   218  }
   219  
   220  // Freq returns the relative frequency center for coefficient i.
   221  // Freq will panic if i is negative or greater than or equal to t.Len().
   222  func (t *CmplxFFT) Freq(i int) float64 {
   223  	if i < 0 || t.Len() <= i {
   224  		panic("fourier: index out of range")
   225  	}
   226  	step := 1 / float64(t.Len())
   227  	if i < (t.Len()-1)/2+1 {
   228  		return step * float64(i)
   229  	}
   230  	return step * float64(i-t.Len())
   231  }
   232  
   233  // ShiftIdx returns a shifted index into a slice of coefficients
   234  // returned by the CmplxFFT so that indexing into the coefficients
   235  // places the zero frequency component at the center of the spectrum.
   236  // ShiftIdx will panic if i is negative or greater than or equal to
   237  // t.Len().
   238  func (t *CmplxFFT) ShiftIdx(i int) int {
   239  	if i < 0 || t.Len() <= i {
   240  		panic("fourier: index out of range")
   241  	}
   242  	h := t.Len() / 2
   243  	if i < h {
   244  		return i + (t.Len()+1)/2
   245  	}
   246  	return i - h
   247  }
   248  
   249  // UnshiftIdx returns inverse of ShiftIdx. UnshiftIdx will panic if i is
   250  // negative or greater than or equal to t.Len().
   251  func (t *CmplxFFT) UnshiftIdx(i int) int {
   252  	if i < 0 || t.Len() <= i {
   253  		panic("fourier: index out of range")
   254  	}
   255  	h := (t.Len() + 1) / 2
   256  	if i < h {
   257  		return i + t.Len()/2
   258  	}
   259  	return i - h
   260  }