HABQT-0.1.0.0: test/ParticleProcessingTests.hs

{-# LANGUAGE NamedFieldPuns #-}

module ParticleProcessingTests
  ( testParticleHandling
  ) where

import Control.Applicative (liftA2)
import Data.Complex (Complex(..))
import Data.List (sort)
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.Data.Particle
import HABQTlib.RandomStates
import qualified Numeric.LinearAlgebra as LA
import StateGenTests (arbPropDensityMatrix)
import qualified Streaming.Prelude as S
import qualified System.Random.MWC as MWC
import qualified Test.QuickCheck as QC
import Test.QuickCheck
  ( Args(..)
  , Positive(..)
  , Property
  , (.&&.)
  , arbitrary
  , quickCheckWith
  , stdArgs
  , suchThat
  )
import TestHelpers

lik :: Double -> Double
lik p = p / (1 - p)

loglik :: Double -> Double
loglik = log . lik

likProp :: PureStateVector -> DensityMatrix -> Property
likProp sv dm =
  let l = pureStateLikelihood sv dm
   in QC.property (l >= 0) .&&. QC.property (l <= 1)

testLikelihood :: IO ()
testLikelihood = do
  putStrLn "Testing likelihood computation for density matrices:"
  quickCheckWith stdArgs {maxSuccess = 1000} likProp

expProp :: QC.Blind Particles -> Property
expProp (QC.Blind pts) =
  let dm = snd . getWDM . V.foldl1' (<+>) . ptsParticles $ pts
      dim = LA.rows . getDensityMatrix $ dm
      z = (dim LA.>< dim) $ repeat 0
      scaleDM :: Weight -> DensityMatrix -> LA.Matrix (Complex Double)
      scaleDM w (DensityMatrix dm'') = LA.scale (w :+ 0) dm''
      dm' = DensityMatrix $ foldOverPts id scaleDM (+) z pts
      matchProp = QC.property $ dm <==> dm'
      validProp = arbPropDensityMatrix dm'
   in matchProp .&&. validProp

testExpectation :: IO ()
testExpectation = do
  putStrLn "Testing particle expectation function:"
  quickCheckWith stdArgs {maxSuccess = 500} expProp

particleUpdateFidR :: PureStateVector -> Particles -> Double
particleUpdateFidR testState@(PureStateVector sv) particles =
  let testPure = DensityMatrix $ sv LA.<> LA.tr sv
      originalEstimate = snd . getWDM $ reduceParticlesToMean particles
      originalFid = fidelityDM testPure originalEstimate
      updatedParticles = updateParticles testState particles
      updatedEstimate = snd . getWDM $ reduceParticlesToMean updatedParticles
      updatedFid = fidelityDM testPure updatedEstimate
   in loglik updatedFid - loglik originalFid

updateProp :: Property
updateProp =
  let gen = do
        (dim, rank) <- arbDimRank
        Positive num <- arbitrary :: QC.Gen (Positive NumberOfParticles)
        stateVector <- QC.resize dim arbitrary
        particles <- arbParticles dim rank (fromIntegral num)
        return (stateVector, particles)
      prop (sv, pts) =
        let wtsf Particles {ptsParticles} =
              V.foldl1' (+) . V.map (fst . getWDM) $ ptsParticles
            ws = wtsf pts
            pts' = updateParticles sv pts
            ws' = wtsf pts'
            wsProp = abs (ws - 1) < 1e-10
            ws'Prop = abs (ws' - 1) < 1e-10
         in wsProp .&&. ws'Prop
      msg = "of updates move the estimate closer"
   in QC.forAll
        gen
        (\(sv, pts) ->
           QC.classify (particleUpdateFidR sv pts >= 0) msg (prop (sv, pts)))

ecdfProp :: Particles -> Property
ecdfProp Particles {ptsParticles} =
  let cdf = ecdf ptsParticles
      hp = V.head cdf > 0
      tp = abs (V.last cdf - 1) < 1e-10
      cdfSorted = V.fromList . sort . V.toList $ cdf
      sorted = cdf == cdfSorted
   in hp .&&. tp .&&. sorted

icdfProp :: Particles -> Property
icdfProp Particles {ptsParticles} =
  QC.ioProperty $ do
    gen <- MWC.createSystemRandom
    x <- MWC.uniform gen
    let cdf = ecdf ptsParticles
        idx = icdf cdf x
        Just idx' = V.findIndex (> x) cdf
    return (idx == idx')

hDistProp :: ParticleHierarchy -> PureStateVector -> Double
hDistProp ph sv =
  let dm = DensityMatrix $ getStateVector sv LA.<> LA.tr (getStateVector sv)
      originalEstimate = getMixedEstimate ph
      updatedEstimate = getMixedEstimate $ updateParticleHierarchy sv ph
      originalFid = fidelityDM dm originalEstimate
      updatedFid = fidelityDM dm updatedEstimate
   in loglik updatedFid - loglik originalFid

hierarchyBatchProp ::
     Int -> Positive Dim -> Positive NumberOfParticles -> IO Property
hierarchyBatchProp batchSize (Positive dim) (Positive num) = do
  ph <- initialiseParticleHierarchy dim num
  passed <-
    S.length_ . S.filter (>= 0) . S.replicateM batchSize $
    hDistProp ph <$> genPureSV dim
  return . QC.property $
    fromIntegral passed / (fromIntegral batchSize :: Double) >= 0.95

hierarchyUpdateProp :: Int -> Property
hierarchyUpdateProp batchSize =
  let gen = do
        dim <- suchThat arbitrary (\x -> x > 1 && x < 20)
        num <- suchThat arbitrary (> 100)
        return (Positive dim, Positive num)
   in QC.forAll
        gen
        (\(dim, num) -> QC.ioProperty $ hierarchyBatchProp batchSize dim num)

effectiveSampleSizeProp :: Property
effectiveSampleSizeProp =
  let genP = do
        (dim, rank) <- arbDimRank
        Positive num <- arbitrary :: QC.Gen (Positive NumberOfParticles)
        arbParticles dim rank (fromIntegral num)
      sizeP ps@(Particles _ _ n _) =
        QC.classify
          (es <= fromIntegral n / 2)
          "of effective sizes are less than half of total particle number"
          ((es <= fromIntegral n) && (es >= 0))
        where
          es = effectiveSize ps
   in QC.forAll genP sizeP

nudgeProp :: Property
nudgeProp =
  let gen = do
        (dim, rank) <- arbDimRank
        dm <- arbDM dim rank
        return (dim, rank, dm)
      nudged (dim, rank, dm) =
        truncateRank rank <$> nudgeParticle dim 1e-2 (mkWDM1 dm)
      prop (dim, rank, dm) =
        QC.ioProperty $ do
          WeighedDensityMatrix (w, ndm) <- nudged (dim, rank, dm)
          let validDM = arbPropDensityMatrix ndm
              preservesRank = getRank dm == getRank ndm
              preservesWeight = w == 1
          return $
            QC.classify
              (fidelityDM ndm dm > 0.99)
              "fidelities exceed 99%"
              (validDM .&&. preservesRank .&&. preservesWeight)
   in QC.forAll gen prop

type ResamplingFunIO = MWC.GenIO -> Particles -> IO Particles

resampleProp :: MWC.GenIO -> ResamplingFunIO -> Property
resampleProp gen resamplingFun =
  let genA = do
        dim <- suchThat (arbitrary :: QC.Gen Dim) (liftA2 (&&) (< 10) (> 1))
        Positive rank <-
          suchThat (arbitrary :: QC.Gen (Positive Rank)) (< Positive dim)
        num <- suchThat arbitrary (> 1000)
        QC.Blind <$> arbParticles dim rank num
      preservesMeanIO :: QC.Blind Particles -> IO Property
      preservesMeanIO (QC.Blind pts) = do
        rpts <- resamplingFun gen pts
        let WeighedDensityMatrix (w1, dm1) = reduceParticlesToMean pts
            WeighedDensityMatrix (w2, dm2) = reduceParticlesToMean rpts
            sumw wacc (WeighedDensityMatrix (w', _)) = wacc + w'
            getpw Particles {ptsParticles = ps} = V.foldl' sumw 0 ps
            w1' = getpw pts
            w2' = getpw rpts
            propWeight = QC.property $ 2 * abs (w1 - w2) / (w1 + w2) < 1e-12
            propWeight' =
              QC.property $ 2 * abs (w1' - w2') / (w1' + w2') < 1e-12
            Particles _ _ num _ = rpts
            propESabs =
              QC.property $ abs (effectiveSize rpts - fromIntegral num) < 1
            propESincrease =
              QC.property $ effectiveSize rpts > effectiveSize pts
            propAll =
              propWeight .&&. propWeight' .&&. propESincrease .&&. propESabs
        return $
          QC.classify
            (fidelityDM dm1 dm2 > 0.99)
            "fidelities between pre- and post-resampling means exceed 99%"
            propAll
   in QC.forAll genA (QC.ioProperty . preservesMeanIO)

testParticleResampling :: IO ()
testParticleResampling = do
  gen <- MWC.createSystemRandom
  putStrLn "Testing multinomial particle resampling without checks:"
  quickCheckWith stdArgs {maxSuccess = 100} $ resampleProp gen resampleMultinom

testParticleNudging :: IO ()
testParticleNudging = do
  putStrLn "Testing particle nudging:"
  quickCheckWith stdArgs {maxSuccess = 1000} nudgeProp

testEffectiveSampleSize :: IO ()
testEffectiveSampleSize = do
  putStrLn "Testing calculation of effective sample size:"
  quickCheckWith stdArgs {maxSuccess = 1000} effectiveSampleSizeProp

testParticleUpdate :: IO ()
testParticleUpdate = do
  putStrLn "Testing update of particles:"
  quickCheckWith stdArgs {maxSuccess = 1000} updateProp

testEcdf :: IO ()
testEcdf = do
  putStrLn "Testing ecdf calculation:"
  quickCheckWith stdArgs {maxSuccess = 10000} ecdfProp
  putStrLn "Testing iecdf calculation:"
  quickCheckWith stdArgs {maxSuccess = 10000} icdfProp

testHierarchyUpdate :: IO ()
testHierarchyUpdate = do
  putStrLn "Testing update of particle hierarchies:"
  quickCheckWith stdArgs {maxSuccess = 50} (hierarchyUpdateProp 100)

testParticleHandling :: IO ()
testParticleHandling = do
  testLikelihood
  testExpectation
  testParticleUpdate
  testEcdf
  testHierarchyUpdate
  testEffectiveSampleSize
  testParticleNudging
  testParticleResampling