This page was generated from the Jupyter notebook benchmark_njit.ipynb.

Benchmark of Numerov integration with and without using numba.njit

import timeit

import numpy as np

from ryd_numerov.rydberg import RydbergState

test_cases: list[tuple[str, int, int, bool]] = [
    # species, n, l, use_njit
    ("H", 100, 80, True),
    ("H", 100, 80, False),
]

def run_benchmark(number: int = 10) -> list[dict]:
    """Run benchmark for different quantum states.

    Args:
        number: Number of times to run each test for averaging

    """
    # run the integration once to compile the numba function
    species, n, l, use_njit = test_cases[0]
    atom = RydbergState(species, n, l, j=l + 0.5)
    atom.create_wavefunction(_use_njit=True)

    results = []
    for species, n, l, use_njit in test_cases:
        # Setup the test function
        stmt = (
            "atom = RydbergState(species, n, l, j=l+0.5)\n"
            "atom.create_grid(dz=1e-3)\n"
            "atom.create_wavefunction(_use_njit=use_njit)"
        )

        # Time the integration multiple times and take average/std
        globals_dict = {"RydbergState": RydbergState, "species": species, "n": n, "l": l, "use_njit": use_njit}
        times = timeit.repeat(stmt=stmt, number=1, repeat=number, globals=globals_dict)
        avg_time = np.mean(times)
        std_time = np.std(times)

        results.append({"species": species, "n": n, "l": l, "use_njit": use_njit, "time": avg_time, "std": std_time})

    return results

results = run_benchmark(number=3)

print("\nBenchmark Results:")
print("-" * 70)
print(f"{'species':>8} {'n':>3} {'l':>3} {'use_njit':>10} {'time (ms)':>10} {'std (ms)':>10}")
print("-" * 70)
for r in results:
    print(
        f"{r['species']:>8} {r['n']:>3} {r['l']:>3} {r['use_njit']!s:>10} "
        f"{r['time'] * 1000:10.2f} {r['std'] * 1000:10.2f}"
    )

Using python implementation of Numerov integration, this is much slower!
Using python implementation of Numerov integration, this is much slower!
Using python implementation of Numerov integration, this is much slower!

Benchmark Results:
----------------------------------------------------------------------
 species   n   l   use_njit  time (ms)   std (ms)
----------------------------------------------------------------------
       H 100  80       True      11.88       0.73
       H 100  80      False     117.49       0.15