Golden.hs

-- | Every function of the model, checked against vectors produced by the
-- pure-Python transcription of the reference implementation.
--
-- The two implementations perform the same floating-point operations in the
-- same order on the same libm, so agreement is expected to the last few bits;
-- the tolerances below are deliberately much tighter than "close enough".
module Test.FSRS.Golden (tests) where

import Control.Monad (unless)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (Assertion, assertFailure, testCase, (@?=))

import FSRS
import Test.FSRS.Gen (approxEqual, relativeError)
import Test.FSRS.GoldenData

-- | Tolerance for everything but the root-finder: bit-for-bit agreement with a
-- little slack for the last ulp.
tol :: Double
tol = 1.0e-12

-- | The interval solver converges by a different route than the Python
-- bisection used to produce the vectors, so it only agrees to about a
-- nanosecond in relative terms.
intervalTol :: Double
intervalTol = 1.0e-9

tests :: TestTree
tests =
  testGroup
    "golden vectors"
    [ testCase "the golden parameter sets are all valid" $
        mapM_
          (\ws -> validateParameters ws @?= [])
          goldenParameterSets
    , testCase "the first golden parameter set is the default one" $
        take 1 goldenParameterSets @?= [parametersToList defaultParameters]
    , labelled "forgetting curve" goldenCurveVectors checkCurve
    , labelled "fast-trace recall" goldenFastRecallVectors checkFastRecall
    , labelled "difficulty" goldenDifficultyVectors checkDifficulty
    , labelled "stability blocks" goldenHalfStabilityVectors checkHalfStability
    , labelled "memory-state transition" goldenStepVectors checkStep
    , labelled "interval inversion" goldenIntervalVectors checkInterval
    , labelled "replayed review histories" goldenReplayVectors checkReplay
    ]

labelled :: String -> [a] -> (a -> Assertion) -> TestTree
labelled name vectors check =
  testCase (name <> " (" <> show (length vectors) <> " vectors)") $
    mapM_ check vectors

paramsAt :: Int -> Parameters
paramsAt i = case mkParameters (goldenParameterSets !! i) of
  Right p -> p
  Left errs -> error ("golden parameter set " <> show i <> " is invalid: " <> show errs)

ratingAt :: Int -> Rating
ratingAt n = case ratingFromInt n of
  Just r -> r
  Nothing -> error ("golden vector has a bad rating: " <> show n)

-- | The vectors spell states out as @(S, D, S_fast)@ triples.
stateOf :: (Double, Double, Double) -> MemoryState
stateOf (s, d, sFast) = MemoryState s d sFast

close :: Double -> String -> Double -> Double -> Assertion
close t label expected actual =
  unless (approxEqual t t actual expected) $
    assertFailure $
      label
        <> "\n  expected: "
        <> show expected
        <> "\n  actual:   "
        <> show actual
        <> "\n  relative error: "
        <> show (relativeError actual expected)

closeState :: Double -> String -> (Double, Double, Double) -> MemoryState -> Assertion
closeState t label (s, d, sFast) actual = do
  close t (label <> " [stability]") s (memoryStability actual)
  close t (label <> " [difficulty]") d (memoryDifficulty actual)
  close t (label <> " [fast stability]") sFast (memoryStabilityFast actual)

checkCurve :: CurveVector -> Assertion
checkCurve v =
  close tol (show v) (cvRetrievability v) $
    retrievability (paramsAt (cvParams v)) (cvElapsedDays v) (stateOf (cvState v))

checkFastRecall :: FastRecallVector -> Assertion
checkFastRecall v =
  close tol (show v) (frRecall v) $
    fastTraceRetrievability (paramsAt (frParams v)) (frElapsedDays v) (frStabilityFast v)

checkDifficulty :: DifficultyVector -> Assertion
checkDifficulty v =
  close tol (show v) (dvNextDifficulty v) $
    case dvDifficulty v of
      Nothing -> initialDifficulty params rating
      Just d -> nextDifficulty params d (dvRetrievability v) rating
  where
    params = paramsAt (dvParams v)
    rating = ratingAt (dvRating v)

checkHalfStability :: HalfStabilityVector -> Assertion
checkHalfStability v =
  close tol (show v) (hsNextStability v) $
    stabilityAfterReview
      block
      (hsStability v)
      (hsDifficulty v)
      (hsRetrievability v)
      (ratingAt (hsRating v))
  where
    params = paramsAt (hsParams v)
    block = (if hsSlowTrace v then slowTraceWeights else fastTraceWeights) params

checkStep :: StepVector -> Assertion
checkStep v = closeState tol (show v) (svNextState v) $
  nextMemoryState
    (paramsAt (svParams v))
    (stateOf <$> svState v)
    (svElapsedDays v)
    (ratingAt (svRating v))

checkInterval :: IntervalVector -> Assertion
checkInterval v =
  close intervalTol (show v) (ivInterval v) $
    nextIntervalDays (paramsAt (ivParams v)) (ivDesiredRetention v) (stateOf (ivState v))

checkReplay :: ReplayVector -> Assertion
checkReplay v = case replayReviews params reviews of
  Nothing -> assertFailure (show v <> ": replayReviews returned Nothing")
  Just actual -> closeState tol (show v) (rvFinalState v) actual
  where
    params = paramsAt (rvParams v)
    reviews = [(dt, ratingAt g) | (dt, g) <- rvReviews v]