packages feed

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()