refactor(dsp): improve findsignal performance and API

src/dsp.jl

@@ -1,8 +1,7 @@
-import DSP
-import DSP: filt, filtfilt, resample, nextfastfft
+import DSP: DSP, filt, filtfilt, resample, nextfastfft
 import Statistics: std
-import Peaks: findmaxima, peakproms!
-import Optim: optimize, minimizer
+import Peaks: argmaxima, peakproms!
+import Optim: optimize, minimizer, BFGS
 
 export fir, removedc, removedc!, demon
 export upconvert, downconvert, rrcosfir, rcosfir
@@ -663,9 +662,10 @@ delay!(x, t::Units.Unitful.Time; kwargs...) = delay!(x, inseconds(t) * framerate
 $(SIGNATURES)
 Finds up to `n` strongest copies of reference signal `r` in signal `s`. The reference
 signal `r` should have a delta-like autocorrelation for this function to work
-well. If the keyword parameter `coarse` is set to `true`, approximate arrival
-times are computed based on a matched filter. If it is set to `false`, an
-iterative optimization is performed to find more accruate arrival times.
+well. If the keyword parameter `finetune` is set to 0, approximate arrival
+times are computed based on a matched filter. If it is set to a positive value, an
+iterative optimization is performed to find more accruate arrival times within
+roughly `finetune` samples of the coarse arrival times.
 
 Returns named tuple `(time=t, amplitude=a, mfo=m)` where `t` is a vector of arrival
 times, `a` is a vector of complex amplitudes of the arrivals, and `m` is the complex
@@ -682,52 +682,47 @@ julia> y4[254:253+length(x4)] += -0.8 * real(x4) # time 0.001544𝓈, index 64.
 julia> y4[513:512+length(x4)] += 0.6 * real(x4) # time 0.003125𝓈, index 129.0
 julia> y = resample(y4, 1//4)
 julia> y .+= 0.1 * randn(length(y))
-julia> findsignal(x, y, 3; coarse=true)
+julia> findsignal(x, y, 3)
 (time = [0.000781, 0.001538, 0.003125], amplitude = [...], mfo=[])
 julia> findsignal(x, y, 3; mfo=true)
 (time = [0.000775, 0.001545, 0.003124], amplitude = [...], mfo=[...])
 ```
 """
-function findsignal(r, s, n=1; prominence=0.2, coarse=false, mfo=false)
+function findsignal(r, s, n=1; prominence=0.0, finetune=2, mfo=false)
  # coarse arrival time estimation
  r = analytic(r)
- r = r / std(r)
  s = analytic(s)
- m = mfilter(r, s) / length(r)
+ m = mfilter(r, s)
+ m ./= (std(r) * length(r))
  T = eltype(m)
- absm = abs.(samples(m))
- p, _ = findmaxima(absm)
- peakproms!(p, absm; minprom=prominence*maximum(absm))
+ m̄ = abs.(samples(m))
+ p = argmaxima(m̄)
+ prominence > 0 && peakproms!(p, m̄; minprom=prominence*maximum(m̄))
  length(p) > length(s)/10 && return (time=Float64[], amplitude=T[], mfo=(mfo ? m : T[]))
- h = absm[p]
- ndx = sortperm(h; rev=true)
- length(ndx) > n && (ndx = ndx[1:n])
+ h = m̄[p]
+ ndx = partialsortperm(h, 1:n; rev=true)
  p = p[ndx]
- if coarse
+ if finetune == 0
  t = time(Float64.(p), s)
  ndx = sortperm(t)
  return (time=t[ndx], amplitude=samples(m[p[ndx]]), mfo=(mfo ? m : T[]))
  end
  # iterative fine arrival time estimation
- margin = 5 # arrival time may vary up to margin from coarse estimates
  i::Int = minimum(p)
- n = maximum(p) - i + length(r) + 2 * margin
- n = nextfastfft(n)
- i = max(1, i - margin)
- N = n
- i + N - 1 > length(s) && (N = length(s) - i + 1)
- X = fft(vcat(samples(r), zeros(n-length(r))))
- soln = let p=p, f=fftfreq(n)
- function reconstruct(v)
- ii = @view v[1:length(p)]
- aa = @views complex.(v[length(p)+1:2*length(p)], v[2*length(p)+1:3*length(p)])
- Z = mapreduce(+, zip(ii, aa)) do (i, a)
+ N::Int = maximum(p) - i + length(r) + 2 * finetune
+ N = nextfastfft(N)
+ i = max(1, i - finetune)
+ X = fft(vcat(samples(r), zeros(N-length(r))))
+ S = fft(vcat(samples(s[i:min(i+N-1,end)]), zeros(max(i+N-1-length(s),0))))
+ soln = let P=length(p), f=fftfreq(N)
+ optimize([p .- i; real.(m[p]); imag.(m[p])], BFGS(); autodiff=:forward) do v
+ ii = @view v[1:P]
+ aa = @views complex.(v[P+1:2P], v[2P+1:3P])
+ X̄ = mapreduce(+, zip(ii, aa)) do (i, a)
  a .* X .* cis.(-2π .* i .* f)
  end
- @view real(ifft(Z))[1:N]
+ sum(abs2, X̄ .- S)
  end
- v0 = [p .- i; real.(m[p]); imag.(m[p])]
- optimize(v -> sum(abs2, reconstruct(v) .- s[i:i+N-1]), v0)
  end
  v = minimizer(soln)
  pp = v[1:length(p)] .+ i

test/runtests.jl

@@ -646,18 +646,18 @@ end
  @test energy(z) ≈ energy(x)
 
  y .+= 0.1 * randn(rng, length(y))
- t, a, m = findsignal(x, y, 3; coarse=false)
- @test t ≈ [0.000775, 0.001545, 0.003124] atol=2e-6
+ t, a, m = findsignal(x, y, 3; finetune=0)
+ @test t ≈ [0.000775, 0.001544, 0.003125] atol=1e-5
  @test real(a) / real(a[1]) ≈ [1.0, -0.8, 0.6] atol=1e-2
  @test length(m) == 0
- t, a, m = findsignal(x, y, 3; coarse=true, mfo=true)
- @test t ≈ [0.000775, 0.001545, 0.003124] atol=1e-5
+ t, a, m = findsignal(x, y, 3; mfo=true)
+ @test t ≈ [0.000775, 0.001544, 0.003125] atol=1e-5
  @test real(a) / real(a[1]) ≈ [1.0, -0.8, 0.6] atol=1e-2
  @test length(m) > 0
 
  y = compose(real(x), time([32.75, 64.25, 129.0], x), [0.8, -0.7, 1.0]; duration=0.2)
- t, a, m = findsignal(x, y, 3; coarse=false)
- @test t ≈ [0.000775, 0.001545, 0.003124] atol=2e-6
+ t, a, m = findsignal(x, y, 3; finetune=0)
+ @test t ≈ [0.000775, 0.001544, 0.003125] atol=1e-5
  @test real(a) / real(a[3]) ≈ [0.8, -0.7, 1.0] atol=1e-2
 
  @test delay!([1,2,3,4], -1) == [2,3,4,0]