packages feed

haskell-fsrs-7.1.0: reference/test_reference.py

"""Self-checks for the pure-Python FSRS-7 reference.

The reference is the source of truth for the golden vectors that the Haskell
test suite is measured against, so it gets a few checks of its own. These are
the model's structural invariants, not a second copy of the Haskell suite:
enough to catch a transcription typo before it is baked into GoldenData.hs.

`test_upstream_*` are the two numeric assertions the reference implementation
makes about FSRS-7 in its own test suite. They are the only fixtures here that
did not come out of this file, so they are what pins the transcription to
upstream rather than to itself.

No test framework needed. Run from the repository root:

    python3 reference/test_reference.py
"""

import math
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from fsrs7_reference import (  # noqa: E402
    DEFAULT_PARAMETERS,
    FAST_TRACE_RATIO,
    LOWER_BOUNDS,
    MAX_DIFFICULTY,
    MIN_DIFFICULTY,
    ORDERING_CONSTRAINTS,
    RETENTION_FLOOR,
    STABILITY_MAX,
    STABILITY_MIN,
    UPPER_BOUNDS,
    clamp,
    clip_parameters,
    fast_component_recall,
    forgetting_curve,
    forgetting_curve_derivative,
    initial_difficulty,
    initial_state,
    next_difficulty,
    next_interval,
    replay,
    stability_after_review,
    step,
)

PARAM_LEN = 34
RATINGS = [1, 2, 3, 4]
RETENTION_CEILING = 1.0 - RETENTION_FLOOR


def parameter_sets(count, seed):
    """The defaults plus `count` random vectors from inside the valid box."""
    rng = random.Random(seed)
    yield DEFAULT_PARAMETERS
    for _ in range(count):
        yield clip_parameters(
            [rng.uniform(lo, hi) for lo, hi in zip(LOWER_BOUNDS, UPPER_BOUNDS)]
        )


def sample_states(seed):
    """Random `(state, delta_t)` pairs, mostly at the model's own trace ratio."""
    rng = random.Random(seed)
    for _ in range(40):
        s = math.exp(rng.uniform(math.log(STABILITY_MIN), math.log(STABILITY_MAX)))
        d = rng.uniform(MIN_DIFFICULTY, MAX_DIFFICULTY)
        if rng.random() < 0.6:
            s_fast = clamp(s * FAST_TRACE_RATIO, STABILITY_MIN, STABILITY_MAX)
        else:
            s_fast = math.exp(rng.uniform(math.log(STABILITY_MIN), math.log(STABILITY_MAX)))
        dt = (
            0.0
            if rng.random() < 0.2
            else math.exp(rng.uniform(math.log(1 / 86400), math.log(3650)))
        )
        yield (s, d, s_fast), dt


def natural_states(seed):
    """States at exactly the fast/slow ratio the model establishes."""
    rng = random.Random(seed)
    for _ in range(40):
        s = math.exp(
            rng.uniform(math.log(STABILITY_MIN / FAST_TRACE_RATIO), math.log(STABILITY_MAX))
        )
        d = rng.uniform(MIN_DIFFICULTY, MAX_DIFFICULTY)
        yield (s, d, s * FAST_TRACE_RATIO)


# ---------------------------------------------------------------------------


def test_parameter_vector_is_well_formed():
    assert len(DEFAULT_PARAMETERS) == PARAM_LEN, len(DEFAULT_PARAMETERS)
    assert len(LOWER_BOUNDS) == PARAM_LEN and len(UPPER_BOUNDS) == PARAM_LEN
    for i, (lo, hi) in enumerate(zip(LOWER_BOUNDS, UPPER_BOUNDS)):
        assert lo <= hi, f"bound {i} is inverted: {lo} > {hi}"
    for i, w in enumerate(DEFAULT_PARAMETERS):
        assert LOWER_BOUNDS[i] <= w <= UPPER_BOUNDS[i], f"default w{i} = {w} is out of bounds"
    for i, j in ORDERING_CONSTRAINTS:
        assert DEFAULT_PARAMETERS[i] <= DEFAULT_PARAMETERS[j], f"w{i} > w{j}"


def test_clipping_is_idempotent_and_lands_in_the_box():
    rng = random.Random(13)
    for _ in range(200):
        w = [rng.uniform(-1000, 1000) for _ in range(PARAM_LEN)]
        once = clip_parameters(w)
        assert clip_parameters(once) == once
        for i, (lo, hi) in enumerate(zip(LOWER_BOUNDS, UPPER_BOUNDS)):
            assert lo <= once[i] <= hi, (i, once[i])
        for i, j in ORDERING_CONSTRAINTS:
            assert once[i] <= once[j], (i, j)


def test_retrievability_is_rescaled_away_from_one():
    # FSRS-7 squeezes R into [1e-5, 1 - 1e-5], so a card reviewed a moment ago
    # is *almost* certain rather than certain.
    for w in parameter_sets(20, seed=1):
        for state, _ in sample_states(seed=2):
            assert abs(forgetting_curve(w, 0.0, state) - RETENTION_CEILING) < 1e-12


def test_retrievability_is_a_decreasing_probability():
    for w in parameter_sets(40, seed=3):
        for state, dt in sample_states(seed=4):
            r = forgetting_curve(w, dt, state)
            assert RETENTION_FLOOR <= r <= RETENTION_CEILING, r
            assert forgetting_curve(w, dt + 1e-3, state) <= r


def test_the_fast_component_is_a_bare_power_law():
    # Unlike the mixture, this one is not rescaled: it hits 1 exactly.
    for w in parameter_sets(20, seed=5):
        for s in (STABILITY_MIN, 0.01, 1.0, 37.0, STABILITY_MAX):
            assert fast_component_recall(w, 0.0, s) == 1.0
            previous = 1.0
            for dt in (1 / 86400, 0.01, 0.5, 1.0, 10.0, 3650.0):
                r = fast_component_recall(w, dt, s)
                assert 0.0 < r <= 1.0, r
                assert r <= previous
                previous = r


def test_the_derivative_matches_a_finite_difference():
    for w in parameter_sets(20, seed=6):
        for state, _ in sample_states(seed=7):
            for t in (0.5, 5.0, 50.0):
                h = 1e-6 * t
                numeric = (
                    forgetting_curve(w, t + h, state) - forgetting_curve(w, t - h, state)
                ) / (2 * h)
                exact = forgetting_curve_derivative(w, t, state)
                assert exact <= 0.0, exact
                assert abs(numeric - exact) <= 1e-7 + 1e-5 * abs(exact), (numeric, exact)


def test_base_weights_are_the_recall_probability_at_t_equals_s():
    # Give both power laws the same base and, at an average difficulty with the
    # two traces level, the mixture collapses to that base at t == s.
    for base in (0.5, 0.7, 0.85):
        w = list(DEFAULT_PARAMETERS)
        w[25] = w[26] = base
        for s in (0.5, 1.0, 37.0, 1000.0):
            expected = base * (1.0 - 2.0 * RETENTION_FLOOR) + RETENTION_FLOOR
            assert abs(forgetting_curve(w, s, (s, 5.0, s)) - expected) < 1e-12


def test_intervals_land_on_the_desired_retention():
    for w in parameter_sets(20, seed=8):
        for dr in (0.7, 0.8, 0.9, 0.95, 0.99):
            for state, _ in sample_states(seed=9):
                t = next_interval(w, dr, state)
                if t <= 1 / 86400 or t >= 36500:
                    continue  # saturated, cannot hit the target
                assert abs(forgetting_curve(w, t, state) - dr) < 1e-9


def test_a_better_rating_is_never_worse_for_the_card():
    for w in parameter_sets(40, seed=10):
        for state, dt in sample_states(seed=11):
            outcomes = [step(w, state, dt, g) for g in RATINGS]
            stabilities = [new_s for new_s, _, _ in outcomes]
            assert stabilities == sorted(stabilities), stabilities


def test_a_lapse_is_the_harshest_answer_when_it_was_surprising():
    # FSRS-7 weights a lapse's difficulty step by `r + 0.1`, so a lapse only
    # outweighs a Hard once the model expected the card to be recalled.
    for w in parameter_sets(20, seed=12):
        for d in (1.0, 3.0, 5.0, 7.5, 10.0):
            for r in (0.4, 0.7, 0.95):
                difficulties = [next_difficulty(w, d, g, r) for g in RATINGS]
                assert difficulties == sorted(difficulties, reverse=True), difficulties
            # ... and a gentler lapse really is gentler.
            assert next_difficulty(w, d, 1, 0.05) <= next_difficulty(w, d, 1, 0.95)


def test_remembering_helps_and_forgetting_does_not():
    for w in parameter_sets(40, seed=13):
        for state, dt in sample_states(seed=14):
            s = state[0]
            for g in (2, 3, 4):
                assert step(w, state, dt, g)[0] >= s * (1 - 1e-12)
            assert step(w, state, dt, 1)[0] <= s * (1 + 1e-12)


def test_post_lapse_stability_ignores_difficulty():
    # The finished model ablated the d ** -fail_d_exp factor from both blocks.
    for w in parameter_sets(20, seed=15):
        for base in (7, 15):
            for s in (0.01, 1.0, 37.0, 1000.0):
                for r in (0.05, 0.5, 0.95):
                    values = {
                        stability_after_review(w, s, d, r, 1, base)
                        for d in (1.0, 3.0, 5.0, 7.0, 10.0)
                    }
                    assert len(values) == 1, values


def test_a_lapse_pulls_the_fast_trace_under_the_slow_one():
    for w in parameter_sets(40, seed=16):
        for state, dt in sample_states(seed=17):
            new_s, _, new_s_short = step(w, state, dt, 1)
            # The cap is applied before the final clamp, so at the very bottom
            # of the range the floor wins.
            cap = max(STABILITY_MIN, FAST_TRACE_RATIO * new_s)
            assert new_s_short <= cap + 1e-12, (new_s_short, cap)


def test_the_state_stays_in_range():
    for w in parameter_sets(40, seed=18):
        for state, dt in sample_states(seed=19):
            for g in RATINGS:
                new_s, new_d, new_s_short = step(w, state, dt, g)
                assert STABILITY_MIN <= new_s <= STABILITY_MAX, new_s
                assert STABILITY_MIN <= new_s_short <= STABILITY_MAX, new_s_short
                assert MIN_DIFFICULTY <= new_d <= MAX_DIFFICULTY, new_d
        for g in RATINGS:
            first_s, first_d, first_s_fast = initial_state(w, g)
            assert first_s == clamp(w[g - 1], STABILITY_MIN, STABILITY_MAX)
            assert first_s_fast == clamp(
                FAST_TRACE_RATIO * first_s, STABILITY_MIN, STABILITY_MAX
            )
            assert MIN_DIFFICULTY <= first_d <= MAX_DIFFICULTY


def test_difficulty_reverts_towards_an_easy_first_review():
    # A Good review leaves the rating delta at zero, so all that is left is the
    # 1% pull towards the *unclamped* init_d(4). That is what pins down which
    # rating the anchor is taken from — with the default weights the anchor sits
    # below the difficulty floor, so the iterate itself only reaches the floor.
    for w in parameter_sets(10, seed=20):
        anchor = initial_difficulty(w, 4)
        for d in (1.0, 2.5, 5.0, 7.5, 10.0):
            expected = clamp(0.01 * anchor + 0.99 * d, MIN_DIFFICULTY, MAX_DIFFICULTY)
            assert abs(next_difficulty(w, d, 3, 0.9) - expected) < 1e-12
        limit = 5.0
        for _ in range(5000):
            limit = next_difficulty(w, limit, 3, 0.9)
        assert abs(limit - clamp(anchor, MIN_DIFFICULTY, MAX_DIFFICULTY)) < 1e-9, limit


def test_difficulty_never_helps_recall_at_the_natural_trace_ratio():
    # Difficulty pulls two ways: it speeds the slow component up, but also
    # shifts mixture weight onto the fast one. At the ratio the model maintains
    # the first wins; for an arbitrarily lopsided pair of traces neither does.
    w = DEFAULT_PARAMETERS
    for s, _, s_fast in natural_states(seed=21):
        for dt in (0.001, 0.5, 7.0, 365.0):
            previous = None
            for d in range(1, 11):
                r = forgetting_curve(w, dt, (s, float(d), s_fast))
                if previous is not None:
                    assert r <= previous + 1e-15, (s, dt, d, previous, r)
                previous = r


# --------------------------------------------------------------------------
# Fixtures taken from the reference implementation's own test suite:
# src/inference.rs of https://github.com/open-spaced-repetition/fsrs-rs/pull/426
# Upstream computes in single precision, hence the loose tolerances.
# --------------------------------------------------------------------------


def test_upstream_memory_state_fsrs7():
    history = [(0.0, 1), (0.0, 3), (1.0, 3), (3.0, 3), (8.0, 3), (21.0, 3)]
    s, d, _ = replay(DEFAULT_PARAMETERS, history)
    assert abs(s - 25.985723) < 1e-4, s
    assert abs(d - 5.877549) < 1e-4, d


def test_upstream_next_interval_fsrs7():
    state = (1.0, 5.0, 1.0)
    intervals = [
        max(1, round(next_interval(DEFAULT_PARAMETERS, i / 10.0, state)))
        for i in range(1, 11)
    ]
    assert intervals == [36500, 36500, 36500, 12813, 843, 92, 13, 2, 1, 1], intervals


# ---------------------------------------------------------------------------


def main():
    checks = [(name, fn) for name, fn in sorted(globals().items()) if name.startswith("test_")]
    failed = 0
    for name, check in checks:
        try:
            check()
        except AssertionError as error:
            failed += 1
            print(f"FAIL {name}: {error}")
        else:
            print(f"ok   {name}")
    print(f"\n{len(checks) - failed} of {len(checks)} checks passed")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())