scripts/test.py

"""CPU-only verification test for Zaremba Density Enumeration"""
print("Testing zaremba-density-cuda...")

def zaremba_count_cpu(digits, max_d):
    """CPU brute-force: enumerate CF denominators from digit set up to max_d."""
    representable = set()
    # BFS on (p_prev, p_curr, q_prev, q_curr) = convergent state
    stack = []
    for a0 in digits:
        # First convergent: a0/1
        stack.append((1, a0, 0, 1))  # (A_{-1}, A_0, B_{-1}, B_0) = (1, a0, 0, 1)
    while stack:
        ap, a, bp, b = stack.pop()
        if b <= max_d and b > 0:
            representable.add(b)
        for d in digits:
            a_new = d * a + ap
            b_new = d * b + bp
            if b_new > max_d:
                continue
            stack.append((a, a_new, b, b_new))
    return len(representable)

# A={1}: denominators are Fibonacci numbers: 1,1,2,3,5,8,13,21,...
# Unique Fibonacci <= 20: {1,2,3,5,8,13} = 6
count = zaremba_count_cpu([1], 20)
print(f"  A={{1}}, N=20: {count} representable (expect 6)")
assert count == 6, f"Expected 6, got {count}"

# A={1,2}: should have few exceptions
count12 = zaremba_count_cpu([1,2], 50)
print(f"  A={{1,2}}, N=50: {count12} representable out of 50")
assert count12 >= 30, "Expected most d<=50 representable for {1,2}"

print(f"\n2/2 tests passed")