"""Does numpy compute the same np.log and np.exp on every x86-64 CPU?

Prints the SIMD features numpy dispatches to, hashes of np.log and np.exp over fixed arrays of
10 million values, and an annualised Sharpe ratio computed from log returns of a fixed balance
curve. The inputs are built with + and * only, so they are the same bits on every machine. Run it
on two machines and compare the lines.

  python numpy_log_cpu.py [save.npz]
"""
import hashlib
import sys

import numpy as np

try:
    from numpy._core._multiarray_umath import __cpu_features__ as features
except ImportError:  # numpy 1.x
    from numpy.core._multiarray_umath import __cpu_features__ as features

levels = ("X86_V3", "X86_V4", "AVX512_ICL", "AVX512_SPR")              # numpy 2.4 and newer
if "X86_V4" not in features:
    levels = ("AVX2", "FMA3", "AVX512F", "AVX512_SKX", "AVX512_ICL", "AVX512_SPR")
active = [k for k in levels if features.get(k)]
rng = np.random.default_rng(20261009)
z = rng.normal(0.0, 0.01, 10_000_000)
x = 1.0 + z                                 # ratios near 1, like balance[t] / balance[t-1]
log_x, exp_z = np.log(x), np.exp(z)
balance = 1e6 * np.cumprod(1.0 + rng.normal(0.0004, 0.01, 2520))   # ten years of daily balances
r = np.log(balance[1:] / balance[:-1])
sharpe = r.mean() / r.std() * np.sqrt(252)
if len(sys.argv) > 1:
    np.savez(sys.argv[1], x=x, z=z, log_x=log_x, exp_z=exp_z)
digest = lambda a: hashlib.sha256(a.tobytes()).hexdigest()[:16]
print("numpy", np.__version__, "| CPU features used:", ", ".join(active))
print("input hash: ", digest(x))
print("np.log hash:", digest(log_x))
print("np.exp hash:", digest(exp_z))
print("Sharpe ratio: %.17g" % sharpe)