Vectorized Koide's coincidence experiment

koide.py

from scipy.constants import physical_constants as pc
import numpy
import scipy.stats
import scipy.linalg


def a():
    def pnorm(v, p):
        return sum([x**p for x in v])**(1/p)

    def f(v):
        return pnorm(v, 1) / pnorm(v, 0.5)

    m_e = pc["electron mass"]
    m_m = pc["muon mass"]
    m_t = pc["tau mass"]

    v0 = [m_e[0], m_m[0], m_t[0]]

    N = 1000
    above = 0
    for _ in range(N):
        r_e = scipy.stats.norm(m_e[0], m_e[2]).rvs()
        r_m = scipy.stats.norm(m_m[0], m_m[2]).rvs()
        r_t = scipy.stats.norm(m_t[0], m_t[2]).rvs()
        t = f([r_e, r_m, r_t])
        if t > 2/3:
            above += 1


def b():
    m_e = pc["electron mass"]
    m_m = pc["muon mass"]
    m_t = pc["tau mass"]

    N = 1000
    r_e = scipy.stats.norm(m_e[0], m_e[2]).rvs(size=N)
    r_m = scipy.stats.norm(m_m[0], m_m[2]).rvs(size=N)
    r_t = scipy.stats.norm(m_t[0], m_t[2]).rvs(size=N)

    r = numpy.stack([r_e, r_m, r_t])

    above = (scipy.linalg.norm(r, 1, axis=0) / scipy.linalg.norm(r, 0.5, axis=0) > 2/3).sum()


def c():
    m_e = pc["electron mass"]
    m_m = pc["muon mass"]
    m_t = pc["tau mass"]

    N = 1000

    r = scipy.stats.norm([m_e[0], m_m[0], m_t[0]], [m_e[2], m_m[2], m_t[2]]).rvs(size=(N, 3))
    above = (scipy.linalg.norm(r, 1, axis=1) / scipy.linalg.norm(r, 0.5, axis=1) > 2/3).sum()


%timeit a()
# 1.97 s ± 19.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit b()
# 2.16 ms ± 16.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit c()
# 878 µs ± 12.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)