1package detect
2
3import (
4 "embed"
5 "encoding/binary"
6 "ewdetect/utils"
7 "math"
8 "os"
9 "strings"
10
11 "ewdetect/config"
12
13 "github.com/argusdusty/gofft"
14 "github.com/rs/zerolog/log"
15 "gonum.org/v1/plot"
16 "gonum.org/v1/plot/plotter"
17 "gonum.org/v1/plot/vg"
18)
19
20//go:embed assets/p_wave_sample.bin
21//go:embed assets/s_wave_sample.bin
22var f embed.FS
23
24var PWaveSample []int32
25var SWaveSample []int32
26
27var ExampleFile []int32
28
29func Init() {
30 pFile, err := f.Open("assets/p_wave_sample.bin")
31 utils.CheckError(err)
32 defer pFile.Close()
33
34 sFile, err := f.Open("assets/s_wave_sample.bin")
35 utils.CheckError(err)
36 defer sFile.Close()
37
38 var val int32
39 for {
40 err := binary.Read(pFile, binary.LittleEndian, &val)
41 if err != nil {
42 break
43 }
44 PWaveSample = append(PWaveSample, val)
45 }
46
47 for {
48 err := binary.Read(sFile, binary.LittleEndian, &val)
49 if err != nil {
50 break
51 }
52 SWaveSample = append(SWaveSample, val)
53 }
54
55 log.Debug().Msg("Wave samples initialized")
56}
57
58func CrossCorrelation(a *[config.BufferSize]int32, b *[]int32, offset int, headPosition int, thresholdMode bool, threshold float64, drawDebugOutput bool, debugOutputFilename string) (int, float64, float64) { // find b's offset in a
59 if len(*a) < len(*b) {
60 log.Fatal().Msg("Analyzed signal must be longer than the reference signal")
61 }
62
63 length := 1
64 for length < len(*a) {
65 length *= 2
66 }
67 x := make([]complex128, length)
68 y := make([]complex128, length)
69
70 for i := range len(*a) {
71 x[i] = complex(float64((*a)[(i+headPosition)%config.BufferSize]), 0)
72 }
73
74 for i := range len(*b) {
75 y[length-i-1] = complex(float64((*b)[i]), 0)
76 }
77 err := gofft.FastConvolve(x, y)
78 utils.CheckError(err)
79
80 if config.Debug && drawDebugOutput {
81 pts := make(plotter.XYs, len(x))
82 for i := range x {
83 pts[i].X = float64(i)
84 pts[i].Y = real(x[i])
85 }
86
87 p := plot.New()
88 p.Add(plotter.NewGrid())
89 line, err := plotter.NewLine(pts)
90 utils.CheckError(err)
91 p.Add(line)
92
93 err = p.Save(8*vg.Inch, 4*vg.Inch, debugOutputFilename)
94 utils.CheckError(err)
95 log.Debug().Str("filename", debugOutputFilename).Msg("Debug plot saved")
96 }
97
98 var dc float64
99 // Calculate DC
100 for i := range x {
101 if i < config.BufferSize {
102 dc += real(x[i])
103 }
104 }
105 dc = dc / float64(config.BufferSize)
106
107 var sumSquares float64
108 // Calculate RMS
109 for i := range x {
110 if i < config.BufferSize {
111 sumSquares += (real(x[i]) - dc) * (real(x[i]) - dc)
112 }
113 }
114 rms := math.Sqrt(sumSquares / float64(config.BufferSize))
115
116 max := real(x[offset])
117 tMax := offset
118 if thresholdMode {
119 for i := range len(x) - offset {
120 if i < config.BufferSize {
121 if real(x[(i+offset)%config.BufferSize])/rms > threshold {
122 max = real(x[(i+offset)%config.BufferSize])
123 tMax = (i + offset) % config.BufferSize
124 break
125 }
126 }
127 }
128 } else {
129 for i := range len(x) - offset {
130 if i < config.BufferSize {
131 if real(x[(i+offset)%config.BufferSize]) > max {
132 max = real(x[(i+offset)%config.BufferSize])
133 tMax = (i + offset) % config.BufferSize
134 }
135 }
136 }
137 }
138
139 log.Debug().
140 Float64("rms", rms).
141 Float64("max", max).
142 Int("tMax", tMax).
143 Msg("Cross correlation completed")
144
145 return tMax, rms, max
146}
147
148func DetectWaves(a *[config.BufferSize]int32, headPosition int) (int, int, float64, float64, float64, float64) {
149 pWaveArrival, pWaveRMS, pWaveMax := CrossCorrelation(a, &PWaveSample, 0, headPosition, true, config.PWaveSensitivity, false, "")
150 sWaveArrival, sWaveRMS, sWaveMax := CrossCorrelation(a, &SWaveSample, pWaveArrival+config.MinWaveDelayInSamples, headPosition, false, 0, false, "")
151
152 log.Debug().
153 Int("pWaveArrival", pWaveArrival).
154 Int("sWaveArrival", sWaveArrival).
155 Float64("pWaveRMS", pWaveRMS).
156 Float64("sWaveRMS", sWaveRMS).
157 Msg("Wave detection completed")
158
159 return pWaveArrival, sWaveArrival, pWaveRMS, sWaveRMS, pWaveMax, sWaveMax
160}
161
162func ThresholdDetectWaves(a *[config.BufferSize]int32, headPosition int, detectorName string, outputDebugGraphs bool) (bool, int, int) {
163 pWaveArrival, sWaveArrival, pWaveRMS, sWaveRMS, pWaveMax, sWaveMax := DetectWaves(a, headPosition)
164 if pWaveMax/pWaveRMS > config.PWaveSensitivity && sWaveMax/sWaveRMS > config.SWaveSensitivity {
165 if config.Debug && outputDebugGraphs {
166 splitPath := strings.Split(detectorName, "/")
167 path_a := "debug/cross-correlation/" + strings.Join(splitPath[:len(splitPath)-1], "/")
168 path_b := "debug/traces/" + strings.Join(splitPath[:len(splitPath)-1], "/")
169 if _, err := os.Stat(path_a); os.IsNotExist(err) {
170 err := os.Mkdir(path_a, os.ModePerm)
171 utils.CheckError(err)
172 err = os.Mkdir(path_b, os.ModePerm)
173 utils.CheckError(err)
174 }
175 CrossCorrelation(a, &PWaveSample, 0, headPosition, true, config.PWaveSensitivity, true, "debug/cross-correlation/"+detectorName+"pwave-correlation.png")
176 CrossCorrelation(a, &SWaveSample, pWaveArrival+config.MinWaveDelayInSamples, headPosition, false, -1, true, "debug/cross-correlation/"+detectorName+"swave-correlation.png")
177 traceOutputFileName := "debug/traces/" + detectorName + "trace.png"
178
179 pts := make(plotter.XYs, len(a))
180 for i := range len(a) {
181 pts[i].X = float64(i)
182 pts[i].Y = float64(a[i])
183 }
184 p := plot.New()
185 p.Add(plotter.NewGrid())
186 line, err := plotter.NewLine(pts)
187 utils.CheckError(err)
188 p.Add(line)
189 err = p.Save(8*vg.Inch, 4*vg.Inch, traceOutputFileName)
190 utils.CheckError(err)
191 log.Debug().Str("filename", traceOutputFileName).Msg("Debug plot saved")
192 log.Debug().Str("detector", detectorName).Msg("Debug graphs generated")
193 }
194 return true, pWaveArrival, sWaveArrival
195 }
196 return false, pWaveArrival, sWaveArrival
197}