align-audio-0.0.0.2: src/Correlate.hs

module Correlate where

import Vector (findPeak, mulConj)
import Common (fourierRC)

import qualified Numeric.FFTW.Rank1 as Trafo1
import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Shape as Shape
import qualified Data.StorableVector.Lazy as SVL
import Data.Array.Comfort.Storable (Array)

import qualified Data.Stream as Stream


{- $setup
>>> import qualified Test.QuickCheck as QC
>>> import qualified Data.StorableVector.Lazy as SVL
>>> import Data.Monoid ((<>))
-}


{- |
Round to next higher number of the form m*2^n with m<16.
This size contains almost only factor 2 and at most one ugly factor < 16
and we add at most 1/8 pad data.

>>> map ceilingFFTSize [1..40]
[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,18,18,20,20,22,22,24,24,26,26,28,28,30,30,32,32,36,36,36,36,40,40,40,40]
prop> \n -> n <= ceilingFFTSize n
prop> \(QC.Positive n) -> ceilingFFTSize n * 8 <= n * 9
-}
ceilingFFTSize :: Int -> Int
ceilingFFTSize n =
   uncurry (*) $
   Stream.head $
   Stream.dropWhile ((>=16).fst) $
      Stream.zip
         (Stream.iterate (\k -> div (k+1) 2) n)
         (Stream.iterate (2*) 1)


correlate ::
   (Class.Real a) =>
   SVL.Vector a -> SVL.Vector a -> Array (Shape.Cyclic Int) a
correlate xs ys =
   let nx = SVL.length xs
       ny = SVL.length ys
       n = ceilingFFTSize (nx+ny)
   in Trafo1.fourierCR $ mulConj (fourierRC n xs) (fourierRC n ys)

{- |
prop> :{
   \(QC.NonNegative preX, QC.NonNegative postX) ->
   \(QC.NonNegative preY, QC.NonNegative postY) ->
   \(QC.NonEmpty xs) ->
   all (0/=)  xs QC.==>

   preX - preY
   QC.===
   let v = SVL.pack SVL.defaultChunkSize xs in
   let pad n = SVL.replicate SVL.defaultChunkSize n (0::Float) in
   determineLag
      (pad preX <> v <> pad postX)
      (pad preY <> v <> pad postY)
:}
-}
determineLag :: (Class.Real a) => SVL.Vector a -> SVL.Vector a -> Int
determineLag xs ys =
   let zs = correlate xs ys
   in case findPeak zs of
         (neg,nonNeg) -> if nonNeg >= SVL.length xs then neg else nonNeg