Q3D-Calibration/qdx/calibration.py

173 lines
5.6 KiB
Python

import os
import csv
import numpy as np
from tqdm import tqdm
from matplotlib import pyplot as plt
from .bind import Bind
from .utils import readBlockData, get_hist
class Calibration(object):
"""Calibrate the detector according to the calibration data"""
def __init__(self):
pass
def __call__(self, file1, file2, bias_b, delta_e, n=6, m=8):
"""Calibration
Parameters
----------
file1/2 : str
data file path of energy 1/2
bias_b : float
bias $b = b_1 + b_2$, in MeV
delta_e : float
delta energy between two beams, in MeV
n : int, optional
number of blocks, default 6
m : int, optional
number of binds, default 8
"""
# Initialization
self.n, self.m = n, m
self.binds = [[Bind(i, j) for j in range(m)] for i in range(n)]
# Read Data
file_list = csv.reader(open(file1, "r"))
pbar = tqdm(desc="Read Data E1", total=len(open(file1, "r").readlines()))
for row in file_list:
pn = int(row[1])
ldata, rdata = readBlockData(row[0], int(row[3]), pn, m)
for i in range(m):
bind = self.binds[pn][i]
bind.add_data(0, ldata[i], rdata[i], float(row[2]))
pbar.update(1)
pbar.close()
file_list = csv.reader(open(file2, "r"))
pbar = tqdm(desc="Read Data E2", total=len(open(file2, "r").readlines()))
for row in file_list:
pn = int(row[1])
ldata, rdata = readBlockData(row[0], int(row[3]), pn, m)
for i in range(m):
bind = self.binds[pn][i]
bind.add_data(1, ldata[i], rdata[i], float(row[2]))
pbar.update(1)
pbar.close()
# Data preprocessing
pbar = tqdm(desc="Bind Process", total=n * m)
for i in range(n):
for j in range(m):
bind: Bind = self.binds[i][j]
bind.clip()
bind.get_line()
bind.get_kb(bias_b, delta_e)
pbar.update(1)
pbar.close()
# Fit
pbar = tqdm(desc="Bind Fit", total=n * m)
for i in range(n):
for j in range(m):
bind: Bind = self.binds[i][j]
bind.get_peak_center()
bind.fit_px()
pbar.update(1)
pbar.close()
def draw_fit(self, path):
"""Draw fit result
Parameters
----------
path : str
save folder, there must be `FIT-LINE`, `GMM`, and `PEAK` 3 subfolders.
"""
pbar = tqdm(desc="Draw Fit Figure", total=self.n * self.m)
for i in range(self.n):
for j in range(self.m):
bind: Bind = self.binds[i][j]
bind.draw_fit_line(os.path.join(path, "FIT-LINE", bind.name + ".png"))
bind.draw_cluster(os.path.join(path, "GMM", bind.name + ".png"))
bind.draw_peak(os.path.join(path, "PEAK", bind.name + ".png"))
pbar.update(1)
pbar.close()
def draw_check(self, path="Check.png"):
"""Draw check figure
Parameters
----------
path : str, optional
save path
"""
pbar = tqdm(desc="Draw Check Figure", total=self.n * self.m)
fig = plt.figure(figsize=(24, 24), dpi=200)
ax1 = fig.add_subplot(3, 1, 1)
ax2 = fig.add_subplot(3, 1, 2)
ax3 = fig.add_subplot(3, 1, 3)
peaks = np.array([])
for i in range(self.n):
engs = np.array([])
for j in range(self.m):
bind = self.binds[i][j]
peaks = np.hstack((np.unique(bind.px[0]), peaks))
eng = bind.predict_energy(bind.x[0], bind.y[0])
pX = bind.predict_px(bind.x[0], bind.y[0])
count, center = get_hist(pX, step=0.5)
engs = np.hstack((engs, eng))
ax2.scatter(pX, eng, s=0.1, color="k")
ax3.scatter(center, count + 3000 * (7 - j), s=0.5, color="k")
pbar.update(1)
count, center = get_hist(engs, step=0.01)
ax1.plot(center, count)
peaks = np.unique(peaks)
for x in peaks:
ax3.vlines(x, 0, 3000 * self.m, color="gray", linestyles="dashed")
for j in range(self.m):
ax3.hlines(
3000 * j,
(np.min(peaks) // 50) * 50,
(np.max(peaks) // 50 + 1) * 50,
color="r",
linestyles="dashdot",
)
ax1.set_xlabel("Energy (MeV)")
ax1.set_ylabel("Count per bin")
ax2.set_xlabel("x (mm)")
ax2.set_ylabel("Energy (MeV)")
ax3.set_xlabel("x (mm)")
ax3.set_ylabel("Count per bin")
fig.savefig(path, facecolor="w", transparent=False)
plt.close()
pbar.close()
def save(self, path="coef.csv"):
"""Save coefficient to file
Parameters
----------
path : str, optional
save path
"""
f = open(path, "w")
for i in range(self.n):
for j in range(self.m):
bind = self.binds[i][j]
f.writelines(
"{:d},{:d},{:.9f},{:.9f},{:.9f},{:.9f},{:.9f}\n".format(
i, j, bind.k1, bind.k2, bind.b, bind.L, bind.C
)
)