Python--Normally学习笔记

Fit of transformed normally 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

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=(9, 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)

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)
truth.cov = np.eye(2)
truth.cov[0, 0] = truth.sr**2
truth.cov[1, 1] = truth.sphi**2
n_tries = 500
n_data = np.unique(np.geomspace(5, 100, 20, dtype=int))

@joblib.delayed
def run(n):
    rng = np.random.default_rng(seed=n)
    n_h = 0
    n_m = 0
    h_lus = []
    m_lus = []
    xs = []
    for i_try in range(n_tries):
        while True:
            d_r = rng.normal(truth.cr, truth.sr, size=n)
            d_phi = rng.normal(truth.cphi, truth.sphi, size=n)
            def nll(cx, cy):
                cr = np.linalg.norm((cx, cy))
                cphi = np.arctan2(cy, cx)
                return -np.sum(multivariate_normal((cr, cphi), truth.cov).logpdf(np.transpose((d_r, d_phi))))
            m = Minuit(nll, cx=0.1, cy=0)
            m.errordef = Minuit.LIKELIHOOD
            try:
                m.migrad()
                if not m.valid:
                    continue
                m.hesse()
                if not m.accurate:
                    continue
                m.minos("cx")
                if m.merrors["cx"].is_valid:
                    break
            except Exception as e:
                print(f"exception in n={n} i_try={i_try}")
                print(e)
        x = m.values["cx"]
        dx = m.errors["cx"]
        me = m.merrors["cx"]
        h_lu = x - dx, x + dx
        m_lu = x + me.lower, x + me.upper
        if h_lu[0] < truth.cx < h_lu[1]:
            n_h += 1
        if m_lu[0] < truth.cx < m_lu[1]:
            n_m += 1
        xs.append(x)
        h_lus.append(h_lu)
        m_lus.append(m_lu)
    x = np.mean(xs)
    h_l, h_u = np.mean(h_lus, axis=0)
    m_l, m_u = np.mean(m_lus, axis=0)
    return n_h, n_m, x, h_l, h_u, m_l, m_u

n_h, n_m, x, hl, hu, ml, mu = np.transpose(joblib.Parallel(-1)(run(x) for x in n_data))
h_pcov = n_h.astype(float) / n_tries
m_pcov = n_m.astype(float) / n_tries

fig, ax = plt.subplots(1, 2, figsize=(9, 4))
plt.sca(ax[0])
plt.fill_between(n_data, hl, hu, alpha=0.5, label="HESSE")
plt.fill_between(n_data, ml, mu, alpha=0.5, label="MINOS")
plt.plot(n_data, x, "-k")
plt.legend()
plt.xlabel("n")
plt.ylabel("cx")
plt.semilogx()
plt.sca(ax[1])
plt.plot(n_data, h_pcov, label="HESSE")
plt.plot(n_data, m_pcov, label="MINOS")
plt.axhline(0.68, ls=":", color="0.5", zorder=0)
plt.xlabel(r"cx")
plt.ylabel("coverage probability")
plt.legend()
plt.ylim(0, 1)
plt.semilogx()