From e6340ab5d3507b2395a258255b10d8e706b311b6 Mon Sep 17 00:00:00 2001 From: "J. David Lee" Date: Wed, 11 Dec 2019 04:19:33 +0100 Subject: [PATCH] WIP: 2 fits, not perfect. --- lib/extractfreq/harmonicfit.go | 66 +++++++++++ lib/extractfreq/parabolicinterp.go | 1 + lib/extractfreq/peaks.go | 118 +++++++++++++++++++ lib/extractfreq/wip.go | 177 ++++++++++++++++++++++++----- lib/extractfreq/wip_test.go | 2 +- 5 files changed, 337 insertions(+), 27 deletions(-) create mode 100644 lib/extractfreq/harmonicfit.go create mode 100644 lib/extractfreq/parabolicinterp.go create mode 100644 lib/extractfreq/peaks.go diff --git a/lib/extractfreq/harmonicfit.go b/lib/extractfreq/harmonicfit.go new file mode 100644 index 0000000..49ad6c8 --- /dev/null +++ b/lib/extractfreq/harmonicfit.go @@ -0,0 +1,66 @@ +package extractfreq + +import ( + "math" + + "gonum.org/v1/gonum/optimize" +) + +func harmonicModel(A, B, n float64) float64 { + return A * n * math.Sqrt(B*n*n+1) +} + +// Fit spectrum peaks to function: +// +// f(n) = A*n * sqrt(1+B*n**2) +func fitSpectrumPeaks( + f1 float64, // Approximate fundamental. + ns []float64, + freqs []float64, // Peak frequencies. + amps []float64, // Peak amplitudes. +) ( + A float64, + B float64, + err error, +) { + // Error function weights each harmonic error by its amplitude. + errFunc := func(x []float64) float64 { + A := math.Abs(x[0]) + B := math.Abs(x[1]) + + sum := float64(0) + + for i, fPeak := range freqs { + n := ns[i] + e := fPeak - harmonicModel(A, B, n) + e *= e + e *= amps[i] + sum += e + } + + return sum + } + + // Initial guess. + x := []float64{f1, 0} + + problem := optimize.Problem{ + Func: errFunc, + } + + result, err := optimize.Minimize(problem, x, nil, nil) + + return result.X[0], result.X[1], err +} + +func ApproxHarmonicNumber(fFun, fPeak float64) float64 { + if fPeak < fFun { + switch math.Round(fFun / fPeak) { + case 2: + return 0.5 + case 4: + return 0.25 + } + } + return math.Round(fPeak / fFun) +} diff --git a/lib/extractfreq/parabolicinterp.go b/lib/extractfreq/parabolicinterp.go new file mode 100644 index 0000000..a04cd46 --- /dev/null +++ b/lib/extractfreq/parabolicinterp.go @@ -0,0 +1 @@ +package extractfreq diff --git a/lib/extractfreq/peaks.go b/lib/extractfreq/peaks.go new file mode 100644 index 0000000..e931ebf --- /dev/null +++ b/lib/extractfreq/peaks.go @@ -0,0 +1,118 @@ +package extractfreq + +import "gonum.org/v1/gonum/floats" + +type mmItem struct { + Val float64 + MaxVal float64 +} + +type moveMax struct { + pushStack []mmItem + popStack []mmItem +} + +func newMoveMax() *moveMax { + return &moveMax{ + pushStack: []mmItem{}, + popStack: []mmItem{}, + } +} + +func (mm *moveMax) Max() float64 { + pushTop := len(mm.pushStack) - 1 + popTop := len(mm.popStack) - 1 + + if pushTop == -1 { + item := mm.popStack[popTop] + return item.MaxVal + } else if popTop == -1 { + item := mm.pushStack[pushTop] + return item.MaxVal + } + + item1 := mm.pushStack[pushTop] + item2 := mm.popStack[popTop] + if item1.MaxVal > item2.MaxVal { + return item1.MaxVal + } else { + return item2.MaxVal + } +} + +func (mm *moveMax) push(stack []mmItem, val float64) []mmItem { + item := mmItem{ + Val: val, + MaxVal: val, + } + + // Ensure max is set correctly. + pushTop := len(stack) - 1 + if pushTop > -1 { + top := stack[pushTop] + if top.MaxVal > item.MaxVal { + item.MaxVal = top.MaxVal + } + } + + // Push. + return append(stack, item) +} + +func (mm *moveMax) Push(val float64) { + mm.pushStack = mm.push(mm.pushStack, val) +} + +func (mm *moveMax) Pop() { + if len(mm.popStack) == 0 { + // Move push stack onto pop stack. + for len(mm.pushStack) > 0 { + item := mm.pushStack[len(mm.pushStack)-1] + mm.pushStack = mm.pushStack[:len(mm.pushStack)-1] + mm.popStack = mm.push(mm.popStack, item.Val) + } + } + mm.popStack = mm.popStack[:len(mm.popStack)-1] +} + +func getPeaks(nHalfWin int, data []float64) (inds []int, amps []float64) { + mm := newMoveMax() + nWin := 2*nHalfWin + 1 + + if len(data) < nWin { + nWin = len(data) + } + + // Push nWin samples. + for i := range data[:nWin] { + mm.Push(data[i]) + } + + for idxPush := nWin + 1; idxPush < len(data); idxPush++ { + mm.Pop() + mm.Push(data[idxPush]) + idx := idxPush - nHalfWin + max := mm.Max() + if data[idx] == max && data[idx+1] != max { + inds = append(inds, idx) + amps = append(amps, data[idx]) + } + } + + return +} + +func sortPeaks(inds []int, amps []float64) ([]int, []float64) { + sort := make([]int, len(amps)) + out := make([]int, len(amps)) + floats.Argsort(amps, sort) + for i, idx := range sort { + out[i] = inds[idx] + } + + for i := 0; i < len(inds)/2; i++ { + out[i], out[len(out)-1-i] = out[len(out)-1-i], out[i] + amps[i], amps[len(amps)-1-i] = amps[len(amps)-1-i], amps[i] + } + return out, amps +} diff --git a/lib/extractfreq/wip.go b/lib/extractfreq/wip.go index 0fa9ff4..202d6e2 100644 --- a/lib/extractfreq/wip.go +++ b/lib/extractfreq/wip.go @@ -7,6 +7,7 @@ import ( "math/cmplx" "git.crumpington.com/public/jlaudio/lib/flac" + "git.crumpington.com/public/jlaudio/lib/util" "gonum.org/v1/gonum/fourier" "gonum.org/v1/plot" "gonum.org/v1/plot/plotter" @@ -48,54 +49,169 @@ func PlotFFT(path string, nSamples int, fFundamental float64) { } // Trim to frequencies of interest. + freq0 := fFundamental / 8.0 + freq1 := 21 * fFundamental + if freq1 > 20000 { + freq1 = 20000 + } + deltaF := freqs[1] - freqs[0] - idx0 := int(20 / deltaF) - idx1 := int(16000 / deltaF) + idx0 := int(freq0 / deltaF) + idx1 := int(freq1 / deltaF) freqs = freqs[idx0:idx1] amps = amps[idx0:idx1] - // Multiply by A-weight function. + // Multiply by A-weight function, then normalize. + ampMax := float64(0) for i := range amps { amps[i] *= AWeight(freqs[i]) + if amps[i] > ampMax { + ampMax = amps[i] + } + } + + for i := range amps { + if amps[i] > 0 { + amps[i] /= ampMax + } } // Find peaks. - dFreq := freqs[1] - freqs[0] - halfWinSize := int(math.Round((fFundamental / dFreq) / 3)) + halfWinSize := int(math.Round((fFundamental / deltaF) / 3)) peakInds, peakAmps := getPeaks(halfWinSize, amps) - peakInds, peakAmps = sortPeaks(peakInds, peakAmps) + peakFreqs := make([]float64, len(peakInds)) + for i, idx := range peakInds { + peakFreqs[i] = freqs[idx] + } + + // First, determine harmonic numbers of each peak. + ns := make([]float64, len(peakFreqs)) + for i, f := range peakFreqs { + ns[i] = ApproxHarmonicNumber(fFundamental, f) + } + + // -------------------------------------------------------------------------- + + // First fit. + + log.Printf("First fit...") + A, B, err := fitSpectrumPeaks(fFundamental, ns, peakFreqs, peakAmps) + if err != nil { + panic(err) + } + + log.Printf("A: %.4f B: %0.4f", A, B) p, err := plot.New() if err != nil { panic(err) } - fMin := float64(20) - fMax := float64(16000) + ptsModel := make(plotter.XYs, len(ns)) + ptsData := make(plotter.XYs, len(ns)) + for i, n := range ns { + ptsModel[i].X = n + ptsModel[i].Y = harmonicModel(A, B, n) + ptsData[i].X = n + ptsData[i].Y = peakFreqs[i] + } + + if err = plotutil.AddScatters(p, "model", ptsModel, "data", ptsData); err != nil { + panic(err) + } + + p.Add(plotter.NewGrid()) + + log.Printf("Saving...") + if err := p.Save(16*vg.Inch, 8*vg.Inch, "fit1.png"); err != nil { + panic(err) + } + + // Remove peaks w/ larger errors: + // TODO: Method? + iOut := 0 + for i := range ns { + f := peakFreqs[i] + fModel := harmonicModel(A, B, ns[i]) + deltaF := math.Abs(f - fModel) + maxDeltaF := 20 * util.CentDeltaFreq(fModel) + if deltaF < maxDeltaF { + ns[iOut] = ns[i] + peakInds[iOut] = peakInds[i] + peakFreqs[iOut] = peakFreqs[i] + peakAmps[iOut] = peakAmps[i] + iOut++ + } + } + + ns = ns[:iOut] + peakInds = peakInds[:iOut] + peakFreqs = peakFreqs[:iOut] + peakAmps = peakAmps[:iOut] + + // Second fit. + + A, B, err = fitSpectrumPeaks(A, ns, peakFreqs, peakAmps) + if err != nil { + panic(err) + } + + log.Printf("A: %.4f B: %0.4f", A, B) + + p, err = plot.New() + if err != nil { + panic(err) + } + + ptsModel = make(plotter.XYs, len(ns)) + ptsData = make(plotter.XYs, len(ns)) + for i, n := range ns { + ptsModel[i].X = n + ptsModel[i].Y = harmonicModel(A, B, n) + ptsData[i].X = n + ptsData[i].Y = peakFreqs[i] + } + + if err = plotutil.AddScatters(p, "model", ptsModel, "data", ptsData); err != nil { + panic(err) + } + + p.Add(plotter.NewGrid()) + + log.Printf("Saving...") + if err := p.Save(16*vg.Inch, 8*vg.Inch, "fit2.png"); err != nil { + panic(err) + } + + // -------------------------------------------------------------------------- + + p, err = plot.New() + if err != nil { + panic(err) + } log.Printf("Making points...") pts := make(plotter.XYs, 0, len(amps)) for i := range amps { - if freqs[i] > fMin && freqs[i] < fMax { - pts = append(pts, plotter.XY{ - X: freqs[i], - Y: amps[i], - }) - } + pts = append(pts, plotter.XY{ + X: freqs[i], + Y: math.Log(amps[i]), + }) } pts2 := make(plotter.XYs, 0, len(amps)) for i := range peakInds { - if freqs[peakInds[i]] > fMin && freqs[peakInds[i]] < fMax { - pts2 = append(pts2, plotter.XY{ - X: freqs[peakInds[i]], - Y: peakAmps[i], - }) - } + pts2 = append(pts2, plotter.XY{ + X: freqs[peakInds[i]], + Y: math.Log(peakAmps[i]), + }) } log.Printf("Plotting...") - if err = plotutil.AddScatters(p, "fft", pts, "max", pts2); err != nil { + if err = plotutil.AddLines(p, "fft", pts); err != nil { + panic(err) + } + if err = plotutil.AddScatters(p, "max", pts2); err != nil { panic(err) } p.Add(plotter.NewGrid()) @@ -106,7 +222,7 @@ func PlotFFT(path string, nSamples int, fFundamental float64) { } for i, idx := range peakInds { - //amp := peakAmps[i] + n := ns[i] // Make window. ampsWin := amps[idx-halfWinSize : idx+halfWinSize] @@ -121,6 +237,15 @@ func PlotFFT(path string, nSamples int, fFundamental float64) { }) } + pts2 := plotter.XYs{ + { + X: harmonicModel(A, B, n), + Y: peakAmps[i], + }, + } + + log.Printf("%f, %f", peakFreqs[i], harmonicModel(A, B, n)) + log.Printf("Plotting...") p, err := plot.New() if err != nil { @@ -130,12 +255,12 @@ func PlotFFT(path string, nSamples int, fFundamental float64) { if err = plotutil.AddLines(p, "fft", pts); err != nil { panic(err) } - - //p.Y.Scale = plot.LogScale{} - //p.X.Scale = plot.LogScale{} + if err = plotutil.AddScatters(p, "model", pts2); err != nil { + panic(err) + } log.Printf("Saving...") - filename := fmt.Sprintf("plot-%03d.png", i) + filename := fmt.Sprintf("plot-%03f.png", n) if err := p.Save(8*vg.Inch, 8*vg.Inch, filename); err != nil { panic(err) } diff --git a/lib/extractfreq/wip_test.go b/lib/extractfreq/wip_test.go index 4cf6370..1668c91 100644 --- a/lib/extractfreq/wip_test.go +++ b/lib/extractfreq/wip_test.go @@ -3,5 +3,5 @@ package extractfreq import "testing" func TestPlotFFT(t *testing.T) { - PlotFFT("test_files/440.flac", 96000, 440) + PlotFFT("test_files/440.flac", 48000, 440) }