sointu/tracker/spectrum.go

package tracker

import (
	"math"
	"math/cmplx"

	"github.com/vsariola/sointu"
)

type (
	SpecAnalyzer struct {
		settings SpecSettings
		broker   *Broker
		temp     specTemp
	}

	SpecSettings struct {
		Channels   SpecChannels
		Smooth     SpecSmooth
		Resolution int
	}

	SpecChannels int
	SpecSmooth   int
	Spectrum     []Decibel
	SpecResult   []Spectrum

	specTemp struct {
		spectra    []Spectrum
		chunk      sointu.AudioBuffer
		weight     []float32    // window weighting function
		normFactor float32      // normalization factor, to account for the windowing
		perm       []int        // bit-reversal permutation table
		tmp        []complex128 // temporary buffer for FFT
	}
)

const (
	SpecResolutionMin = 7
	SpecResolutionMax = 16
)

const (
	SpecChannelsOff       SpecChannels = iota // no spectrum analysis is done to save CPU resources
	SpecChannelsCombined                      // calculate a single combined spectrum for both channels
	SpecChannelsSeparated                     // calculate separate spectrums for left and right channels
	SpecChannelsLeft                          // calculate spectrum only for the left channel
	SpecChannelsRight                         // calculate spectrum only for the right channel
)

const (
	SpecSmoothSlow SpecSmooth = iota
	SpecSmoothMedium
	SpecSmoothFast
)

var spectrumSmoothingMap map[SpecSmooth]float32 = map[SpecSmooth]float32{
	SpecSmoothSlow:   0.05,
	SpecSmoothMedium: 0.2,
	SpecSmoothFast:   1.0,
}

func (s *SpecAnalyzer) Run() {
	for {
		select {
		case <-s.broker.CloseSpecAn:
			close(s.broker.FinishedSpecAn)
			return
		case msg := <-s.broker.ToSpecAn:
			s.handleMsg(msg)
		}
	}
}

func (s *SpecAnalyzer) handleMsg(msg any) {
	switch m := msg.(type) {
	case SpecSettings:
		if s.settings != m {
			s.init(m)
		}
	case sointu.AudioBuffer:
		s.update(m)
	default:
		// unknown message type; ignore
	}
}

func (a *SpecAnalyzer) init(s SpecSettings) {
	s.Resolution = min(max(s.Resolution, SpecResolutionMin), SpecResolutionMax)
	a.settings = s
	n := 1 << s.Resolution
	a.temp = specTemp{
		spectra: make([]Spectrum, 0),
		chunk:   make(sointu.AudioBuffer, n),
		weight:  make([]float32, n),
		perm:    make([]int, n),
		tmp:     make([]complex128, n),
	}
	for i := range n {
		// Hanning window
		w := float32(0.5 * (1 - math.Cos(2*math.Pi*float64(i)/float64(n-1))))
		a.temp.weight[i] = w
		a.temp.normFactor += w
		// initialize the bit-reversal permutation table
		a.temp.perm[i] = i
	}
	// compute the bit-reversal permutation
	for i, j := 1, 0; i < n; i++ {
		bit := n >> 1
		for ; j&bit != 0; bit >>= 1 {
			j ^= bit
		}
		j ^= bit
		if i < j {
			a.temp.perm[i], a.temp.perm[j] = a.temp.perm[j], a.temp.perm[i]
		}
	}
}

func (sd *spectrumDetector) update(input []float32) {
	c := sd.tmp
	for i := range sd.tmp {
		p := sd.perm[i]
		c[i] = complex(float64(input[p]*sd.window[p]), 0)
	}
	n := len(c)
	for l := 2; l <= n; l <<= 1 {
		ang := 2 * math.Pi / float64(l)
		wlen := complex(math.Cos(ang), math.Sin(ang))
		for i := 0; i < n; i += l {
			w := complex(1, 0)
			for j := 0; j < l/2; j++ {
				u := c[i+j]
				v := c[i+j+l/2] * w
				c[i+j] = u + v
				c[i+j+l/2] = u - v
				w *= wlen
			}
		}
	}
	for i := range input {
		a := cmplx.Abs(c[i])
		power := float32(a*a) / sd.normFactor
		sd.spectrum[i] += sd.alpha * (power - sd.spectrum[i])
	}
}