ウェーブレット変換のデモコード

cwt_demo.go

package main

import (
	"fmt"
	"github.com/mjibson/go-dsp/fft"
	"math"
	"math/cmplx"
)

// Morlet represents the Morlet wavelet function.
type Morlet struct {
	W0 float64 // frequency
}

// FrequencyRep returns wavelet in frequency domain.
func (m *Morlet) FrequencyRep(omega []float64, scale float64) []float64 {
	rep := make([]float64, len(omega))

	sw := make([]float64, len(omega))
	for i := range sw {
		sw[i] = omega[i] * scale
	}

	norm := math.Pow(math.Pi, -0.25)
	for i := range rep {
		if sw[i] > 0 {
			rep[i] = norm * math.Exp(-0.5*math.Pow(sw[i]-m.W0, 2.0))
		}
	}

	return rep
}

// Wavelet represents Continuous Wavelet Analysis.
type Wavelet struct {
	Dt, Dj float64
	Scales []float64
	Basis  *Morlet
}

// CWT performs Continuous Wavelet Transform.
func (w *Wavelet) CWT(x []float64) [][]complex128 {
	spectrogram := make([][]complex128, len(x))
	for i := range spectrogram {
		spectrogram[i] = make([]complex128, len(w.Scales))
	}

	y := fft.FFTReal(x)
	omega := angularFrequency(len(y), w.Dt)

	for j, scale := range w.Scales {
		norm := math.Sqrt(2.0 * math.Pi * scale / w.Dt)

		waveletInFreqDomain := w.Basis.FrequencyRep(omega, scale)

		product := make([]complex128, len(x))
		for n := 0; n < len(y); n++ {
			product[n] = complex(norm*waveletInFreqDomain[n], 0) * y[n]
		}
		convolved := fft.IFFT(product)

		for n := range convolved {
			spectrogram[n][j] = convolved[n]
		}
	}

	return spectrogram
}

func angularFrequency(N int, dt float64) []float64 {
	omega := make([]float64, N)

	for i := 0; i <= N/2; i++ {
		omega[i] = 2.0 * math.Pi * float64(i) / (dt * float64(N))
	}
	for i := N/2 + 1; i < N; i++ {
		omega[i] = -2.0 * math.Pi * float64(i) / (dt * float64(N))
	}

	return omega
}

func createCentScales(baseFreq float64, numOctaves, centInterval int) []float64 {
	numFreqBins := 1200 * numOctaves / centInterval
	scales := make([]float64, numFreqBins)

	for s := 0; s < numFreqBins; s++ {
		order := float64(s) * float64(centInterval) / 1200.0
		freq := baseFreq * math.Pow(2.0, order)
		scales[s] = 1.0 / freq
	}

	return scales
}

func main() {
	w := &Wavelet{
		Dt:     1.0 / 10000.0,
		Dj:     25.0 / 1200.0,
		Scales: createCentScales(55.0, 8, 25),
		Basis:  &Morlet{W0: 6.0}}

	result := w.CWT([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 4, 3})

	printMatrixAsGnuplotFormat(result)
}

func printMatrixAsGnuplotFormat(matrix [][]complex128) {
	fmt.Println("#", len(matrix[0]), len(matrix))
	for i, vec := range matrix {
		for j, val := range vec {
			fmt.Println(i, j, cmplx.Abs(val))
		}
		fmt.Println("")
	}
}