Parameters.hs

{-# LANGUAGE DerivingStrategies #-}

-- | The 35 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..w15@ | stability update, long-term block                          |
-- +-----------+------------------------------------------------------------+
-- | @w16..w24@| stability update, short-term (same-day) block              |
-- +-----------+------------------------------------------------------------+
-- | @w25,w26@ | the long-\/short-term transition function                   |
-- +-----------+------------------------------------------------------------+
-- | @w27..w34@| the two-component forgetting curve                         |
-- +-----------+------------------------------------------------------------+
--
-- The two stability blocks have identical shapes, which is why
-- 'longTermWeights' and 'shortTermWeights' both produce a t'StabilityWeights'.
module FSRS.Parameters
  ( -- * Parameters
    Parameters
  , parameterCount
  , defaultParameters
  , mkParameters
  , clampParameters
  , parametersToList
  , parameterAt

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

    -- * Weight blocks
  , initialStabilityWeight
  , initialDifficultyBase
  , initialDifficultyRate
  , difficultyDelta
  , transitionRate
  , transitionAmplitude
  , StabilityWeights (..)
  , longTermWeights
  , shortTermWeights
  , CurveWeights (..)
  , curveWeights
  ) where

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

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)

-- | @35@. FSRS-6 had 21 weights; FSRS-7 adds a second stability block, the
-- transition function and the six extra forgetting-curve weights.
parameterCount :: Int
parameterCount = 35

-- | 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 match @models\/fsrs_v7.py@ in
-- <https://github.com/open-spaced-repetition/srs-benchmark srs-benchmark>.
-- Note that the @Default Parameters@ section of that repository's README still
-- lists @1.15@ for @w15@ and @w24@; that block of the README has not been
-- updated since 2026-03-18, while the model was changed to @1.3@ three days
-- later. The values below follow the model.
defaultParameters :: Parameters
defaultParameters = Parameters (listArray (0, parameterCount - 1) ws)
  where
    ws =
      [ -- Initial stability, indexed by rating - 1
        0.041
      , 2.4175
      , 4.1283
      , 11.9709
      , -- Difficulty
        5.6385
      , 0.4468
      , 3.262
      , -- Stability, long-term block
        2.3054
      , 0.1688
      , 1.3325
      , 0.3524
      , 0.0049
      , 0.7503
      , 0.0896
      , 0.6625
      , 1.3
      , -- Stability, short-term block
        0.882
      , 0.3072
      , 3.5875
      , 0.303
      , 0.0107
      , 0.2279
      , 2.6413
      , 0.5594
      , 1.3
      , -- Long-/short-term transition function
        2.5
      , 1.0
      , -- Forgetting curve
        0.0723
      , 0.1634
      , 0.5
      , 0.9555
      , 0.2245
      , 0.6232
      , 0.1362
      , 0.3862
      ]

-- | 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, long-term block
    (0.0, 4.0)
  , (0.0, 1.2)
  , (0.3, 3.0)
  , (0.01, 1.5)
  , (0.001, 0.9)
  , (0.1, 1.0)
  , (0.0, 3.5)
  , (0.0, 1.0)
  , (1.0, 7.0)
  , -- Stability, short-term block
    (0.0, 4.0)
  , (0.0, 2.0)
  , (0.5, 6.0)
  , (0.001, 1.5)
  , (0.001, 2.0)
  , (0.001, 1.0)
  , (0.0, 5.0)
  , (0.0, 1.0)
  , (1.0, 7.0)
  , -- Transition function
    (2.5, 15.0)
  , (0.0, 1.0)
  , -- Forgetting curve
    (0.01, 0.25)
  , (0.01, 0.95)
  , (0.5, 0.85)
  , (0.5, 0.99)
  , (0.01, 1.0)
  , (0.1, 1.0)
  , (0.0, 0.9)
  , (0.1, 1.1)
  ]

-- | Pairs @(i, j)@ for which @w[i] <= w[j]@ must hold.
orderingConstraints :: [(Int, Int)]
orderingConstraints = [(0, 1), (1, 2), (2, 3), (27, 28), (29, 30)]

-- | 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 = []
    orderErrors = mapMaybe checkOrder orderingConstraints
    checkOrder (i, j)
      | wi <= wj = Nothing
      | otherwise = Just (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: bounds first, in index
-- order, so that the ordering constraints are resolved against already-clamped
-- neighbours. 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) clamped))
  where
    clamped = foldl step [] (zip3 [0 :: Int ..] ws parameterBounds)
    -- `acc` is the already-clamped prefix, in order.
    step acc (i, w, (lo, hi)) = acc <> [clamp lo' hi' w]
      where
        lo' = maximum (lo : [acc !! j | (j, k) <- orderingConstraints, k == i])
        hi' = max lo' hi
    clamp lo hi x
      | isNaN x = lo
      | otherwise = min hi (max lo x)

-- | @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

-- | @w25@: how fast a review stops counting as same-day.
transitionRate :: Parameters -> Double
transitionRate p = parameterAt p 25

-- | @w26@: how much of the short-term behaviour applies at zero elapsed time.
transitionAmplitude :: Parameters -> Double
transitionAmplitude p = parameterAt p 26

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

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

-- | @w7..w15@: the block used for reviews separated by a real interval.
longTermWeights :: Parameters -> StabilityWeights
longTermWeights = stabilityWeightsAt 7

-- | @w16..w24@: the block used for same-day reviews.
shortTermWeights :: Parameters -> StabilityWeights
shortTermWeights = stabilityWeightsAt 16

-- | @w27..w34@: the mixture of two power laws that makes up the FSRS-7
-- forgetting curve.
data CurveWeights = CurveWeights
  { cwDecay1 :: !Double
  -- ^ @negate w27@ — already carries the minus sign the formula wants.
  , cwDecay2 :: !Double
  -- ^ @negate w28@.
  , cwBase1 :: !Double
  -- ^ @w29@: the recall probability of the first component at @t == s@.
  , cwBase2 :: !Double
  -- ^ @w30@: likewise for the second component.
  , cwWeight1 :: !Double
  -- ^ @w31@.
  , cwWeight2 :: !Double
  -- ^ @w32@.
  , cwStabilityPower1 :: !Double
  -- ^ @w33@; applied as @s ** negate cwStabilityPower1@.
  , cwStabilityPower2 :: !Double
  -- ^ @w34@; applied as @s ** cwStabilityPower2@.
  }
  deriving stock (Eq, Show)

-- | Read the forgetting-curve block out of a parameter vector.
curveWeights :: Parameters -> CurveWeights
curveWeights p =
  CurveWeights
    { cwDecay1 = negate (parameterAt p 27)
    , cwDecay2 = negate (parameterAt p 28)
    , cwBase1 = parameterAt p 29
    , cwBase2 = parameterAt p 30
    , cwWeight1 = parameterAt p 31
    , cwWeight2 = parameterAt p 32
    , cwStabilityPower1 = parameterAt p 33
    , cwStabilityPower2 = parameterAt p 34
    }