fib_bench.py

#!/usr/bin/env python3
"""Benchmark different Fibonacci implementations and plot timings.

Methods:
 1. naive recursive
 2. recursive with memoization (top-down DP)
 3. bottom-up array (iterative DP)
4. matrix exponentiation (O(log n))

Usage: python fib_bench.py N1 N2 N3 ... [--repeat R] [--output out.png]

Notes:
 - Naive recursion grows exponentially; by default it's skipped for n > 35 unless --allow-slow is set.
 - Times are averaged over --repeat runs (default 3) for faster methods; naive uses 1 run by default.
"""

from __future__ import annotations

import argparse
import math
import time
from typing import Callable, Dict, List, Optional

import matplotlib.pyplot as plt
import numpy as np


def fib_recursive(n: int) -> int:
    if n < 2:
        return n
    return fib_recursive(n - 1) + fib_recursive(n - 2)


def fib_memo(n: int, memo: Optional[Dict[int, int]] = None) -> int:
    if memo is None:
        memo = {}
    if n < 2:
        return n
    if n in memo:
        return memo[n]
    memo[n] = fib_memo(n - 1, memo) + fib_memo(n - 2, memo)
    return memo[n]


def fib_bottom_up(n: int) -> int:
    if n < 2:
        return n
    a, b = 0, 1
    for _ in range(2, n + 1):
        a, b = b, a + b
    return b


def mat_mult(A: List[List[int]], B: List[List[int]]) -> List[List[int]]:
    return [
        [A[0][0] * B[0][0] + A[0][1] * B[1][0], A[0][0] * B[0][1] + A[0][1] * B[1][1]],
        [A[1][0] * B[0][0] + A[1][1] * B[1][0], A[1][0] * B[0][1] + A[1][1] * B[1][1]],
    ]


def mat_pow(mat: List[List[int]], n: int) -> List[List[int]]:
    # Exponentiation by squaring
    result = [[1, 0], [0, 1]]
    base = mat
    while n > 0:
        if n & 1:
            result = mat_mult(result, base)
        base = mat_mult(base, base)
        n >>= 1
    return result


def fib_matrix(n: int) -> int:
    # Using matrix [[1,1],[1,0]]^n gives F_{n+1} and F_n
    if n < 2:
        return n
    M = [[1, 1], [1, 0]]
    Mn = mat_pow(M, n - 1)
    # Mn * [F1, F0] = [F_n, F_{n-1}] where [F1, F0] = [1, 0]
    return Mn[0][0]


def time_function(fn: Callable[[int], int], n: int, repeat: int = 1) -> float:
    # Run fn(n) repeat times and return average elapsed seconds
    # For functions that may mutate shared state (memo), caller must handle it
    times = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn(n)
        t1 = time.perf_counter()
        times.append(t1 - t0)
    return float(sum(times) / len(times))


def benchmark(ns: List[int], repeat: int = 3, max_recursion_n: int = 35, allow_slow: bool = False):
    methods = [
        ("naive_recursive", lambda n: fib_recursive(n)),
        ("memo_recursive", lambda n: fib_memo(n)),
        ("bottom_up", lambda n: fib_bottom_up(n)),
        ("matrix_pow", lambda n: fib_matrix(n)),
    ]

    results: Dict[str, List[float]] = {name: [] for name, _ in methods}

    for n in ns:
        for name, fn in methods:
            if name == "naive_recursive" and (n > max_recursion_n and not allow_slow):
                results[name].append(float('nan'))
                continue

            # For memo_recursive, ensure a fresh memo per call
            if name == "memo_recursive":
                wrapped = lambda x, f=fn: f(x)
            else:
                wrapped = fn

            this_repeat = 1 if name == "naive_recursive" else repeat
            try:
                t = time_function(wrapped, n, repeat=this_repeat)
            except RecursionError:
                t = float('nan')
            results[name].append(t)

    return results


def plot_results(ns: List[int], results: Dict[str, List[float]], out_path: str):
    # prefer a clean style if available; fall back to default
    try:
        plt.style.use('seaborn-darkgrid')
    except Exception:
        try:
            plt.style.use('ggplot')
        except Exception:
            pass
    fig, ax = plt.subplots(figsize=(10, 6))

    names_map = {
        'naive_recursive': 'Naive Rec (exp)',
        'memo_recursive': 'Memoized Rec (DP)',
        'bottom_up': 'Bottom-up (array)',
        'matrix_pow': 'Matrix pow (log n)'
    }

    for name, times in results.items():
        y = np.array(times, dtype=np.float64)
        ax.plot(ns, y, marker='o', label=names_map.get(name, name))

    ax.set_xlabel('n')
    ax.set_ylabel('Time (seconds)')
    ax.set_title('Fibonacci: time comparison of methods')
    ax.legend()
    ax.set_xticks(ns)

    # If range is wide, use log scale for y to show trends
    finite = np.isfinite(np.array([t for vals in results.values() for t in vals], dtype=np.float64))
    if finite.any():
        times_all = np.array([t for vals in results.values() for t in vals], dtype=np.float64)
        times_all = times_all[np.isfinite(times_all)]
        if times_all.max() / max(times_all.min(), 1e-12) > 1000:
            ax.set_yscale('log')

    plt.tight_layout()
    plt.savefig(out_path)
    print(f"Saved plot to: {out_path}")


def parse_args():
    p = argparse.ArgumentParser(description='Benchmark Fibonacci implementations')
    p.add_argument('ns', metavar='N', type=int, nargs='+', help='n values to benchmark')
    p.add_argument('--repeat', '-r', type=int, default=3, help='repeat runs for averaging (except naive which defaults to 1)')
    p.add_argument('--output', '-o', default='fib_bench.png', help='output plot filename')
    p.add_argument('--max-recursive', type=int, default=35, help='max n for naive recursion (skipped above this unless --allow-slow)')
    p.add_argument('--allow-slow', action='store_true', help='allow running naive recursion for large n (may take very long)')
    return p.parse_args()


def main():
    args = parse_args()
    ns = sorted(set(args.ns))
    print(f"Benchmarking n values: {ns}")
    results = benchmark(ns, repeat=args.repeat, max_recursion_n=args.max_recursive, allow_slow=args.allow_slow)

    print("Results (seconds):")
    for name, times in results.items():
        print(f"  {name}: {times}")

    plot_results(ns, results, args.output)


if __name__ == '__main__':
    main()