sin_noise_peak_plot.go

package main

import (
	"fmt"
	"image/color"
	"math"
	"math/rand/v2"

	"github.com/mpiannucci/peakdetect"

	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/plotutil"
	"gonum.org/v1/plot/vg"
	"gonum.org/v1/plot/vg/draw"
)

func sinfunc(n int) plotter.XYs {
	xy := make(plotter.XYs, n)
	xmax := float64(2.0)
	for i := 0; i < n; i++ {
		xy[i].X = xmax / float64(n) * float64(i)
		xy[i].Y = math.Sin(2 * math.Pi * xy[i].X)
	}
	return xy
}

func addnoise(xy plotter.XYs) plotter.XYs {
	max := 0.2
	min := -0.2
	for i := 0; i < len(xy); i++ {
		r := min + rand.Float64()*(max-min)
		xy[i].Y += r
	}
	return xy
}

func main() {
	N := 1024
	s0 := sinfunc(N)
	s0 = addnoise(s0)

	y := make([]float64, len(s0))
	for i := 0; i < len(y); i++ {
		_, y[i] = s0.XY(i)
	}

	// Returns four arrays -> min indices, min values, max indices, max values
	min_indices, min_values, max_indices, max_values := peakdetect.PeakDetect(y, 0.4)

	max_xy := make(plotter.XYs, len(max_indices))
	for i := 0; i < len(max_indices); i++ {
		fmt.Printf("peak %d: %v, %v\n", i, max_indices[i], max_values[i])
		max_xy[i].X, max_xy[i].Y = s0.XY(max_indices[i])
	}
	min_xy := make(plotter.XYs, len(min_indices))
	for i := 0; i < len(min_indices); i++ {
		fmt.Printf("valley %d: %v, %v\n", i, min_indices[i], min_values[i])
		min_xy[i].X, min_xy[i].Y = s0.XY(min_indices[i])
	}

	p := plot.New()
	p.Title.Text = "Peak detect example"
	p.X.Label.Text = "X"
	p.Y.Label.Text = "Y"

	p.Add(plotter.NewGrid())

	err := plotutil.AddLinePoints(p, "Sin", s0)
	if err != nil {
		panic(err)
	}

	max_scatter, err := plotter.NewScatter(max_xy)
	if err != nil {
		panic(err)
	}

	min_scatter, err := plotter.NewScatter(min_xy)
	if err != nil {
		panic(err)
	}

	max_scatter.GlyphStyle.Color = color.RGBA{R: 0, G: 0, B: 255, A: 255}
	max_scatter.GlyphStyle.Shape = draw.SquareGlyph{}
	max_scatter.GlyphStyle.Radius = vg.Points(7)
	p.Add(max_scatter)

	min_scatter.GlyphStyle.Color = color.RGBA{R: 0, G: 255, B: 0, A: 255}
	min_scatter.GlyphStyle.Shape = draw.SquareGlyph{}
	min_scatter.GlyphStyle.Radius = vg.Points(7)
	p.Add(min_scatter)

	// Save the plot to a PNG file.
	if err := p.Save(8*vg.Inch, 8*vg.Inch, "sin_noise_peak_plot.png"); err != nil {
		panic(err)
	}
}