haskell-fsrs-7.1.0: reference/fsrs7_reference.py
"""Pure-Python transcription of the FSRS-7 memory model.
Transcribed line-by-line from the *finished* FSRS-7, whose reference
implementation is the Rust model the algorithm's author signed off on:
https://github.com/Expertium/fsrs-rs-speed-autoresearch
fsrs-rs/src/model.rs (PARAM_LEN = 34, "DUAL-TRACE ... finished FSRS-7")
fsrs-rs/src/inference.rs (DEFAULT_PARAMETERS)
and its scalar port in the fsrs-rs pull request that upstreams it:
https://github.com/open-spaced-repetition/fsrs-rs/pull/426
src/model_v7.rs, src/parameter_clipper_v7.rs
The reference is written with tensors; this file reproduces the exact same
scalar arithmetic with the standard library only, so that it can be used to
generate golden vectors for the Haskell port without pulling in torch.
What the finished model changed relative to the 35-parameter draft:
* The single memory trace became *two*. A card carries a slow stability and a
fast one, and the forgetting curve is a mixture of a fast and a slow recall
component. This replaces the draft's `transition_function`, which blended a
long- and a short-term stability update by elapsed time; the two weights that
parameterised it are gone.
* Post-lapse stability lost its `d ** -fail_d_exp` factor in both stability
blocks — it is now difficulty-independent — so each block has eight weights
instead of nine.
* The forgetting curve gained three weights that modulate it by difficulty and
stability, and it now depends on difficulty at all, which it never did before.
* A lapse's difficulty step is weighted by how surprising the lapse was.
* Retrievability is rescaled into `[1e-5, 1 - 1e-5]`, so a card reviewed at
`t == 0` has recall probability `1 - 1e-5` rather than exactly `1`.
"""
import math
# Indices 31..33 follow the reference's all-positive storage convention: they
# are stored shifted so that their range starts at zero, and the formulas below
# offset them back (`w - 0.5` for the difficulty weight, `w - 0.3` for the two
# decay modulators).
DEFAULT_PARAMETERS = [
# 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,
# Forgetting-curve modulation by difficulty and stability
0.3734, 0.5333, 0.3048,
]
LOWER_BOUNDS = [
0.0001, 0.0001, 0.0001, 0.0001,
1.0, 0.001, 0.1,
0.0, 0.0, 0.3, 0.01, 0.1, 0.0, 0.0, 1.0,
0.0, 0.0, 0.5, 0.001, 0.001, 0.0, 0.0, 1.0,
0.01, 0.01, 0.2, 0.5, 0.01, 0.1, 0.0, 0.1,
0.0, 0.0, 0.0,
]
UPPER_BOUNDS = [
50.0, 100.0, 100.0, 100.0,
10.0, 4.0, 4.0,
4.0, 1.2, 3.0, 1.5, 1.0, 3.5, 1.0, 7.0,
4.0, 2.0, 6.0, 1.5, 1.0, 5.0, 1.0, 7.0,
0.25, 0.95, 0.85, 0.99, 1.0, 1.0, 0.9, 1.1,
1.0, 0.6, 0.6,
]
STABILITY_MIN = 0.0001 # S_MIN
STABILITY_MAX = 36500.0 # S_MAX
MIN_DIFFICULTY = 1.0 # D_MIN
MAX_DIFFICULTY = 10.0 # D_MAX
MIN_INTERVAL = 1.0 / 86400.0 # one second, in days
MAX_INTERVAL = 36500.0 # one hundred years, in days
# Where each trace's block of eight stability weights starts.
SLOW_TRACE_BASE = 7
FAST_TRACE_BASE = 15
# The fast trace starts at, and after a lapse is capped at, this fraction of
# the slow one.
FAST_TRACE_RATIO = 0.8
# Retrievability is squeezed into [RETENTION_FLOOR, 1 - RETENTION_FLOOR] so
# that neither it nor its logarithm can saturate.
RETENTION_FLOOR = 1e-5
# The exponent of `factor1` is built in log space and capped here so that both
# the value and its gradient stay finite.
LOG_FACTOR_CAP = 60.0
# The magnitude of either decay is held inside this range.
DECAY_MIN = 0.01
DECAY_MAX = 0.95
def clamp(x, lo, hi):
return min(max(x, lo), hi)
# --------------------------------------------------------------------------
# Forgetting curve: a mixture of a fast-trace and a slow-trace power law.
# --------------------------------------------------------------------------
def fast_component_recall(w, t, s_fast):
"""The fast trace's own recall probability after `t` days.
Shared by the forgetting curve, as its fast component, and by the
fast-trace stability update, which reads this rather than the mixture.
"""
t = max(t, 0.0)
s_fast = clamp(s_fast, STABILITY_MIN, STABILITY_MAX)
decay1 = -clamp(w[23] * s_fast ** (w[33] - 0.3), DECAY_MIN, DECAY_MAX)
factor1 = math.exp(min(math.log(w[25]) / decay1, LOG_FACTOR_CAP)) - 1.0
return (1.0 + factor1 * (t / s_fast)) ** decay1
def _curve_parts(w, t, state):
"""The pieces of the mixture, shared by the curve and its derivative."""
t = max(t, 0.0)
s = clamp(state[0], STABILITY_MIN, STABILITY_MAX)
d = clamp(state[1], MIN_DIFFICULTY, MAX_DIFFICULTY)
s_fast = clamp(state[2], STABILITY_MIN, STABILITY_MAX)
# Fast component: decay modulated by the fast stability itself.
decay1 = -clamp(w[23] * s_fast ** (w[33] - 0.3), DECAY_MIN, DECAY_MAX)
factor1 = math.exp(min(math.log(w[25]) / decay1, LOG_FACTOR_CAP)) - 1.0
inner1 = 1.0 + factor1 * (t / s_fast)
# Slow component: difficulty rescales *time* rather than the decay, so a
# hard card experiences time faster but decays with the same slope.
decay2 = -clamp(w[24], DECAY_MIN, DECAY_MAX)
factor2 = w[26] ** (1.0 / decay2) - 1.0
d_timescale = math.exp((d - 5.0) * (w[32] - 0.3))
inner2 = 1.0 + factor2 * d_timescale * (t / s)
# Mixture weights, one keyed to each trace; the slow one is also modulated
# by difficulty.
weight1 = w[27] * s_fast ** -w[29]
weight2 = w[28] * s ** w[30] * math.exp((d - 5.0) * (w[31] - 0.5))
return (
decay1, factor1, inner1, s_fast,
decay2, factor2 * d_timescale, inner2, s,
weight1, weight2,
)
def forgetting_curve(w, t, state):
"""Probability of recall after `t` days, given a `(S, D, S_fast)` state."""
(
decay1, _factor1, inner1, _s_fast,
decay2, _factor2, inner2, _s,
weight1, weight2,
) = _curve_parts(w, t, state)
r1 = inner1 ** decay1
r2 = inner2 ** decay2
retention = (weight1 * r1 + weight2 * r2) / (weight1 + weight2)
return retention * (1.0 - 2.0 * RETENTION_FLOOR) + RETENTION_FLOOR
def forgetting_curve_derivative(w, t, state):
"""dR/dt of `forgetting_curve`; always <= 0 for in-bounds parameters."""
(
decay1, factor1, inner1, s_fast,
decay2, factor2, inner2, s,
weight1, weight2,
) = _curve_parts(w, t, state)
d1 = decay1 * inner1 ** (decay1 - 1.0) * (factor1 / s_fast)
d2 = decay2 * inner2 ** (decay2 - 1.0) * (factor2 / s)
derivative = (weight1 * d1 + weight2 * d2) / (weight1 + weight2)
return derivative * (1.0 - 2.0 * RETENTION_FLOOR)
def next_interval(w, desired_retention, state):
"""Invert the forgetting curve: the t with R(t, state) = desired_retention.
The mixture has no closed-form inverse. R is strictly decreasing in t, so a
bisection on [MIN_INTERVAL, MAX_INTERVAL] converges to full double
precision; this is the value the Haskell root-finder is checked against.
"""
lo, hi = MIN_INTERVAL, MAX_INTERVAL
if forgetting_curve(w, lo, state) <= desired_retention:
return lo
if forgetting_curve(w, hi, state) >= desired_retention:
return hi
for _ in range(200):
mid = math.sqrt(lo * hi) # bisect in log space
if mid <= lo or mid >= hi:
break
if forgetting_curve(w, mid, state) > desired_retention:
lo = mid
else:
hi = mid
return math.sqrt(lo * hi)
# --------------------------------------------------------------------------
# Difficulty
# --------------------------------------------------------------------------
def initial_difficulty(w, rating):
"""Unclamped; `step` clamps the result to [1, 10]."""
return w[4] - math.exp(w[5] * (rating - 1)) + 1.0
def linear_damping(delta_difficulty, difficulty):
return delta_difficulty * (10.0 - difficulty) / 9.0
def mean_reversion(init, current):
return 0.01 * init + 0.99 * current
def next_difficulty(w, difficulty, rating, retrievability):
"""Difficulty after a review.
A lapse's step is weighted by how surprising the lapse was: forgetting a
card the model expected you to recall says more about the card's difficulty
than forgetting one that was long overdue.
"""
delta_d = -w[6] * (rating - 3)
if rating == 1:
delta_d *= retrievability + 0.1
new_d = difficulty + linear_damping(delta_d, difficulty)
reverted = mean_reversion(initial_difficulty(w, 4), new_d)
return clamp(reverted, MIN_DIFFICULTY, MAX_DIFFICULTY)
# --------------------------------------------------------------------------
# Stability
# --------------------------------------------------------------------------
def stability_after_review(w, s, d, r, rating, base):
"""One trace's stability update; `base` is SLOW_TRACE_BASE or FAST_TRACE_BASE.
Post-lapse stability is difficulty-independent in the finished model: the
`d ** -fail_d_exp` factor of the draft was ablated.
"""
w_sinc_base = w[base]
w_sinc_s_exp = w[base + 1]
w_sinc_r_mult = w[base + 2]
w_fail_mult = w[base + 3]
w_fail_s_exp = w[base + 4]
w_fail_r_mult = w[base + 5]
w_hard = w[base + 6]
w_easy = w[base + 7]
hard_penalty = w_hard if rating == 2 else 1.0
easy_bonus = w_easy if rating == 4 else 1.0
new_s_fail = (
w_fail_mult
* ((s + 1.0) ** w_fail_s_exp - 1.0)
* math.exp((1.0 - r) * w_fail_r_mult)
)
pls = min(s, new_s_fail)
s_inc = 1.0 + (
math.exp(w_sinc_base - 1.5)
* (11.0 - d)
* s ** -w_sinc_s_exp
* (math.exp((1.0 - r) * w_sinc_r_mult) - 1.0)
* hard_penalty
* easy_bonus
)
new_s_success = max(pls, s * s_inc)
result = new_s_success if rating > 1 else pls
return clamp(result, STABILITY_MIN, STABILITY_MAX)
def next_state(w, state, delta_t, rating):
"""Both traces and the difficulty, after a review of an existing card.
The slow trace updates from the mixed retrievability; the fast trace
updates from its *own* recall component, and on a lapse is pulled down to
at most `FAST_TRACE_RATIO` of the post-lapse slow stability.
"""
delta_t = max(delta_t, 0.0)
s, d, s_fast = state
r = forgetting_curve(w, delta_t, state)
new_s_long = stability_after_review(w, s, d, r, rating, SLOW_TRACE_BASE)
r_fast = fast_component_recall(w, delta_t, s_fast)
new_s_short = stability_after_review(w, s_fast, d, r_fast, rating, FAST_TRACE_BASE)
if rating == 1:
new_s_short = min(new_s_short, FAST_TRACE_RATIO * new_s_long)
new_d = next_difficulty(w, d, rating, r)
return (new_s_long, new_d, new_s_short)
# --------------------------------------------------------------------------
# The state transition
# --------------------------------------------------------------------------
def initial_state(w, rating):
"""The state a card is born with, given its first rating."""
s = clamp(w[rating - 1], STABILITY_MIN, STABILITY_MAX)
d = clamp(initial_difficulty(w, rating), MIN_DIFFICULTY, MAX_DIFFICULTY)
s_fast = clamp(FAST_TRACE_RATIO * s, STABILITY_MIN, STABILITY_MAX)
return (s, d, s_fast)
def step(w, state, delta_t, rating):
"""`state` is None for the very first review, else an (S, D, S_fast) triple."""
if state is None:
return initial_state(w, rating)
s, d, s_fast = next_state(w, state, delta_t, rating)
return (
clamp(s, STABILITY_MIN, STABILITY_MAX),
clamp(d, MIN_DIFFICULTY, MAX_DIFFICULTY),
clamp(s_fast, STABILITY_MIN, STABILITY_MAX),
)
def replay(w, reviews):
"""Fold `step` over a list of (delta_t, rating) pairs."""
state = None
for delta_t, rating in reviews:
state = step(w, state, delta_t, rating)
return state
# --------------------------------------------------------------------------
# Parameter clipping
# --------------------------------------------------------------------------
# Pairs (i, j) for which w[i] <= w[j] must hold. The reference restores these
# *after* the per-parameter box clamps, by pulling the larger index up.
ORDERING_CONSTRAINTS = [(0, 1), (1, 2), (2, 3), (25, 26)]
def clip_parameters(w):
"""The reference clipper: box clamps first, then monotonicity."""
out = [clamp(v, lo, hi) for v, lo, hi in zip(w, LOWER_BOUNDS, UPPER_BOUNDS)]
for lower, upper in ORDERING_CONSTRAINTS:
out[upper] = max(out[upper], out[lower])
return out