WIP: 2 fits, not perfect.

WIP
J. David Lee 2019-12-11 04:19:33 +01:00
parent 72af0bdd82
commit e6340ab5d3
5 changed files with 337 additions and 27 deletions

View File

@ -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)
}

View File

@ -0,0 +1 @@
package extractfreq

118
lib/extractfreq/peaks.go Normal file
View File

@ -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
}

View File

@ -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)
}

View File

@ -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)
}