packages feed

arrayfire-0.8.0.0: test/ArrayFire/SignalSpec.hs

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
module ArrayFire.SignalSpec where

import qualified ArrayFire       as A
import           Data.Complex
import           Test.Hspec
import           Test.Hspec.QuickCheck (prop)
import           Test.QuickCheck       (NonEmptyList (..), choose, forAll, vectorOf)

-- | Check all elements of two Double arrays are within tolerance.
shouldBeApproxD
  :: A.Array Double
  -> A.Array Double
  -> Expectation
shouldBeApproxD actual expected =
  zipWith (\a e -> abs (a - e))
    (A.toList @Double actual)
    (A.toList @Double expected)
  `shouldSatisfy` all (< 1e-6)

-- | Check all elements of two Complex Double arrays are within tolerance.
shouldBeApproxC
  :: A.Array (Complex Double)
  -> A.Array (Complex Double)
  -> Expectation
shouldBeApproxC actual expected =
  zipWith (\a e -> magnitude (a - e))
    (A.toList @(Complex Double) actual)
    (A.toList @(Complex Double) expected)
  `shouldSatisfy` all (< 1e-10)

spec :: Spec
spec =
  describe "Signal" $ do

    describe "fft" $ do
      it "fftInPlace runs without error" $ do
        A.fftInPlace (A.scalar @(Complex Double) (1 :+ 0)) 1.0
          `shouldReturn` ()

      it "transform of a Dirac delta is a flat spectrum" $ do
        A.fft (A.mkArray @(Complex Double) [4] [1,0,0,0]) 1.0 4
          `shouldBeApproxC`
          A.mkArray @(Complex Double) [4] [1,1,1,1]

      it "transform of all-ones concentrates all energy at DC" $ do
        A.fft (A.mkArray @(Complex Double) [4] [1,1,1,1]) 1.0 4
          `shouldBeApproxC`
          A.mkArray @(Complex Double) [4] [4,0,0,0]

      it "normalization factor scales the output" $ do
        A.fft (A.mkArray @(Complex Double) [4] [1,0,0,0]) 2.0 4
          `shouldBeApproxC`
          A.mkArray @(Complex Double) [4] [2,2,2,2]

      it "ifft . fft is the identity" $ do
        let n     = 8
            input = A.mkArray @(Complex Double) [n] (map (:+ 0) [1..8])
        A.ifft (A.fft input 1.0 n) (1.0 / fromIntegral n) n
          `shouldBeApproxC` input

      it "fft output_size pads with zeros when larger than input" $ do
        -- 4-point FFT of a 2-point signal padded to 4: input [1,1,0,0]
        A.fft (A.mkArray @(Complex Double) [2] [1,1]) 1.0 4
          `shouldBeApproxC`
          A.fft (A.mkArray @(Complex Double) [4] [1,1,0,0]) 1.0 4

    describe "fft2" $ do
      it "2D transform of a Dirac delta is a flat spectrum" $ do
        A.fft2 (A.mkArray @(Complex Double) [4,4] (1 : replicate 15 0)) 1.0 4 4
          `shouldBeApproxC`
          A.mkArray @(Complex Double) [4,4] (replicate 16 1)

      it "ifft2 . fft2 is the identity" $ do
        let input = A.mkArray @(Complex Double) [4,4] (map (:+ 0) [1..16])
        A.ifft2 (A.fft2 input 1.0 4 4) (1.0 / 16) 4 4
          `shouldBeApproxC` input

      it "2D transform of all-ones concentrates all energy at DC" $ do
        A.fft2 (A.mkArray @(Complex Double) [4,4] (replicate 16 1)) 1.0 4 4
          `shouldBeApproxC`
          A.mkArray @(Complex Double) [4,4] (16 : replicate 15 0)

    describe "fft2_inplace" $ do
      it "runs without error" $ do
        A.fft2_inplace (A.mkArray @(Complex Double) [4,4] (map (:+ 0) [1..16])) 1.0
          `shouldReturn` ()

    describe "fft3" $ do
      it "3D transform of a Dirac delta is a flat spectrum" $ do
        A.fft3 (A.mkArray @(Complex Double) [4,4,4] (1 : replicate 63 0)) 1.0 4 4 4
          `shouldBeApproxC`
          A.mkArray @(Complex Double) [4,4,4] (replicate 64 1)

      it "ifft3 . fft3 is the identity" $ do
        let input = A.mkArray @(Complex Double) [4,4,4] (map (:+ 0) [1..64])
        A.ifft3 (A.fft3 input 1.0 4 4 4) (1.0 / 64) 4 4 4
          `shouldBeApproxC` input

    describe "fft3_inplace" $ do
      it "runs without error" $ do
        A.fft3_inplace (A.mkArray @(Complex Double) [4,4,4] (map (:+ 0) [1..64])) 1.0
          `shouldReturn` ()

    describe "ifft_inplace" $ do
      it "runs without error" $ do
        A.ifft_inplace (A.mkArray @(Complex Double) [4] (map (:+ 0) [1..4])) 1.0
          `shouldReturn` ()

    describe "ifft2_inplace" $ do
      it "runs without error" $ do
        A.ifft2_inplace (A.mkArray @(Complex Double) [4,4] (map (:+ 0) [1..16])) 1.0
          `shouldReturn` ()

    describe "ifft3_inplace" $ do
      it "runs without error" $ do
        A.ifft3_inplace (A.mkArray @(Complex Double) [4,4,4] (map (:+ 0) [1..64])) 1.0
          `shouldReturn` ()

    describe "fftr2c / fftc2r" $ do
      it "fftr2c output has (n/2+1) complex elements" $ do
        let n = 8
            out = A.fftr2c (A.mkArray @Double [n] [1..8]) 1.0 n
        A.getElements out `shouldBe` (n `div` 2 + 1)

      it "fftc2r recovers even-length real signal" $ do
        let n    = 8
            inp  = A.mkArray @Double [n] [1..8]
            spec' = A.fftr2c inp 1.0 n
            -- norm = 1/n so that r2c * c2r = identity
            out  = A.fftc2r spec' (1.0 / fromIntegral n) False
        out `shouldBeApproxD` inp

      it "fft2r2c output first dim is (n/2+1)" $ do
        let n = 8
            out = A.fft2r2c (A.mkArray @Double [n,n] (replicate (n*n) 1.0)) 1.0 n n
            (d0, _, _, _) = A.getDims out
        d0 `shouldBe` (n `div` 2 + 1)

      it "fft3r2c runs without error" $ do
        let n = 4
            out = A.fft3r2c (A.mkArray @Double [n,n,n] (replicate (n*n*n) 1.0)) 1.0 n n n
        A.getElements out `shouldSatisfy` (> 0)

    describe "approx1" $ do
      it "matches docstring example with Cubic interpolation" $ do
        let input     = A.vector @Float 3 [10,20,30]
            positions = A.vector @Float 5 [0.0, 0.5, 1.0, 1.5, 2.0]
            result    = A.approx1 input positions A.Cubic 0.0
        zipWith (\a e -> abs (a - e))
          (A.toList @Float result)
          (A.toList @Float (A.mkArray @Float [5] [10.0, 13.75, 20.0, 26.25, 30.0]))
          `shouldSatisfy` all (< 1e-4)

      it "Nearest interpolation returns nearest sample value" $ do
        let input     = A.vector @Float 3 [10,20,30]
            positions = A.vector @Float 3 [0.0, 1.0, 2.0]
        zipWith (\a e -> abs (a - e))
          (A.toList @Float (A.approx1 input positions A.Nearest 0.0))
          (A.toList @Float (A.mkArray @Float [3] [10.0, 20.0, 30.0]))
          `shouldSatisfy` all (< 1e-4)

      it "out-of-bounds positions use the fill value" $ do
        let input     = A.vector @Double 3 [10,20,30]
            positions = A.vector @Double 1 [-1.0]
        A.approx1 input positions A.Linear 0.0
          `shouldBeApproxD` A.mkArray @Double [1] [0.0]

    describe "approx2" $ do
      it "matches docstring example with Cubic interpolation" $ do
        let input  = A.matrix @Float (3,3) [[1,1,1],[2,2,2],[3,3,3]]
            pos1   = A.matrix @Float (2,2) [[0.5,1.5],[0.5,1.5]]
            pos2   = A.matrix @Float (2,2) [[0.5,0.5],[1.5,1.5]]
            result = A.approx2 input pos1 pos2 A.Cubic 0.0
        zipWith (\a e -> abs (a - e))
          (A.toList @Float result)
          (A.toList @Float (A.mkArray @Float [2,2] [1.375, 2.625, 1.375, 2.625]))
          `shouldSatisfy` all (< 1e-4)

    describe "convolve1" $ do
      it "convolving with unit delta is identity" $ do
        let sig    = A.mkArray @Double [5] [1,2,3,4,5]
            delta  = A.mkArray @Double [1] [1]
        A.convolve1 sig delta A.ConvDefault A.ConvDomainSpatial
          `shouldBeApproxD` sig

      it "ConvExpand output length is signal_len + filter_len - 1" $ do
        let sig = A.mkArray @Double [5] [1,2,3,4,5]
            flt = A.mkArray @Double [3] [1,0,0]
            out = A.convolve1 sig flt A.ConvExpand A.ConvDomainSpatial
        A.getElements out `shouldBe` 7

      it "ConvDomainAuto matches ConvDomainSpatial result" $ do
        let sig = A.mkArray @Double [8] [1,2,3,4,5,6,7,8]
            flt = A.mkArray @Double [3] [1,2,1]
        A.convolve1 sig flt A.ConvDefault A.ConvDomainAuto
          `shouldBeApproxD`
          A.convolve1 sig flt A.ConvDefault A.ConvDomainSpatial

    describe "convolve2" $ do
      it "convolving with unit 2D delta is identity" $ do
        let img   = A.mkArray @Double [4,4] [1..16]
            delta = A.mkArray @Double [1,1] [1]
        A.convolve2 img delta A.ConvDefault A.ConvDomainSpatial
          `shouldBeApproxD` img

    describe "convolve2Sep" $ do
      it "separable convolution matches full 2D convolution with outer-product kernel" $ do
        let img  = A.mkArray @Double [4,4] [1..16]
            colF = A.mkArray @Double [1] [1]
            rowF = A.mkArray @Double [1] [1]
        A.convolve2Sep colF rowF img A.ConvDefault
          `shouldBeApproxD` img

    describe "fftConvolve2" $ do
      it "result matches spatial convolve2 for a simple kernel" $ do
        let img = A.mkArray @Double [8,8] [1..64]
            flt = A.mkArray @Double [3,3] [0,0,0, 0,1,0, 0,0,0]
        A.fftConvolve2 img flt A.ConvDefault
          `shouldBeApproxD`
          A.convolve2 img flt A.ConvDefault A.ConvDomainSpatial

    describe "fir" $ do
      it "passthrough filter (b=[1]) returns input unchanged" $ do
        let sig = A.mkArray @Double [5] [1,2,3,4,5]
            b   = A.mkArray @Double [1] [1]
        A.fir b sig `shouldBeApproxD` sig

    describe "iir" $ do
      it "all-feedforward / no-feedback is equivalent to FIR" $ do
        let sig = A.mkArray @Double [5] [1,2,3,4,5]
            b   = A.mkArray @Double [1] [1]
            a   = A.mkArray @Double [1] [1]
        A.iir b a sig `shouldBeApproxD` sig

    describe "medFilt1" $ do
      it "constant signal is unchanged by any kernel" $ do
        let sig = A.mkArray @Double [7] (replicate 7 3.0)
        A.medFilt1 sig 3 A.PadZero `shouldBeApproxD` sig

    describe "medFilt2" $ do
      it "constant image is unchanged by any kernel" $ do
        let img = A.mkArray @Double [5,5] (replicate 25 7.0)
        A.medFilt2 img 3 3 A.PadSym `shouldBeApproxD` img

    describe "convolve3" $ do
      it "convolving with unit 3D delta is identity" $ do
        let vol   = A.mkArray @Double [4,4,4] [1..64]
            delta = A.mkArray @Double [1,1,1] [1]
        A.convolve3 vol delta A.ConvDefault A.ConvDomainSpatial
          `shouldBeApproxD` vol

    describe "fft2C2r" $ do
      it "fft2r2c . fft2C2r is the identity for an even-size 2D signal" $ do
        let n   = 8
            inp = A.mkArray @Double [n,n] [1..fromIntegral (n*n)]
            c2r = A.fft2C2r (A.fft2r2c inp 1.0 n n) (1.0 / fromIntegral (n*n)) False
        c2r `shouldBeApproxD` inp

    describe "fft3C2r" $ do
      it "fft3r2c . fft3C2r is the identity for an even-size 3D signal" $ do
        let n   = 4
            inp = A.mkArray @Double [n,n,n] [1..fromIntegral (n*n*n)]
            c2r = A.fft3C2r (A.fft3r2c inp 1.0 n n n) (1.0 / fromIntegral (n*n*n)) False
        c2r `shouldBeApproxD` inp

    describe "setFFTPlanCacheSize" $ do
      it "runs without error" $ do
        A.setFFTPlanCacheSize 4 `shouldReturn` ()

    describe "FFT properties" $ do
      -- ifft . fft = id for arbitrary complex signals of power-of-2 length
      prop "ifft . fft = id (arbitrary complex signal)" $
        forAll (choose (1 :: Int, 6)) $ \k ->
          forAll (vectorOf (2^k) (choose (-10, 10 :: Double))) $ \xs ->
            let n     = 2^k
                input = A.mkArray @(A.Complex Double) [n] (map (:+ 0) xs)
                out   = A.ifft (A.fft input 1.0 n) (1.0 / fromIntegral n) n
            in zipWith (\a e -> magnitude (a - e))
                 (A.toList @(A.Complex Double) out)
                 (A.toList @(A.Complex Double) input)
               `shouldSatisfy` all (< 1e-9)

      -- FFT linearity: fft(a + b) = fft(a) + fft(b)
      prop "fft is linear: fft(a+b) = fft(a) + fft(b)" $
        forAll (choose (1 :: Int, 5)) $ \k ->
          forAll (vectorOf (2^k) (choose (-5, 5 :: Double))) $ \as_ ->
          forAll (vectorOf (2^k) (choose (-5, 5 :: Double))) $ \bs_ ->
            let n   = 2^k
                a   = A.mkArray @(A.Complex Double) [n] (map (:+ 0) as_)
                b   = A.mkArray @(A.Complex Double) [n] (map (:+ 0) bs_)
                lhs = A.toList @(A.Complex Double) (A.fft (a + b) 1.0 n)
                rhs = zipWith (+)
                        (A.toList @(A.Complex Double) (A.fft a 1.0 n))
                        (A.toList @(A.Complex Double) (A.fft b 1.0 n))
            in zipWith (\l r -> magnitude (l - r)) lhs rhs
               `shouldSatisfy` all (< 1e-9)

      -- Parseval's theorem: ||x||^2 = (1/N) * ||X||^2
      prop "Parseval's theorem holds for arbitrary signals" $
        forAll (choose (1 :: Int, 6)) $ \k ->
          forAll (vectorOf (2^k) (choose (-10, 10 :: Double))) $ \xs ->
            let n       = 2^k
                input   = A.mkArray @(A.Complex Double) [n] (map (:+ 0) xs)
                tEnergy = sum (map (\x -> x*x) xs)
                xf      = A.fft input 1.0 n
                fEnergy = (1.0 / fromIntegral n) *
                            sum (map (\c -> realPart c * realPart c + imagPart c * imagPart c)
                                     (A.toList @(A.Complex Double) xf))
            in abs (tEnergy - fEnergy) < 1e-6 + 1e-6 * abs tEnergy

      -- convolve1 with unit delta is identity for arbitrary signals
      prop "convolve1 with unit delta is identity" $ \(NonEmpty xs) ->
        let sig   = A.mkArray @Double [length xs] xs
            delta = A.mkArray @Double [1] [1]
            out   = A.convolve1 sig delta A.ConvDefault A.ConvDomainSpatial
        in zipWith (\a e -> abs (a - e))
             (A.toList @Double out)
             (A.toList @Double sig)
           `shouldSatisfy` all (< 1e-9)

      -- fftr2c . fftc2r round-trip for arbitrary even-length real signals
      prop "fftc2r . fftr2c = id for even-length real signals" $
        forAll (choose (1 :: Int, 5)) $ \k ->
          forAll (vectorOf (2^k) (choose (-10, 10 :: Double))) $ \xs ->
            let n   = 2^k
                inp = A.mkArray @Double [n] xs
                out = A.fftc2r (A.fftr2c inp 1.0 n) (1.0 / fromIntegral n) False
            in zipWith (\a e -> abs (a - e))
                 (A.toList @Double out)
                 xs
               `shouldSatisfy` all (< 1e-9)