repro.py

import time

import numpy as np
from numba import njit
from scipy.optimize import minimize

SQRT_2PI = np.sqrt(2 * np.pi)


def ldnorm(x, mu, sd):
    return x.size * np.log(sd * SQRT_2PI) + sum(((x - mu) / sd) ** 2 / 2)


def mk(k):
    x = np.random.normal(5, 3, k)
    return lambda theta: ldnorm(x, *theta)


def theta():
    return np.random.gamma(3, (2, 1))


def main(fn, n, *, perf_counter=time.perf_counter):
    time_opt = 0
    time_numba = 0
    ncalls = 0

    start = perf_counter()

    def tramp(theta):
        nonlocal time_numba, ncalls
        t0 = perf_counter()
        result = fn(theta)
        time_numba += perf_counter() - t0
        ncalls += 1
        return result

    for _ in range(n):
        theta_i = theta()
        t0 = perf_counter()
        minimize(
            tramp,
            theta_i,
            bounds=(
                (-np.inf, np.inf),
                (1e-9, np.inf),
            ),
        )
        time_opt += perf_counter() - t0

    duration = perf_counter() - start
    other = duration - time_opt
    time_opt -= time_numba
    fntime = time_numba / ncalls

    print(
        f"{n=} avfn={fntime*1e3:.3f} ms/{ncalls=} numba={time_numba/duration:.1%} optimize={time_opt/duration:.1%} other={other/duration:.1%}"
    )


if __name__ == "__main__":
    fn = mk(10000)
    main(fn, 1)
    ldnorm = njit(ldnorm)
    # force JIT compile
    fn(theta())
    main(fn, 1)
    # start profiling after JIT compilation
    import scalene; scalene.scalene_profiler.start()  # noqa
    main(fn, 100)