Parameters.hs

{-# LANGUAGE DerivingStrategies #-}

-- | The 34 weights of FSRS-7, and typed views onto the blocks they form.
--
-- FSRS-7 groups its weights like this:
--
-- +-----------+------------------------------------------------------------+
-- | @w0..w3@  | initial stability, indexed by rating                       |
-- +-----------+------------------------------------------------------------+
-- | @w4..w6@  | difficulty                                                 |
-- +-----------+------------------------------------------------------------+
-- | @w7..w14@ | stability update, slow trace                               |
-- +-----------+------------------------------------------------------------+
-- | @w15..w22@| stability update, fast trace                               |
-- +-----------+------------------------------------------------------------+
-- | @w23..w30@| the two-component forgetting curve                         |
-- +-----------+------------------------------------------------------------+
-- | @w31..w33@| how difficulty and stability modulate that curve           |
-- +-----------+------------------------------------------------------------+
--
-- The two stability blocks have identical shapes, which is why
-- 'slowTraceWeights' and 'fastTraceWeights' both produce a t'StabilityWeights'.
module FSRS.Parameters
  ( -- * Parameters
    Parameters
  , parameterCount
  , defaultParameters
  , mkParameters
  , clampParameters
  , parametersToList
  , parameterAt

    -- * Validation
  , ParameterError (..)
  , parameterBounds
  , orderingConstraints
  , validateParameters

    -- * Weight blocks
  , initialStabilityWeight
  , initialDifficultyBase
  , initialDifficultyRate
  , difficultyDelta
  , StabilityWeights (..)
  , slowTraceWeights
  , fastTraceWeights
  , CurveWeights (..)
  , curveWeights
  ) where

import Data.Array.Unboxed (UArray, bounds, elems, listArray, (!))

import FSRS.Types (Rating, ratingToInt)

-- | An FSRS-7 parameter vector: exactly 'parameterCount' weights.
--
-- Build one with 'mkParameters' or 'clampParameters'; the constructor is
-- hidden so that the length invariant cannot be broken.
newtype Parameters = Parameters (UArray Int Double)
  deriving stock (Eq, Ord)

instance Show Parameters where
  showsPrec d p =
    showParen (d > 10) $
      showString "mkParameters " . showsPrec 11 (parametersToList p)

-- | @34@. FSRS-6 had 21 weights. FSRS-7 keeps two stability blocks of eight
-- weights each — one per memory trace — and spends eleven on a forgetting
-- curve that mixes the two traces and is modulated by difficulty.
parameterCount :: Int
parameterCount = 34

-- | The 0-based weight at the given index. Indices outside
-- @[0, 'parameterCount')@ are a programmer error and raise an exception.
parameterAt :: Parameters -> Int -> Double
parameterAt (Parameters a) i
  | i >= lo && i <= hi = a ! i
  | otherwise =
      error $
        "FSRS.Parameters.parameterAt: index " <> show i <> " out of range " <> show (lo, hi)
  where
    (lo, hi) = bounds a

-- | The weights as a plain list, in index order.
parametersToList :: Parameters -> [Double]
parametersToList (Parameters a) = elems a

-- | The default FSRS-7 parameters, obtained by the upstream authors through
-- multi-user optimisation.
--
-- These are the @DEFAULT_PARAMETERS@ of the finished FSRS-7 reference
-- implementation, <https://github.com/Expertium/fsrs-rs-speed-autoresearch>.
defaultParameters :: Parameters
defaultParameters = Parameters (listArray (0, parameterCount - 1) ws)
  where
    ws =
      [ -- Initial stability, indexed by rating - 1
        0.1104
      , 2.2395
      , 3.9221
      , 11.7841
      , -- Difficulty
        6.1686
      , 0.6457
      , 3.6807
      , -- Stability, slow trace
        1.9795
      , 0.0
      , 1.3826
      , 0.7024
      , 0.5999
      , 0.8146
      , 0.6398
      , 1.0
      , -- Stability, fast trace
        1.3207
      , 0.6707
      , 3.8668
      , 0.4416
      , 0.0934
      , 1.8631
      , 0.6162
      , 1.0869
      , -- Forgetting curve
        0.1567
      , 0.0801
      , 0.2421
      , 0.9464
      , 0.1433
      , 0.7145
      , 0.0
      , 0.5667
      , -- Modulation of the curve by difficulty and stability
        0.3734
      , 0.5333
      , 0.3048
      ]

-- | Why a list of weights is not a valid parameter vector.
data ParameterError
  = -- | Got this many weights instead of 'parameterCount'.
    WrongParameterCount !Int
  | -- | @index@, @value@, @lower bound@, @upper bound@.
    ParameterOutOfBounds !Int !Double !Double !Double
  | -- | @w[i] <= w[j]@ is required but was violated: @i@, @j@, @w[i]@, @w[j]@.
    ParameterOutOfOrder !Int !Int !Double !Double
  | -- | @index@, the offending @NaN@ or infinity.
    ParameterNotFinite !Int !Double
  deriving stock (Eq, Show)

-- | The inclusive @(lower, upper)@ bound of every weight, in index order.
--
-- Taken from the parameter clipper the upstream optimiser applies after each
-- gradient step, so any parameter set produced by a real optimiser satisfies
-- them.
parameterBounds :: [(Double, Double)]
parameterBounds =
  [ -- Initial stability
    (1.0e-4, 50.0)
  , (1.0e-4, 100.0)
  , (1.0e-4, 100.0)
  , (1.0e-4, 100.0)
  , -- Difficulty
    (1.0, 10.0)
  , (0.001, 4.0)
  , (0.1, 4.0)
  , -- Stability, slow trace
    (0.0, 4.0)
  , (0.0, 1.2)
  , (0.3, 3.0)
  , (0.01, 1.5)
  , (0.1, 1.0)
  , (0.0, 3.5)
  , (0.0, 1.0)
  , (1.0, 7.0)
  , -- Stability, fast trace
    (0.0, 4.0)
  , (0.0, 2.0)
  , (0.5, 6.0)
  , (0.001, 1.5)
  , (0.001, 1.0)
  , (0.0, 5.0)
  , (0.0, 1.0)
  , (1.0, 7.0)
  , -- Forgetting curve
    (0.01, 0.25)
  , (0.01, 0.95)
  , (0.2, 0.85)
  , (0.5, 0.99)
  , (0.01, 1.0)
  , (0.1, 1.0)
  , (0.0, 0.9)
  , (0.1, 1.1)
  , -- Modulation of the curve
    (0.0, 1.0)
  , (0.0, 0.6)
  , (0.0, 0.6)
  ]

-- | Pairs @(i, j)@ for which @w[i] <= w[j]@ must hold: the initial stabilities
-- rise with the rating, and the slow component of the forgetting curve starts
-- above the fast one.
orderingConstraints :: [(Int, Int)]
orderingConstraints = [(0, 1), (1, 2), (2, 3), (25, 26)]

-- | Every problem with a candidate weight list, or @[]@ if there is none.
validateParameters :: [Double] -> [ParameterError]
validateParameters ws
  | n /= parameterCount = [WrongParameterCount n]
  | otherwise = boundErrors <> orderErrors
  where
    n = length ws
    boundErrors = concat (zipWith3 check [0 ..] ws parameterBounds)
    check i w (lo, hi)
      | isNaN w || isInfinite w = [ParameterNotFinite i w]
      | w < lo || w > hi = [ParameterOutOfBounds i w lo hi]
      | otherwise = []
    -- Only meaningful once every weight is known to be finite; otherwise a
    -- NaN would be reported twice, once as itself and once as out of order.
    orderErrors
      | null boundErrors = concatMap checkOrder orderingConstraints
      | otherwise = []
    checkOrder (i, j)
      | wi <= wj = []
      | otherwise = [ParameterOutOfOrder i j wi wj]
      where
        wi = ws !! i
        wj = ws !! j

-- | Build a parameter vector, rejecting anything out of bounds.
mkParameters :: [Double] -> Either [ParameterError] Parameters
mkParameters ws = case validateParameters ws of
  [] -> Right (Parameters (listArray (0, parameterCount - 1) ws))
  errs -> Left errs

-- | Build a parameter vector by pulling every weight into its valid range,
-- exactly as the upstream optimiser's clipper does: every weight is clamped to
-- its own bounds first, and only then are the ordering constraints restored by
-- raising the /larger/ index. Only a wrong number of weights can still fail.
clampParameters :: [Double] -> Either [ParameterError] Parameters
clampParameters ws
  | length ws /= parameterCount = Left [WrongParameterCount (length ws)]
  | otherwise = Right (Parameters (listArray (0, parameterCount - 1) ordered))
  where
    boxed = zipWith clampInto ws parameterBounds
    clampInto x (lo, hi)
      | isNaN x = lo
      | otherwise = min hi (max lo x)
    -- The constraints are listed in increasing order and chain through the
    -- initial stabilities, so a single left-to-right pass settles them.
    ordered = foldl raise boxed orderingConstraints
    raise acc (i, j) = setAt j (max (acc !! j) (acc !! i)) acc
    setAt i x xs = take i xs <> [x] <> drop (i + 1) xs

-- | @w[rating - 1]@: the stability a card is born with after its first review.
initialStabilityWeight :: Parameters -> Rating -> Double
initialStabilityWeight p r = parameterAt p (ratingToInt r - 1)

-- | @w4@.
initialDifficultyBase :: Parameters -> Double
initialDifficultyBase p = parameterAt p 4

-- | @w5@.
initialDifficultyRate :: Parameters -> Double
initialDifficultyRate p = parameterAt p 5

-- | @w6@.
difficultyDelta :: Parameters -> Double
difficultyDelta p = parameterAt p 6

-- | One of the two eight-weight blocks that drive the stability update, one
-- per memory trace.
data StabilityWeights = StabilityWeights
  { swIncreaseBase :: !Double
  -- ^ @w7@ \/ @w15@, used as @exp (base - 1.5)@.
  , swIncreaseStabilityExponent :: !Double
  -- ^ @w8@ \/ @w16@.
  , swIncreaseRetrievabilityFactor :: !Double
  -- ^ @w9@ \/ @w17@.
  , swFailureFactor :: !Double
  -- ^ @w10@ \/ @w18@.
  , swFailureStabilityExponent :: !Double
  -- ^ @w11@ \/ @w19@.
  , swFailureRetrievabilityFactor :: !Double
  -- ^ @w12@ \/ @w20@.
  , swHardPenalty :: !Double
  -- ^ @w13@ \/ @w21@.
  , swEasyBonus :: !Double
  -- ^ @w14@ \/ @w22@.
  }
  deriving stock (Eq, Show)

stabilityWeightsAt :: Int -> Parameters -> StabilityWeights
stabilityWeightsAt base p =
  StabilityWeights
    { swIncreaseBase = at 0
    , swIncreaseStabilityExponent = at 1
    , swIncreaseRetrievabilityFactor = at 2
    , swFailureFactor = at 3
    , swFailureStabilityExponent = at 4
    , swFailureRetrievabilityFactor = at 5
    , swHardPenalty = at 6
    , swEasyBonus = at 7
    }
  where
    at k = parameterAt p (base + k)

-- | @w7..w14@: the block that updates the slow trace.
slowTraceWeights :: Parameters -> StabilityWeights
slowTraceWeights = stabilityWeightsAt 7

-- | @w15..w22@: the block that updates the fast trace.
fastTraceWeights :: Parameters -> StabilityWeights
fastTraceWeights = stabilityWeightsAt 15

-- | @w23..w33@: the mixture of two power laws that makes up the FSRS-7
-- forgetting curve, and the three weights that modulate it.
--
-- The last three are stored the way the reference stores them — shifted so
-- that their range starts at zero — and the formulas in "FSRS.Algorithm"
-- offset them back. The field documentation gives the offset each one carries.
data CurveWeights = CurveWeights
  { cwDecayBase1 :: !Double
  -- ^ @w23@: sets the fast component's decay, which is also scaled by the
  -- fast stability, so unlike every earlier version of FSRS the decay is not
  -- a constant.
  , cwDecay2 :: !Double
  -- ^ @w24@: the slow component's decay magnitude.
  , cwBase1 :: !Double
  -- ^ @w25@: the fast component's recall probability at @t == s_fast@.
  , cwBase2 :: !Double
  -- ^ @w26@: likewise for the slow component at @t == s@.
  , cwWeight1 :: !Double
  -- ^ @w27@: how much the fast component counts in the mixture.
  , cwWeight2 :: !Double
  -- ^ @w28@: likewise for the slow component.
  , cwStabilityPower1 :: !Double
  -- ^ @w29@; applied as @s_fast ** negate cwStabilityPower1@.
  , cwStabilityPower2 :: !Double
  -- ^ @w30@; applied as @s ** cwStabilityPower2@.
  , cwDifficultyWeight :: !Double
  -- ^ @w31@; the mixture weight of the slow component is scaled by
  -- @exp ((d - 5) * (cwDifficultyWeight - 0.5))@.
  , cwDifficultyDecay :: !Double
  -- ^ @w32@; the slow component sees time scaled by
  -- @exp ((d - 5) * (cwDifficultyDecay - 0.3))@, so a hard card experiences
  -- time faster while decaying with the same slope.
  , cwStabilityDecay1 :: !Double
  -- ^ @w33@; the fast component's decay is scaled by
  -- @s_fast ** (cwStabilityDecay1 - 0.3)@.
  }
  deriving stock (Eq, Show)

-- | Read the forgetting-curve block out of a parameter vector.
curveWeights :: Parameters -> CurveWeights
curveWeights p =
  CurveWeights
    { cwDecayBase1 = parameterAt p 23
    , cwDecay2 = parameterAt p 24
    , cwBase1 = parameterAt p 25
    , cwBase2 = parameterAt p 26
    , cwWeight1 = parameterAt p 27
    , cwWeight2 = parameterAt p 28
    , cwStabilityPower1 = parameterAt p 29
    , cwStabilityPower2 = parameterAt p 30
    , cwDifficultyWeight = parameterAt p 31
    , cwDifficultyDecay = parameterAt p 32
    , cwStabilityDecay1 = parameterAt p 33
    }