Python--iminuit相关学习笔记（七）

一、Poisson distributed data

import numpy as np
from matplotlib import pyplot as plt
from scipy.stats import multivariate_normal, poisson
from argparse import Namespace
from iminuit import Minuit
from functools import lru_cache
import joblib

@lru_cache(10000)
def run(k):
    m = Minuit(lambda lambd: -poisson.logpmf(k, lambd), lambd=k + 1)
    m.limits["lambd"] = (0, None)
    m.errordef = Minuit.LIKELIHOOD
    m.migrad()
    m.hesse()
    m.minos()
    assert m.valid

    p = m.values["lambd"]
    dp = m.errors["lambd"]
    pm = max(p + m.merrors["lambd"].lower, 0.0), p + m.merrors["lambd"].upper

    r = p, dp, *pm
    return r


rng = np.random.default_rng(seed=1)
nmc = 5000
mu = 10 ** np.linspace(-1, 2, 100)

pcov = {
    "HESSE": np.empty_like(mu),
    "MINOS": np.empty_like(mu),
}

for i, mui in enumerate(mu):

    nh = 0
    nm = 0
    for imc in range(nmc):
        k = rng.poisson(mui)

        p, dp, pd, pu = run(k)

        if p - dp < mui < p + dp:
            nh += 1
        if pd < mui < pu:
            nm += 1

    pcov["HESSE"][i] = nh / nmc
    pcov["MINOS"][i] = nm / nmc

fig, ax = plt.subplots(1, 2, figsize=(10, 4))

plt.sca(ax[0])
n = np.arange(101)
interval = {
    "HESSE": np.empty((len(n), 2)),
    "MINOS": np.empty((len(n), 2)),
}
for i, k in enumerate(n):
    p, dp, pd, pu = run(k)
    interval["HESSE"][i] = (p - dp, p + dp)
    interval["MINOS"][i] = (pd, pu)

for algo, vals in interval.items():
    plt.plot(n, vals[:, 0] - n, color="C0" if algo == "HESSE" else "C1", label=algo)
    plt.plot(n, vals[:, 1] - n, color="C0" if algo == "HESSE" else "C1", label=algo)
plt.xlabel("$k$")
plt.ylabel(r"$\lambda^d - \hat\lambda$, $\lambda^u - \hat\lambda$")

plt.sca(ax[1])
for algo, vals in pcov.items():
    plt.plot(mu, vals, label=algo)

plt.axhline(0.68, ls=":", color="0.5", zorder=0)
plt.xlabel(r"$\lambda$")
plt.ylabel("coverage probability")
plt.legend()
plt.semilogx()

二、Fit of transformed normally distributed data

rng = np.random.default_rng(1)

truth = Namespace(cr=0.1, cphi=0, sr=0.1, sphi=2)
truth.cx = truth.cr * np.cos(truth.cphi)
truth.cy = truth.cr * np.sin(truth.cphi)

d_r = rng.normal(truth.cr, truth.sr, size=5)
d_phi = rng.normal(truth.cphi, truth.sphi, size=d_r.shape)

cov = np.eye(2)
cov[0, 0] = truth.sr ** 2
cov[1, 1] = truth.sphi ** 2


def nll(cx, cy):
    cr = np.linalg.norm((cx, cy))
    cphi = np.arctan2(cy, cx)
    return -np.sum(
        multivariate_normal((cr, cphi), cov).logpdf(np.transpose((d_r, d_phi)))
    )


m = Minuit(nll, cx=0.1, cy=0)
m.errordef = Minuit.LIKELIHOOD
m.migrad()
m.hesse()
m.minos()
display(m.fmin, m.params)
display(m.merrors)

plt.figure(figsize=(14, 4))

plt.subplot(121, polar=True)
plt.plot(d_phi, d_r, ".")

plt.subplot(122, aspect="equal")
m.draw_mncontour("cx", "cy", size=100)
plt.plot(m.values["cx"], m.values["cy"], "+k")
plt.xlim(-0.1, 0.2)
plt.ylim(-0.1, 0.2)

三、Acceleration with Numba

from iminuit import Minuit
import numpy as np
import numba as nb
import math
from scipy.stats import expon, norm
from matplotlib import pyplot as plt
from argparse import Namespace

np.random.seed(1)  # fix seed

# true parameters for signal and background
truth = Namespace(n_sig=2000, f_bkg=10, sig=(5.0, 0.5), bkg=(0.0, 4.0))
n_bkg = truth.n_sig * truth.f_bkg

# make a data set
x = np.empty(truth.n_sig + n_bkg)

# fill m variables
x[: truth.n_sig] = norm(*truth.sig).rvs(truth.n_sig)
x[truth.n_sig :] = expon(*truth.bkg).rvs(n_bkg)

# cut a range in x
xrange = np.array((1.0, 9.0))
ma = (xrange[0] < x) & (x < xrange[1])
x = x[ma]

plt.hist(
    (x[truth.n_sig :], x[: truth.n_sig]),
    bins=50,
    stacked=True,
    label=("background", "signal"),
)
plt.xlabel("x")
plt.legend()