Совр дет в физике частиц ФИАН / Калориметр / calorimeter

#!/usr/bin/env python3
import ROOT
from array import array
import math

ROOT.gROOT.SetBatch(True)
ROOT.gStyle.SetOptFit(1111)

# ---------------- параметры ----------------
energies = list(range(800, 4001, 100))
particles = [
    "kaonm", "kaonp", "pim", "pip",
    "gamma", "ep", "proton", "anti_proton"
]

TRIM_RIGHT = 0.15       # удалять 15% самых больших значений; левый хвост не трогаем
NORM = 10000.0          # используется только если это действительно нормировка вашей гистограммы
RES_REL_ERR = 0.00    # заданная относительная ошибка для графика sigma/E

# ---------------- статистические функции ----------------
def histogram_counts_and_values(h):
    """
    Получить центры бинов и целочисленные частоты.
    Предполагаются невзвешенные события и целочисленные bin contents.
    Underflow/overflow не включаются.
    """
    values, counts = [], []
    for i in range(1, h.GetNbinsX() + 1):
        c = float(h.GetBinContent(i))
        if c > 0:
            values.append(float(h.GetBinCenter(i)))
            counts.append(int(round(c)))
    return values, counts


def median_and_err(h):
    """
    Медиана ROOT-гистограммы и асимптотическая стандартная ошибка:
        SE(median) ~= 1 / (2 f(m) sqrt(N)).
    Плотность оценивается частотой бина медианы / (N * bin_width).
    Учитываются только обычные бины, как и при оценке плотности.
    """
    prob = array('d', [0.5])
    xq = array('d', [0.0])
    h.GetQuantiles(1, xq, prob)
    med = float(xq[0])

    values, counts = histogram_counts_and_values(h)
    N = sum(counts)
    if N <= 0:
        return med, 0.0

    axis = h.GetXaxis()
    ibin = int(h.FindBin(med))
    if ibin < 1 or ibin > h.GetNbinsX():
        return med, 0.0

    bw = float(axis.GetBinWidth(ibin))
    if bw <= 0:
        return med, 0.0

    count_med = float(h.GetBinContent(ibin))
    # Нормированная эмпирическая плотность в бине медианы.
    f_med = count_med / (N * bw)

    # Если бин медианы пуст (например, особенности интерполяции GetQuantiles),
    # оцениваем плотность по ближайшим непустым бинам.
    if f_med <= 0:
        neighbors = []
        for j in (ibin - 1, ibin + 1):
            if 1 <= j <= h.GetNbinsX():
                cj = float(h.GetBinContent(j))
                wj = float(axis.GetBinWidth(j))
                if cj > 0 and wj > 0:
                    neighbors.append(cj / (N * wj))
        if not neighbors:
            return med, 0.0
        f_med = sum(neighbors) / len(neighbors)

    return med, 1.0 / (2.0 * f_med * math.sqrt(N))


def usech_sprava_aver_and_err(h, trim_right=0.15):
    """
    Среднее после удаления только верхнего хвоста (усечение справа).
    Оценка ошибки учитывает оценивание правой квантильной границы.

    Для фиксированного alpha и верхнего квантиля q = F^{-1}(1-alpha),
    influence function для условного среднего E[X | X <= q]:
      IF(x) = (min(x,q) - mu_trim + alpha*(q - mu_trim)) / (1-alpha)
    Эквивалентная асимптотическая дисперсия:
      Var(aver_trim) ~= Var(IF(X))/N.

    На гистограмме значения внутри бинов аппроксимируются центрами бинов.
    При дискретных/грубых данных квантиль может совпадать с большим числом
    наблюдений; формула тогда является приближением.
    """
    if not 0.0 <= trim_right < 1.0:
        raise ValueError("trim_right must be in [0, 1)")

    values, counts = histogram_counts_and_values(h)
    N = sum(counts)
    if N <= 0:
        return 0.0, 0.0

    # Сортированный компактный массив (значение, частота), без разворачивания
    # каждого события в отдельный элемент.
    pairs = sorted(zip(values, counts), key=lambda p: p[0])
    n_keep_target = (1.0 - trim_right) * N

    if trim_right == 0:
        total = sum(v * c for v, c in pairs)
        aver = total / N
        # Обычная выборочная дисперсия с частотами.
        if N < 2:
            return aver, 0.0
        ss = sum(c * (v - aver) ** 2 for v, c in pairs)
        sample_var = ss / (N - 1)
        return aver, math.sqrt(sample_var / N)

    # Оставляем ровно N*(1-alpha) наблюдений в смысле частот;
    # если граница попадает внутрь бина, используем дробную частоту
    # на граничном бине для стабильной оценки по гистограмме.
    kept = []
    remaining = n_keep_target
    for v, c in pairs:
        if remaining <= 0:
            break
        take = min(float(c), remaining)
        if take > 0:
            kept.append((v, take))
            remaining -= take

    n_kept = sum(c for _, c in kept)
    if n_kept <= 0:
        return 0.0, 0.0

    aver_trim = sum(v * c for v, c in kept) / n_kept

    # Верхняя граница усечения q: значение, на котором проходит cut.
    q = kept[-1][0]
    alpha_eff = 1.0 - n_kept / N
    p = 1.0 - alpha_eff
    if p <= 0:
        return aver_trim, 0.0

    # Винзоризируем только правый хвост: все x > q заменяем на q.
    # Для усечённого справа среднего асимптотически:
    # Var(T) ~= Var(min(X, q)) / (N * p^2), p = 1 - alpha.
    # Граничное значение q оценивается по той же выборке.
    winsor_sum = 0.0
    for v, c in pairs:
        winsor_sum += min(v, q) * c
    winsor_aver = winsor_sum / N

    winsor_ss = sum(
        c * (min(v, q) - winsor_aver) ** 2
        for v, c in pairs
    )
    winsor_var = winsor_ss / (N - 1) if N > 1 else 0.0
    se = math.sqrt(max(winsor_var, 0.0) / (N * p * p))

    return aver_trim, se


# ---------------- вспомогательные функции фитов ----------------
def fit_line(gr, name, color):
    xmin = min(gr.GetX()[i] for i in range(gr.GetN()))
    xmax = max(gr.GetX()[i] for i in range(gr.GetN()))
    f = ROOT.TF1(name, "pol1", xmin, xmax)
    f.SetLineColor(color)
    f.SetLineWidth(2)
    gr.Fit(f, "R EX0")
    return f


def fit_resolution(gr, name, color):
    xmin = min(gr.GetX()[i] for i in range(gr.GetN()))
    xmax = max(gr.GetX()[i] for i in range(gr.GetN()))
    f = ROOT.TF1(name, "[0] + [1]/sqrt(x)", xmin, xmax)
    f.SetLineColor(color)
    f.SetLineWidth(2)
    f.SetParName(0, "a")
    f.SetParName(1, "b")
    f.SetParameter(0, 0.0)
    f.SetParameter(1, 0.1)
    gr.Fit(f, "R")
    return f


def report(fit, tag):
    a1 = fit.GetParameter(1)
    a0 = fit.GetParameter(0)
    a1_err = fit.GetParError(1)
    a0_err = fit.GetParError(0)
    chi2 = fit.GetChisquare()
    ndf = fit.GetNDF()
    print(f"\n--- Фит E(Q) [{tag}] ---")
    print(f"E(Q) = ({a1:.5f} ± {a1_err:.5f})·Q + ({a0:.5f} ± {a0_err:.5f})")
    print(f"chi2/ndf = {chi2:.2f}/{ndf} = {chi2/max(ndf, 1):.3f}")
    return a1, a0, a1_err, a0_err


def report_res(fit, tag):
    a = fit.GetParameter(0)
    aerr = fit.GetParError(0)
    b = fit.GetParameter(1)
    berr = fit.GetParError(1)
    chi2 = fit.GetChisquare()
    ndf = fit.GetNDF()
    print(f"\n--- Фит σ/E = a + b/sqrt(E) [{tag}] ---")
    print(f"a = {a:.6g} ± {aerr:.3g}")
    print(f"b = {b:.6g} ± {berr:.3g}")
    print(f"chi2/ndf = {chi2:.2f}/{ndf} = {chi2/max(ndf, 1):.3f}")


# ---------------- основной цикл по частицам ----------------
for particle in particles:
    print(f"\n==================== {particle} ====================")
    output_file = f"1{particle}.root"

    E_list = []
    med_list, med_err_list = [], []
    taver_list, taver_err_list = [], []

    for E in energies:
        fname = f"{particle}_{E}.root"
        f = ROOT.TFile.Open(fname, "READ")
        if not f or f.IsZombie():
            print(f"[WARN] не открыть {fname}")
            continue

        h0 = f.Get("Egap")
        if not h0:
            print(f"[WARN] нет Egap в {fname}")
            f.Close()
            continue

        h = h0.Clone(f"Egap_{particle}_{E}")
        h.SetDirectory(0)
        f.Close()

        med, med_err = median_and_err(h)
        taver, taver_err = usech_sprava_aver_and_err(h, TRIM_RIGHT)

        E_list.append(float(E))
        med_list.append(med)
        med_err_list.append(med_err)
        taver_list.append(taver)
        taver_err_list.append(taver_err)

        print(
            f"E = {E:4d} МэВ | median = {med:8.3f} ± {med_err:6.3f} | "
            f"usech_sprava (15%) = {taver:8.3f} ± {taver_err:6.3f}"
        )

    if len(E_list) < 2:
        print(f"[ERROR] мало точек для {particle}, пропускаем")
        continue

    n = len(E_list)

    # ---------------- графики E vs Q ----------------
    gr_med = ROOT.TGraphErrors(
        n,
        array('d', med_list), array('d', E_list),
        array('d', med_err_list), array('d', [0.0] * n)
    )
    gr_med.SetTitle(f"E = f(Median Q) [{particle}];Median Q (MeV);E (MeV)")
    gr_med.SetMarkerStyle(21)
    gr_med.SetMarkerColor(ROOT.kBlue)

    gr_tm = ROOT.TGraphErrors(
        n,
        array('d', taver_list), array('d', E_list),
        array('d', taver_err_list), array('d', [0.0] * n)
    )
    gr_tm.SetTitle(
        f"E = g(Usech_sprava aver Q) [{particle}];"
        "Usech_sprava aver Q (MeV);E (MeV)"
    )
    gr_tm.SetMarkerStyle(22)
    gr_tm.SetMarkerColor(ROOT.kGreen + 2)

    # ---------------- фиты E(Q) = a1*Q + a0 ----------------
    fit_med = fit_line(gr_med, f"fit_median_{particle}", ROOT.kRed)
    fit_tm = fit_line(gr_tm, f"fit_usech_sprava_{particle}", ROOT.kMagenta + 1)

    a1_m, a0_m, a1e_m, a0e_m = report(fit_med, f"median, {particle}")
    a1_t, a0_t, a1e_t, a0e_t = report(fit_tm, f"usech_sprava, {particle}")

    # ---------------- распространение погрешностей ----------------
    # Здесь сохраняется прежняя схема пользователя, но отмечаем её ограничение:
    # для полной ошибки предсказания нужно учитывать covariance(a0,a1).
    sigma_tot_med, sigma_tot_tm = [], []
    for i in range(n):
        s_med_stat = abs(a1_m) * med_err_list[i]
        s_med_fit = abs(a1e_m) * med_list[i]
        sigma_tot_med.append(math.sqrt(s_med_stat**2 + s_med_fit**2))

        s_tm_stat = abs(a1_t) * taver_err_list[i]
        s_tm_fit = abs(a1e_t) * taver_list[i]
        sigma_tot_tm.append(math.sqrt(s_tm_stat**2 + s_tm_fit**2))

    # ---------------- графики sigma_tot / E vs E ----------------
    ey_res_med = [
        max(RES_REL_ERR * sigma_tot_med[i] / E_list[i], 1e-6)
        for i in range(n)
    ]
    ey_res_tm = [
        max(RES_REL_ERR * sigma_tot_tm[i] / E_list[i], 1e-6)
        for i in range(n)
    ]

    gr_res_med = ROOT.TGraphErrors(
        n,
        array('d', E_list),
        array('d', [sigma_tot_med[i] / E_list[i] for i in range(n)]),
        array('d', [0.0] * n),
        array('d', ey_res_med)
    )
    gr_res_med.SetTitle(f"Rel. error (Median) [{particle}];E (MeV);#sigma/E")
    gr_res_med.SetMarkerStyle(21)
    gr_res_med.SetMarkerColor(ROOT.kBlue)

    gr_res_tm = ROOT.TGraphErrors(
        n,
        array('d', E_list),
        array('d', [sigma_tot_tm[i] / E_list[i] for i in range(n)]),
        array('d', [0.0] * n),
        array('d', ey_res_tm)
    )
    gr_res_tm.SetTitle(
        f"Rel. error (usech_sprava aver) [{particle}];E (MeV);#sigma/E"
    )
    gr_res_tm.SetMarkerStyle(22)
    gr_res_tm.SetMarkerColor(ROOT.kGreen + 2)

    fit_res_med = fit_resolution(
        gr_res_med, f"fit_res_median_{particle}", ROOT.kRed
    )
    fit_res_tm = fit_resolution(
        gr_res_tm, f"fit_res_usech_sprava_{particle}", ROOT.kMagenta + 1
    )

    report_res(fit_res_med, f"median, {particle}")
    report_res(fit_res_tm, f"usech_sprava, {particle}")

    # ---------------- сохранить результаты ----------------
    out = ROOT.TFile(output_file, "RECREATE")
    gr_med.Write("E_vs_Median")
    fit_med.Write("fit_median")
    gr_tm.Write("E_vs_Usech_sprava")
    fit_tm.Write("fit_usech_sprava")
    gr_res_med.Write("Resolution_Median")
    fit_res_med.Write("fit_res_median")
    gr_res_tm.Write("Resolution_Usech_sprava")
    fit_res_tm.Write("fit_res_usech_sprava")
    c_res_tm = ROOT.TCanvas(f"c_Resolution_Usech_sprava_{particle}", "", 800, 600)
    gr_res_tm.Draw("AP")
    fit_res_tm.Draw("SAME")
    c_res_tm.Write("canvas_Resolution_Usech_sprava")
    out.Close()
    

    print(f"\nСохранено в {output_file}")