WIP: 2 fits, not perfect.
parent
72af0bdd82
commit
e6340ab5d3
|
@ -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)
|
||||
}
|
|
@ -0,0 +1 @@
|
|||
package extractfreq
|
|
@ -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
|
||||
}
|
|
@ -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)
|
||||
}
|
||||
|
|
|
@ -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)
|
||||
}
|
||||
|
|
Reference in New Issue