Algorithm.hs

{-# LANGUAGE BangPatterns #-}

-- | The FSRS-7 memory model.
--
-- This module is a direct transcription of the reference implementation in
-- <https://github.com/open-spaced-repetition/srs-benchmark srs-benchmark>
-- (@models\/fsrs_v7.py@ and @models\/fsrs_v7_interval_penalty.py@). Every
-- function here is pure and total.
--
-- What changed in FSRS-7 relative to FSRS-6:
--
-- * The forgetting curve is a stability-weighted mixture of /two/ power laws
--   instead of one, which is why it has eight weights and no closed-form
--   inverse (see 'nextIntervalDays').
-- * The stability update is computed twice — once with a long-term weight
--   block and once with a short-term one — and the two are blended by
--   'transitionCoefficient', a smooth function of the elapsed time. FSRS-6
--   switched between two separate formulas on a same-day\/not-same-day flag.
-- * Intervals are genuinely continuous. Ten minutes is @10 \/ 1440@ days and
--   the model is meant to be evaluated at such values, not at whole days.
module FSRS.Algorithm
  ( -- * Forgetting curve
    retrievability
  , retrievabilityDerivative
  , nextIntervalDays

    -- * Difficulty
  , initialDifficulty
  , nextDifficulty

    -- * Stability
  , initialStability
  , stabilityAfterReview
  , transitionCoefficient
  , nextStability

    -- * State transition
  , nextMemoryState
  , replayReviews

    -- * Bounds
  , stabilityMin
  , stabilityMax
  , difficultyMin
  , difficultyMax
  , minimumIntervalDays
  , maximumIntervalDays
  ) where

import FSRS.Parameters
  ( CurveWeights (..)
  , Parameters
  , StabilityWeights (..)
  , curveWeights
  , difficultyDelta
  , initialDifficultyBase
  , initialDifficultyRate
  , initialStabilityWeight
  , longTermWeights
  , shortTermWeights
  , transitionAmplitude
  , transitionRate
  )
import FSRS.Types
  ( Days
  , Difficulty
  , MemoryState (..)
  , Rating (..)
  , Retrievability
  , Stability
  , ratingToInt
  )

-- | The smallest stability the model will report, @1e-4@ days (about nine
-- seconds). This is the @s_min@ the upstream configuration uses for FSRS-7,
-- which is always run with second-precision intervals.
stabilityMin :: Stability
stabilityMin = 1.0e-4

-- | The largest stability the model will report: @36500@ days, a century.
stabilityMax :: Stability
stabilityMax = 36500.0

-- | @1@.
difficultyMin :: Difficulty
difficultyMin = 1.0

-- | @10@.
difficultyMax :: Difficulty
difficultyMax = 10.0

-- | One second, expressed in days — the shortest interval 'nextIntervalDays'
-- will return.
minimumIntervalDays :: Days
minimumIntervalDays = 1.0 / 86400.0

-- | A century in days — the longest interval 'nextIntervalDays' will return.
maximumIntervalDays :: Days
maximumIntervalDays = 36500.0

clampTo :: Double -> Double -> Double -> Double
clampTo lo hi x = min hi (max lo x)

-- ---------------------------------------------------------------------------
-- Forgetting curve
-- ---------------------------------------------------------------------------

-- | The probability of recalling a card @t@ days after the last review, given
-- its stability.
--
-- FSRS-7 mixes two power laws whose relative weight depends on the stability
-- itself, so — unlike every earlier version — the retrievability at @t == s@
-- is not a fixed @0.9@ but drifts with @s@.
--
-- @t@ must be non-negative and @s@ strictly positive.
retrievability :: Parameters -> Days -> Stability -> Retrievability
retrievability params t s = (weight1 * r1 + weight2 * r2) / (weight1 + weight2)
  where
    cw = curveWeights params
    tOverS = t / s
    powerLaw base decay = (1 + factor * tOverS) ** decay
      where
        factor = base ** (1 / decay) - 1
    r1 = powerLaw (cwBase1 cw) (cwDecay1 cw)
    r2 = powerLaw (cwBase2 cw) (cwDecay2 cw)
    weight1 = cwWeight1 cw * s ** negate (cwStabilityPower1 cw)
    weight2 = cwWeight2 cw * s ** cwStabilityPower2 cw

-- | @d\/dt@ of 'retrievability'. Never positive for in-bounds parameters.
retrievabilityDerivative :: Parameters -> Days -> Stability -> Double
retrievabilityDerivative params t s =
  (weight1 * d1 + weight2 * d2) / (weight1 + weight2)
  where
    cw = curveWeights params
    tOverS = t / s
    slope base decay = decay * inner ** (decay - 1) * (factor / s)
      where
        factor = base ** (1 / decay) - 1
        inner = 1 + factor * tOverS
    d1 = slope (cwBase1 cw) (cwDecay1 cw)
    d2 = slope (cwBase2 cw) (cwDecay2 cw)
    weight1 = cwWeight1 cw * s ** negate (cwStabilityPower1 cw)
    weight2 = cwWeight2 cw * s ** cwStabilityPower2 cw

-- | The interval, in days, after which a card of the given stability will have
-- decayed to exactly the desired retention.
--
-- The FSRS-7 forgetting curve has no closed-form inverse, so this is a
-- root-find: Newton's method in @log t@ (which is what makes the problem
-- well-conditioned — see the upstream @fsrs_v7_interval_penalty@ module),
-- safeguarded by a bracketing bisection so that it cannot diverge.
--
-- The result is always within @['minimumIntervalDays', 'maximumIntervalDays']@;
-- a desired retention that is unreachable within that window saturates at the
-- nearer end.
nextIntervalDays :: Parameters -> Retrievability -> Stability -> Days
nextIntervalDays params target s
  | not (target > 0) = maximumIntervalDays -- also catches NaN
  | target >= 1 = minimumIntervalDays
  | retrievability params lo s <= target = lo
  | retrievability params hi s >= target = hi
  | otherwise = exp (search (0 :: Int) (log lo) (log hi) u0)
  where
    lo = minimumIntervalDays
    hi = maximumIntervalDays
    -- R(s, s) is near the interesting range, so start there.
    u0 = clampTo (log lo) (log hi) (log s)

    maxIterations = 200 :: Int
    -- Full double precision in log space.
    tolerance = 1.0e-15

    search !n !a !b !u
      | n >= maxIterations = u
      | fu == 0 = u
      | converged = next
      | otherwise = search (n + 1) a' b' next
      where
        t = exp u
        fu = retrievability params t s - target
        -- R is strictly decreasing in t, so the sign of `fu` says which side
        -- of the root `u` is on.
        (a', b') = if fu > 0 then (u, b) else (a, u)
        slope = retrievabilityDerivative params t s * t
        newton = u - fu / slope
        next
          | isNaN newton || isInfinite newton || newton <= a' || newton >= b' =
              0.5 * (a' + b')
          | otherwise = newton
        converged = abs (next - u) <= tolerance * max 1 (abs u)

-- ---------------------------------------------------------------------------
-- Difficulty
-- ---------------------------------------------------------------------------

-- | Difficulty before clamping. The mean reversion in 'nextDifficulty' pulls
-- towards the /unclamped/ value for 'Easy', which is why this is kept
-- separate.
rawInitialDifficulty :: Parameters -> Rating -> Double
rawInitialDifficulty params rating =
  initialDifficultyBase params
    - exp (initialDifficultyRate params * fromIntegral (ratingToInt rating - 1))
    + 1

-- | The difficulty a card is born with, given the rating of its first review.
initialDifficulty :: Parameters -> Rating -> Difficulty
initialDifficulty params =
  clampTo difficultyMin difficultyMax . rawInitialDifficulty params

-- | Difficulty after a review: a rating-driven step, damped so that it slows
-- down as difficulty approaches its maximum, then reverted 1% of the way
-- towards the difficulty an 'Easy' first review would have produced.
nextDifficulty :: Parameters -> Difficulty -> Rating -> Difficulty
nextDifficulty params d rating =
  clampTo difficultyMin difficultyMax (meanReversion damped)
  where
    delta = negate (difficultyDelta params) * fromIntegral (ratingToInt rating - 3)
    damped = d + delta * (10 - d) / 9
    meanReversion current = 0.01 * rawInitialDifficulty params Easy + 0.99 * current

-- ---------------------------------------------------------------------------
-- Stability
-- ---------------------------------------------------------------------------

-- | The stability a card is born with, given the rating of its first review.
initialStability :: Parameters -> Rating -> Stability
initialStability = initialStabilityWeight

-- | One half of the stability update.
--
-- Called twice per review — once with 'FSRS.Parameters.longTermWeights' and
-- once with 'FSRS.Parameters.shortTermWeights' — and the two results are
-- blended by 'nextStability'.
--
-- On a lapse the new stability is the post-lapse stability, which can never
-- exceed the old one. On a success it is the old stability scaled by a factor
-- that grows with how overdue the card was, shrinks as the card gets easier
-- and more stable, and is scaled again by the hard penalty or the easy bonus.
stabilityAfterReview
  :: StabilityWeights
  -> MemoryState
  -> Retrievability
  -- ^ The retrievability at review time, from 'retrievability'.
  -> Rating
  -> Stability
stabilityAfterReview w (MemoryState s d) r rating
  | rating == Again = postLapse
  | otherwise = max postLapse (s * increase)
  where
    postLapse = min s failureStability
    failureStability =
      swFailureFactor w
        * d ** negate (swFailureDifficultyExponent w)
        * ((s + 1) ** swFailureStabilityExponent w - 1)
        * exp ((1 - r) * swFailureRetrievabilityFactor w)
    hardPenalty = if rating == Hard then swHardPenalty w else 1
    easyBonus = if rating == Easy then swEasyBonus w else 1
    increase =
      1
        + exp (swIncreaseBase w - 1.5)
          * (11 - d)
          * s ** negate (swIncreaseStabilityExponent w)
          * (exp ((1 - r) * swIncreaseRetrievabilityFactor w) - 1)
          * hardPenalty
          * easyBonus

-- | How much of a long-term review this is: @0@ for a review that happens at
-- the same instant as the previous one, tending to @1@ as the gap grows.
--
-- This is the piece that replaces FSRS-6's hard same-day\/not-same-day split.
-- The elapsed time is expected to be non-negative.
transitionCoefficient :: Parameters -> Days -> Double
transitionCoefficient params deltaT =
  1 - transitionAmplitude params * exp (negate (transitionRate params) * deltaT)

-- | Stability after a review, blending the long- and short-term updates.
nextStability
  :: Parameters
  -> MemoryState
  -> Days
  -- ^ Days since the previous review.
  -> Rating
  -> Stability
nextStability params state deltaT rating =
  coefficient * longTerm + (1 - coefficient) * shortTerm
  where
    r = retrievability params deltaT (memoryStability state)
    longTerm = stabilityAfterReview (longTermWeights params) state r rating
    shortTerm = stabilityAfterReview (shortTermWeights params) state r rating
    coefficient = transitionCoefficient params deltaT

-- ---------------------------------------------------------------------------
-- State transition
-- ---------------------------------------------------------------------------

-- | Advance a card's memory state by one review.
--
-- Pass 'Nothing' for a card that has never been reviewed; the elapsed time is
-- then ignored and the state is read straight off the initial-stability and
-- initial-difficulty weights.
nextMemoryState
  :: Parameters
  -> Maybe MemoryState
  -- ^ The state before the review, or 'Nothing' for a brand-new card.
  -> Days
  -- ^ Days since the previous review.
  -> Rating
  -> MemoryState
nextMemoryState params before deltaT rating =
  MemoryState
    { memoryStability = clampTo stabilityMin stabilityMax s
    , memoryDifficulty = clampTo difficultyMin difficultyMax d
    }
  where
    (s, d) = case before of
      Nothing ->
        ( initialStability params rating
        , initialDifficulty params rating
        )
      Just state ->
        ( nextStability params state deltaT rating
        , nextDifficulty params (memoryDifficulty state) rating
        )

-- | Fold a whole review history into a memory state.
--
-- Each element is @(days since the previous review, rating)@; the elapsed time
-- of the first review is ignored. 'Nothing' for an empty history.
replayReviews :: Parameters -> [(Days, Rating)] -> Maybe MemoryState
replayReviews params = go Nothing
  where
    go before [] = before
    go before ((deltaT, rating) : rest) =
      let next = nextMemoryState params before deltaT rating
       in next `seq` go (Just next) rest