Algorithm.hs

{-# LANGUAGE BangPatterns #-}

-- | The FSRS-7 memory model.
--
-- This module is a direct transcription of the finished FSRS-7, whose
-- reference implementation is the Rust model its author signed off on:
-- <https://github.com/Expertium/fsrs-rs-speed-autoresearch> (see
-- @fsrs-rs\/src\/model.rs@), upstreamed in
-- <https://github.com/open-spaced-repetition/fsrs-rs/pull/426 fsrs-rs#426>.
-- Every function here is pure and total.
--
-- What changed in FSRS-7 relative to FSRS-6:
--
-- * A card carries /two/ memory traces instead of one — see t'MemoryState'.
--   The forgetting curve is a mixture of a fast and a slow component, one per
--   trace, and each trace is updated by its own block of weights. This is what
--   makes same-day reviews behave differently from spaced ones; FSRS-6
--   switched between two formulas on a same-day\/not-same-day flag.
-- * The forgetting curve depends on /difficulty/, which it never did before:
--   a hard card experiences time faster, and counts the slow trace for more.
-- * Post-lapse stability no longer depends on difficulty.
-- * A lapse's difficulty step is weighted by how surprising the lapse was —
--   see 'nextDifficulty'.
-- * Retrievability is squeezed into @[1e-5, 1 - 1e-5]@, so a card reviewed a
--   moment ago has a retrievability just short of @1@.
-- * 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
  , fastTraceRetrievability
  , nextIntervalDays

    -- * Difficulty
  , initialDifficulty
  , nextDifficulty

    -- * Stability
  , initialStability
  , stabilityAfterReview
  , nextStability

    -- * State transition
  , initialMemoryState
  , nextMemoryState
  , replayReviews

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

import FSRS.Parameters
  ( CurveWeights (..)
  , Parameters
  , StabilityWeights (..)
  , curveWeights
  , difficultyDelta
  , fastTraceWeights
  , initialDifficultyBase
  , initialDifficultyRate
  , initialStabilityWeight
  , slowTraceWeights
  )
import FSRS.Types
  ( Days
  , Difficulty
  , MemoryState (..)
  , Rating (..)
  , Retrievability
  , Stability
  , ratingToInt
  )

-- | The smallest stability the model will report, @1e-4@ days (about nine
-- seconds). FSRS-7 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

-- | @1e-5@. 'retrievability' is rescaled into
-- @['retrievabilityFloor', 1 - 'retrievabilityFloor']@ so that neither it nor
-- its logarithm can saturate.
retrievabilityFloor :: Retrievability
retrievabilityFloor = 1.0e-5

-- | @0.8@. The fast trace starts at this fraction of the slow one, and after a
-- lapse is pulled back down to at most this fraction of it.
fastTraceRatio :: Double
fastTraceRatio = 0.8

-- | The magnitude of either decay is held inside @[0.01, 0.95]@.
decayMin, decayMax :: Double
decayMin = 0.01
decayMax = 0.95

-- | The exponent that builds the fast component's shape factor is capped here,
-- in log space, so that both the value and its gradient stay finite.
logFactorCap :: Double
logFactorCap = 60.0

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

clampStability :: Stability -> Stability
clampStability = clampTo stabilityMin stabilityMax

clampDifficulty :: Difficulty -> Difficulty
clampDifficulty = clampTo difficultyMin difficultyMax

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

-- | The pieces of the two-component mixture, shared by 'retrievability' and
-- 'retrievabilityDerivative' so that the two cannot drift apart.
--
-- Each component is @inner ** decay@ where @inner = 1 + rate * t@, which makes
-- both the value and its slope one-liners.
data CurveParts = CurveParts
  { cpDecay1 :: !Double
  , cpRate1 :: !Double
  , cpInner1 :: !Double
  , cpDecay2 :: !Double
  , cpRate2 :: !Double
  , cpInner2 :: !Double
  , cpWeight1 :: !Double
  , cpWeight2 :: !Double
  }

curveParts :: Parameters -> Days -> MemoryState -> CurveParts
curveParts params t0 state =
  CurveParts
    { cpDecay1 = decay1
    , cpRate1 = rate1
    , cpInner1 = 1 + rate1 * t
    , cpDecay2 = decay2
    , cpRate2 = rate2
    , cpInner2 = 1 + rate2 * t
    , cpWeight1 = cwWeight1 cw * sFast ** negate (cwStabilityPower1 cw)
    , cpWeight2 =
        cwWeight2 cw
          * s ** cwStabilityPower2 cw
          * exp ((d - 5) * (cwDifficultyWeight cw - 0.5))
    }
  where
    cw = curveWeights params
    t = max 0 t0
    s = clampStability (memoryStability state)
    d = clampDifficulty (memoryDifficulty state)
    sFast = clampStability (memoryStabilityFast state)

    -- Fast component: the decay is scaled by the fast stability itself.
    decay1 = fastDecay cw sFast
    rate1 = shapeFactor1 cw decay1 / sFast

    -- Slow component: difficulty rescales *time* rather than the decay, so a
    -- hard card experiences time faster but decays with the same slope.
    decay2 = negate (clampTo decayMin decayMax (cwDecay2 cw))
    factor2 = cwBase2 cw ** (1 / decay2) - 1
    timescale = exp ((d - 5) * (cwDifficultyDecay cw - 0.3))
    rate2 = factor2 * timescale / s

fastDecay :: CurveWeights -> Stability -> Double
fastDecay cw sFast =
  negate . clampTo decayMin decayMax $
    cwDecayBase1 cw * sFast ** (cwStabilityDecay1 cw - 0.3)

-- | @base1 ** (1 \/ decay1) - 1@, built in log space and capped so that a
-- decay near zero cannot overflow it.
shapeFactor1 :: CurveWeights -> Double -> Double
shapeFactor1 cw decay1 = exp (min logFactorCap (log (cwBase1 cw) / decay1)) - 1

-- | Rescale a raw mixture value into
-- @['retrievabilityFloor', 1 - 'retrievabilityFloor']@.
rescale :: Double -> Double
rescale x = x * (1 - 2 * retrievabilityFloor) + retrievabilityFloor

-- | The probability of recalling a card @t@ days after the last review.
--
-- FSRS-7 mixes a fast-trace and a slow-trace power law, so this needs the
-- whole memory state — both stabilities /and/ the difficulty. The weight of
-- each component depends on its own stability, so — unlike every earlier
-- version — the retrievability at @t == s@ is not a fixed @0.9@.
--
-- @t@ is treated as @0@ if negative; the state is clamped into range, so any
-- finite state is safe.
retrievability :: Parameters -> Days -> MemoryState -> Retrievability
retrievability params t state = rescale ((w1 * r1 + w2 * r2) / (w1 + w2))
  where
    cp = curveParts params t state
    r1 = cpInner1 cp ** cpDecay1 cp
    r2 = cpInner2 cp ** cpDecay2 cp
    w1 = cpWeight1 cp
    w2 = cpWeight2 cp

-- | @d\/dt@ of 'retrievability'. Never positive for in-bounds parameters.
retrievabilityDerivative :: Parameters -> Days -> MemoryState -> Double
retrievabilityDerivative params t state =
  (w1 * d1 + w2 * d2) / (w1 + w2) * (1 - 2 * retrievabilityFloor)
  where
    cp = curveParts params t state
    d1 = cpDecay1 cp * cpInner1 cp ** (cpDecay1 cp - 1) * cpRate1 cp
    d2 = cpDecay2 cp * cpInner2 cp ** (cpDecay2 cp - 1) * cpRate2 cp
    w1 = cpWeight1 cp
    w2 = cpWeight2 cp

-- | The fast trace's /own/ recall probability — the mixture's fast component
-- on its own, unrescaled and with no contribution from the slow trace.
--
-- This is what drives the fast trace's stability update in 'nextStability':
-- the fast trace responds to how well /it/ was doing, not to the mixture.
fastTraceRetrievability :: Parameters -> Days -> Stability -> Retrievability
fastTraceRetrievability params t0 sFast0 = (1 + rate * t) ** decay
  where
    cw = curveWeights params
    t = max 0 t0
    sFast = clampStability sFast0
    decay = fastDecay cw sFast
    rate = shapeFactor1 cw decay / sFast

-- | The interval, in days, after which a card in the given state 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), 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. Since retrievability tops out just below @1@, a target of @1@
-- saturates rather than diverging.
nextIntervalDays :: Parameters -> Retrievability -> MemoryState -> Days
nextIntervalDays params target state
  | not (target > 0) = maximumIntervalDays -- also catches NaN
  | target >= 1 = minimumIntervalDays
  | curve lo <= target = lo
  | curve hi >= target = hi
  | otherwise = exp (search (0 :: Int) (log lo) (log hi) u0)
  where
    curve t = retrievability params t state
    lo = minimumIntervalDays
    hi = maximumIntervalDays
    -- R around the slow stability is near the interesting range, so start
    -- there.
    u0 = clampTo (log lo) (log hi) (log (clampStability (memoryStability state)))

    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 = curve t - 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 state * 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 = clampDifficulty . 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.
--
-- A lapse's step is weighted by how surprising the lapse was, which is why
-- this takes the retrievability at review time: forgetting a card the model
-- expected you to recall says more about the card than forgetting one that was
-- long overdue.
nextDifficulty
  :: Parameters
  -> Difficulty
  -> Retrievability
  -- ^ The retrievability at review time, from 'retrievability'.
  -> Rating
  -> Difficulty
nextDifficulty params d r rating = clampDifficulty (meanReversion damped)
  where
    step = negate (difficultyDelta params) * fromIntegral (ratingToInt rating - 3)
    delta = if rating == Again then step * (r + 0.1) else step
    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 trace's stability update.
--
-- Called twice per review — once with 'FSRS.Parameters.slowTraceWeights' and
-- once with 'FSRS.Parameters.fastTraceWeights' — on that trace's own stability
-- and its own retrievability.
--
-- On a lapse the new stability is the post-lapse stability, which can never
-- exceed the old one and, in the finished FSRS-7, does not depend on
-- difficulty. 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
  -> Stability
  -- ^ The stability of the trace being updated.
  -> Difficulty
  -> Retrievability
  -- ^ That trace's retrievability at review time.
  -> Rating
  -> Stability
stabilityAfterReview w s d r rating
  | rating == Again = clampStability postLapse
  | otherwise = clampStability (max postLapse (s * increase))
  where
    postLapse = min s failureStability
    failureStability =
      swFailureFactor 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

-- | Both stabilities after a review, as @(slow, fast)@.
--
-- The slow trace is updated from the mixed 'retrievability'; the fast trace
-- from its own 'fastTraceRetrievability'. On a lapse the fast trace is pulled
-- down to at most 'fastTraceRatio' of the post-lapse slow stability, so that
-- forgetting a card cannot leave it with a large short-term boost.
nextStability
  :: Parameters
  -> MemoryState
  -> Days
  -- ^ Days since the previous review.
  -> Rating
  -> (Stability, Stability)
nextStability params state deltaT rating = (slow, fast)
  where
    (slow, fast, _) = advance params state deltaT rating

-- | The whole update in one pass: both stabilities and the mixed
-- retrievability that the difficulty step also needs, so that the forgetting
-- curve is evaluated exactly once per review.
advance
  :: Parameters
  -> MemoryState
  -> Days
  -> Rating
  -> (Stability, Stability, Retrievability)
advance params state deltaT rating = (slow, fast, r)
  where
    t = max 0 deltaT
    d = clampDifficulty (memoryDifficulty state)
    r = retrievability params t state
    slow =
      stabilityAfterReview
        (slowTraceWeights params)
        (clampStability (memoryStability state))
        d
        r
        rating
    rFast = fastTraceRetrievability params t (memoryStabilityFast state)
    raw =
      stabilityAfterReview
        (fastTraceWeights params)
        (clampStability (memoryStabilityFast state))
        d
        rFast
        rating
    fast
      | rating == Again = min raw (fastTraceRatio * slow)
      | otherwise = raw

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

-- | The state a card is born with, given the rating of its first review. The
-- fast trace starts at 'fastTraceRatio' of the slow one.
initialMemoryState :: Parameters -> Rating -> MemoryState
initialMemoryState params rating =
  MemoryState
    { memoryStability = clampStability s
    , memoryDifficulty = initialDifficulty params rating
    , memoryStabilityFast = clampStability (fastTraceRatio * s)
    }
  where
    s = initialStability params rating

-- | 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 comes from 'initialMemoryState'.
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 = case before of
  Nothing -> initialMemoryState params rating
  Just state ->
    MemoryState
      { memoryStability = clampStability slow
      , memoryDifficulty =
          nextDifficulty params (clampDifficulty (memoryDifficulty state)) r rating
      , memoryStabilityFast = clampStability fast
      }
    where
      (slow, fast, r) = advance params state deltaT 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