haskell-fsrs-7.0.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,
LOWER_BOUNDS,
UPPER_BOUNDS,
LONG_TERM_BASE,
SHORT_TERM_BASE,
clamp,
forgetting_curve,
initial_difficulty,
next_difficulty,
next_interval,
stability_after_review,
step,
transition_function,
)
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) + "]"
# --------------------------------------------------------------------------
# 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))
# Restore the ordering constraints the clipper enforces.
for i in (1, 2, 3):
w[i] = max(w[i], w[i - 1])
w[28] = max(w[28], w[27])
w[30] = max(w[30], w[29])
return [round(v, 6) for v in 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, 4.1283, 15.0, 100.0, 1000.0, 36500.0]
DIFFICULTIES = [1.0, 2.5, 4.194588083372719, 7.0, 10.0]
HALF_DIFFICULTIES = [1.0, 4.194588083372719, 10.0]
HALF_STABILITIES = [0.01, 1.0, 4.1283, 1000.0]
HALF_RETRIEVABILITIES = [0.05, 0.5, 1.0]
STEP_STABILITIES = [0.0001, 1.0, 4.1283, 1000.0, 36500.0]
STEP_DIFFICULTIES = [1.0, 4.194588083372719, 10.0]
STEP_DELTAS = [0.0, 10.0 / 1440.0, 0.5, 1.0, 45.0, 400.0]
DELTAS = [0.0, 1.0 / 1440.0, 10.0 / 1440.0, 0.5, 1.0, 6.0, 45.0, 400.0]
RETENTIONS = [0.7, 0.8, 0.85, 0.9, 0.95, 0.97, 0.99]
RATINGS = [1, 2, 3, 4]
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)],
[(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 official benchmark 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(" , DifficultyVector (..)")
add(" , goldenDifficultyVectors")
add(" , HalfStabilityVector (..)")
add(" , goldenHalfStabilityVectors")
add(" , TransitionVector (..)")
add(" , goldenTransitionVectors")
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, s)@, the probability of recall.")
add("data CurveVector = CurveVector")
add(" { cvParams :: !Int")
add(" , cvElapsedDays :: !Double")
add(" , cvStability :: !Double")
add(" , cvRetrievability :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for t in TIMES:
for s in STABILITIES:
rows.append(
f"CurveVector {pi} {hs(t)} {hs(s)} {hs(forgetting_curve(w, t, s))}"
)
emit_list(add, "goldenCurveVectors", rows)
# ----------------------------------------------------------- difficulty
add("-- | Initial and subsequent difficulty.")
add("data DifficultyVector = DifficultyVector")
add(" { dvParams :: !Int")
add(" , dvDifficulty :: !(Maybe Double)")
add(" -- ^ 'Nothing' for the initial difficulty of a brand-new card.")
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 {g} {hs(v)}")
for d in DIFFICULTIES:
for g in RATINGS:
v = clamp(next_difficulty(w, d, g), 1.0, 10.0)
rows.append(
f"DifficultyVector {pi} (Just {hs(d)}) {g} {hs(v)}"
)
emit_list(add, "goldenDifficultyVectors", rows)
# ------------------------------------------------------- half stability
add("-- | One half of the stability update, before the two are blended.")
add("data HalfStabilityVector = HalfStabilityVector")
add(" { hsParams :: !Int")
add(" , hsLongTerm :: !Bool")
add(" -- ^ 'True' for the long-term block, 'False' for the short-term one.")
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 HALF_RETRIEVABILITIES:
for g in RATINGS:
for long, base in ((True, LONG_TERM_BASE), (False, SHORT_TERM_BASE)):
v = stability_after_review(w, s, d, r, g, base)
rows.append(
f"HalfStabilityVector {pi} {long} {hs(s)} {hs(d)} "
f"{hs(r)} {g} {hs(v)}"
)
emit_list(add, "goldenHalfStabilityVectors", rows)
# --------------------------------------------------------- transition
add("-- | The long-\\/short-term blending coefficient.")
add("data TransitionVector = TransitionVector")
add(" { tvParams :: !Int")
add(" , tvElapsedDays :: !Double")
add(" , tvCoefficient :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for t in TIMES:
rows.append(
f"TransitionVector {pi} {hs(t)} {hs(transition_function(w, t))}"
)
emit_list(add, "goldenTransitionVectors", rows)
# --------------------------------------------------------------- step
add("-- | A full memory-state transition.")
add("data StepVector = StepVector")
add(" { svParams :: !Int")
add(" , svState :: !(Maybe (Double, Double))")
add(" -- ^ @(stability, difficulty)@; 'Nothing' for the first review.")
add(" , svElapsedDays :: !Double")
add(" , svRating :: !Int")
add(" , svNextState :: !(Double, Double)")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for g in RATINGS:
s1, d1 = step(w, None, 0.0, g)
rows.append(f"StepVector {pi} Nothing 0.0 {g} ({hs(s1)}, {hs(d1)})")
for s in STEP_STABILITIES:
for d in STEP_DIFFICULTIES:
for dt in STEP_DELTAS:
for g in RATINGS:
s1, d1 = step(w, (s, d), dt, g)
rows.append(
f"StepVector {pi} (Just ({hs(s)}, {hs(d)})) {hs(dt)} {g} "
f"({hs(s1)}, {hs(d1)})"
)
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(" , ivStability :: !Double")
add(" , ivInterval :: !Double")
add(" }")
add(" deriving stock (Eq, Show)")
add("")
rows = []
for pi, w in enumerate(PARAMETER_SETS):
for dr in RETENTIONS:
for s in STABILITIES:
rows.append(
f"IntervalVector {pi} {hs(dr)} {hs(s)} {hs(next_interval(w, dr, s))}"
)
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)")
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(state[0])}, {hs(state[1])})"
)
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()