haskell-fsrs-7.1.0: reference/gen_golden.py
"""Generate test/Test/FSRS/GoldenData.hs from the Python FSRS-7 reference.
Run from the repository root:
python3 reference/gen_golden.py
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fsrs7_reference import ( # noqa: E402
DEFAULT_PARAMETERS,
FAST_TRACE_BASE,
FAST_TRACE_RATIO,
LOWER_BOUNDS,
SLOW_TRACE_BASE,
STABILITY_MAX,
STABILITY_MIN,
UPPER_BOUNDS,
clamp,
clip_parameters,
fast_component_recall,
forgetting_curve,
initial_difficulty,
next_difficulty,
next_interval,
stability_after_review,
step,
)
OUT = os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"test",
"Test",
"FSRS",
"GoldenData.hs",
)
def hs(x):
"""Render a Python float as a Haskell `Double` literal."""
x = float(x)
s = repr(x)
if "e" in s or "E" in s:
mantissa, exponent = s.replace("E", "e").split("e")
if "." not in mantissa:
mantissa += ".0"
return f"{mantissa}e{int(exponent)}"
if "." not in s:
s += ".0"
return s
def hs_list(xs):
return "[" + ", ".join(xs) + "]"
def hs_triple(t):
return "(" + ", ".join(hs(v) for v in t) + ")"
# --------------------------------------------------------------------------
# Alternative parameter sets: a deterministic LCG walk over the valid box.
# --------------------------------------------------------------------------
def make_parameter_set(seed):
state = seed
w = []
for lo, hi in zip(LOWER_BOUNDS, UPPER_BOUNDS):
state = (state * 6364136223846793005 + 1442695040888963407) % (2**64)
u = (state >> 11) / float(2**53)
w.append(lo + u * (hi - lo))
return [round(v, 6) for v in clip_parameters(w)]
PARAMETER_SETS = [
DEFAULT_PARAMETERS,
make_parameter_set(1),
make_parameter_set(2026),
]
TIMES = [0.0, 1.0 / 86400.0, 1.0 / 1440.0, 10.0 / 1440.0, 0.25, 1.0, 3.0, 7.5, 30.0, 365.0, 3650.0]
STABILITIES = [0.0001, 0.01, 0.5, 1.0, 3.9221, 15.0, 100.0, 1000.0, 36500.0]
DIFFICULTIES = [1.0, 2.5, 3.5307239812763354, 7.0, 10.0]
RATINGS = [1, 2, 3, 4]
RETRIEVABILITIES = [0.05, 0.5, 0.95]
RETENTIONS = [0.7, 0.8, 0.85, 0.9, 0.95, 0.97, 0.99]
HALF_DIFFICULTIES = [1.0, 3.5307239812763354, 10.0]
HALF_STABILITIES = [0.01, 1.0, 3.9221, 1000.0]
FAST_STABILITIES = [0.0001, 0.01, 1.0, 3.9221, 100.0, 36500.0]
def fast_of(s, ratio):
return clamp(s * ratio, STABILITY_MIN, STABILITY_MAX)
# States for the forgetting curve: the whole stability range at three
# difficulties with the natural fast/slow ratio, plus a block that moves the
# fast trace independently of the slow one.
CURVE_STATES = [
(s, d, fast_of(s, FAST_TRACE_RATIO))
for s in STABILITIES
for d in (1.0, 5.0, 10.0)
] + [
(s, d, fast_of(s, ratio))
for s in (0.5, 3.9221, 100.0)
for d in (2.5, 7.0)
for ratio in (0.05, 1.0, 20.0)
]
INTERVAL_STATES = [
(s, d, fast_of(s, FAST_TRACE_RATIO)) for s in STABILITIES for d in (1.0, 5.0, 10.0)
] + [(3.9221, 5.0, fast_of(3.9221, ratio)) for ratio in (0.05, 1.0, 20.0)]
STEP_STATES = [
(s, d, fast_of(s, ratio))
for s in (0.0001, 1.0, 3.9221, 1000.0, 36500.0)
for d in (1.0, 3.5307239812763354, 10.0)
for ratio in (FAST_TRACE_RATIO, 5.0)
]
STEP_DELTAS = [0.0, 10.0 / 1440.0, 0.5, 1.0, 45.0, 400.0]
REVIEW_SEQUENCES = [
[(0.0, 3)],
[(0.0, 1)],
[(0.0, 4)],
[(0.0, 3), (10.0 / 1440.0, 3)],
[(0.0, 1), (1.0 / 1440.0, 3), (10.0 / 1440.0, 3), (1.0, 3)],
[(0.0, 3), (1.0, 3), (3.0, 3), (8.0, 3), (21.0, 3)],
# The fixture the reference implementation's own test suite asserts on.
[(0.0, 1), (0.0, 3), (1.0, 3), (3.0, 3), (8.0, 3), (21.0, 3)],
[(0.0, 2), (1.0, 2), (2.0, 2), (3.0, 2)],
[(0.0, 4), (15.0, 4), (90.0, 4), (365.0, 4)],
[(0.0, 3), (5.0, 1), (10.0 / 1440.0, 3), (1.0, 3), (4.0, 4)],
[(0.0, 1), (0.0, 1), (0.0, 1), (0.0, 3)],
[(0.0, 3), (100.0, 1), (0.5, 2), (2.0, 3), (6.0, 4), (30.0, 1)],
[(0.0, 2), (0.25, 3), (0.75, 4), (2.5, 1), (0.01, 3), (7.0, 3)],
[(0.0, 3)] + [(2.0 ** k, 3) for k in range(10)],
[(0.0, 4)] + [(1.0, r) for r in (4, 3, 2, 1, 3, 4, 2, 1)],
]
def render():
"""Build the contents of GoldenData.hs and return it as a string."""
lines = []
add = lines.append
add("{-# LANGUAGE DerivingStrategies #-}")
add("")
add("-- This module is nothing but a few thousand literals; optimising it")
add("-- costs compile time and buys nothing.")
add("{-# OPTIONS_GHC -O0 #-}")
add("")
add("-- | Golden vectors for the FSRS-7 implementation.")
add("--")
add("-- Generated by @reference\\/gen_golden.py@ from the pure-Python")
add("-- transcription of the reference model. Do not edit by hand;")
add("-- regenerate with @python3 reference\\/gen_golden.py@.")
add("module Test.FSRS.GoldenData")
add(" ( goldenParameterSets")
add(" , CurveVector (..)")
add(" , goldenCurveVectors")
add(" , FastRecallVector (..)")
add(" , goldenFastRecallVectors")
add(" , DifficultyVector (..)")
add(" , goldenDifficultyVectors")
add(" , HalfStabilityVector (..)")
add(" , goldenHalfStabilityVectors")
add(" , StepVector (..)")
add(" , goldenStepVectors")
add(" , IntervalVector (..)")
add(" , goldenIntervalVectors")
add(" , ReplayVector (..)")
add(" , goldenReplayVectors")
add(" ) where")
add("")
add("-- | The parameter sets the vectors below refer to by index.")
add("-- Index 0 is the FSRS-7 default parameter set.")
add("goldenParameterSets :: [[Double]]")
add("goldenParameterSets =")
for i, w in enumerate(PARAMETER_SETS):
prefix = " [ " if i == 0 else " , "
add(prefix + hs_list([hs(v) for v in w]))
add(" ]")
add("")
# ---------------------------------------------------------------- curve
add("-- | @R(t, state)@, the probability of recall. The state is spelled out")
add("-- as a @(stability, difficulty, fast stability)@ triple.")
add("data CurveVector = CurveVector")
add(" { cvParams :: !Int")
add(" , cvElapsedDays :: !Double")
add(" , cvState :: !(Double, Double, Double)")
add(" , cvRetrievability :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for t in TIMES:
for st in CURVE_STATES:
rows.append(
f"CurveVector {pi} {hs(t)} {hs_triple(st)} "
f"{hs(forgetting_curve(w, t, st))}"
)
emit_list(add, "goldenCurveVectors", rows)
# ----------------------------------------------------------- fast recall
add("-- | The fast trace's own recall probability, which drives its own")
add("-- stability update. This is the piece that replaced the draft's")
add("-- long-\\/short-term transition function.")
add("data FastRecallVector = FastRecallVector")
add(" { frParams :: !Int")
add(" , frElapsedDays :: !Double")
add(" , frStabilityFast :: !Double")
add(" , frRecall :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for t in TIMES:
for s in FAST_STABILITIES:
rows.append(
f"FastRecallVector {pi} {hs(t)} {hs(s)} "
f"{hs(fast_component_recall(w, t, s))}"
)
emit_list(add, "goldenFastRecallVectors", rows)
# ----------------------------------------------------------- difficulty
add("-- | Initial and subsequent difficulty. The retrievability only")
add("-- matters on a lapse, where it weights the step.")
add("data DifficultyVector = DifficultyVector")
add(" { dvParams :: !Int")
add(" , dvDifficulty :: !(Maybe Double)")
add(" -- ^ 'Nothing' for the initial difficulty of a brand-new card.")
add(" , dvRetrievability :: !Double")
add(" , dvRating :: !Int")
add(" , dvNextDifficulty :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for g in RATINGS:
v = clamp(initial_difficulty(w, g), 1.0, 10.0)
rows.append(f"DifficultyVector {pi} Nothing 1.0 {g} {hs(v)}")
for d in DIFFICULTIES:
for r in RETRIEVABILITIES:
for g in RATINGS:
v = next_difficulty(w, d, g, r)
rows.append(
f"DifficultyVector {pi} (Just {hs(d)}) {hs(r)} {g} {hs(v)}"
)
emit_list(add, "goldenDifficultyVectors", rows)
# ------------------------------------------------------- half stability
add("-- | One trace's stability update, before the lapse cap is applied.")
add("data HalfStabilityVector = HalfStabilityVector")
add(" { hsParams :: !Int")
add(" , hsSlowTrace :: !Bool")
add(" -- ^ 'True' for the slow trace's block, 'False' for the fast one's.")
add(" , hsStability :: !Double")
add(" , hsDifficulty :: !Double")
add(" , hsRetrievability :: !Double")
add(" , hsRating :: !Int")
add(" , hsNextStability :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for s in HALF_STABILITIES:
for d in HALF_DIFFICULTIES:
for r in RETRIEVABILITIES:
for g in RATINGS:
for slow, base in ((True, SLOW_TRACE_BASE), (False, FAST_TRACE_BASE)):
v = stability_after_review(w, s, d, r, g, base)
rows.append(
f"HalfStabilityVector {pi} {slow} {hs(s)} {hs(d)} "
f"{hs(r)} {g} {hs(v)}"
)
emit_list(add, "goldenHalfStabilityVectors", rows)
# --------------------------------------------------------------- step
add("-- | A full memory-state transition.")
add("data StepVector = StepVector")
add(" { svParams :: !Int")
add(" , svState :: !(Maybe (Double, Double, Double))")
add(" -- ^ @(stability, difficulty, fast stability)@; 'Nothing' for the")
add(" -- first review.")
add(" , svElapsedDays :: !Double")
add(" , svRating :: !Int")
add(" , svNextState :: !(Double, Double, Double)")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for g in RATINGS:
rows.append(
f"StepVector {pi} Nothing 0.0 {g} {hs_triple(step(w, None, 0.0, g))}"
)
for st in STEP_STATES:
for dt in STEP_DELTAS:
for g in RATINGS:
rows.append(
f"StepVector {pi} (Just {hs_triple(st)}) {hs(dt)} {g} "
f"{hs_triple(step(w, st, dt, g))}"
)
emit_list(add, "goldenStepVectors", rows)
# ----------------------------------------------------------- intervals
add("-- | The interval that lands exactly on the desired retention.")
add("data IntervalVector = IntervalVector")
add(" { ivParams :: !Int")
add(" , ivDesiredRetention :: !Double")
add(" , ivState :: !(Double, Double, Double)")
add(" , ivInterval :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for dr in RETENTIONS:
for st in INTERVAL_STATES:
rows.append(
f"IntervalVector {pi} {hs(dr)} {hs_triple(st)} "
f"{hs(next_interval(w, dr, st))}"
)
emit_list(add, "goldenIntervalVectors", rows)
# -------------------------------------------------------------- replay
add("-- | A whole review history folded into a final memory state.")
add("data ReplayVector = ReplayVector")
add(" { rvParams :: !Int")
add(" , rvReviews :: ![(Double, Int)]")
add(" -- ^ @(days since the previous review, rating)@.")
add(" , rvFinalState :: !(Double, Double, Double)")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for seq_ in REVIEW_SEQUENCES:
state = None
for dt, g in seq_:
state = step(w, state, dt, g)
reviews = hs_list([f"({hs(dt)}, {g})" for dt, g in seq_])
assert state is not None
rows.append(f"ReplayVector {pi} {reviews} {hs_triple(state)}")
emit_list(add, "goldenReplayVectors", rows)
return "\n".join(lines).rstrip() + "\n"
def main(out=OUT):
text = render()
with open(out, "w") as fh:
fh.write(text)
print(f"wrote {out} ({len(text.splitlines())} lines)")
def emit_list(add, name, rows):
ty = name[len("golden"):]
ty = ty[0].upper() + ty[1:]
ty = ty[: -len("Vectors")] + "Vector"
add(f"{name} :: [{ty}]")
add(f"{name} =")
for i, row in enumerate(rows):
add((" [ " if i == 0 else " , ") + row)
add(" ]")
add("")
if __name__ == "__main__":
main()