161 lines
4.9 KiB
Python
161 lines
4.9 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
|
|
|
|
Parameters
|
|
----------
|
|
bias_b : float
|
|
bias $b = b_1 + b_2$
|
|
delta_e : float
|
|
delta energy between two beams
|
|
n : int, optional
|
|
number of blocks, default 6
|
|
m : int, optional
|
|
number of binds, default 8
|
|
"""
|
|
|
|
def __init__(self, bias_b, delta_e, n=6, m=8):
|
|
self.bias = bias_b
|
|
self.delta_e = delta_e
|
|
self.n = n
|
|
self.m = m
|
|
|
|
self.binds = [[Bind(i, j) for j in range(m)] for i in range(n)]
|
|
|
|
def __call__(self, file1, file2):
|
|
"""Calibration
|
|
|
|
Parameters
|
|
----------
|
|
file1/2 : str
|
|
data file path of energy 1/2
|
|
"""
|
|
# 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], pn, self.m)
|
|
for i in range(self.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 E1", total=len(open(file2, "r").readlines()))
|
|
for row in file_list:
|
|
pn = int(row[1])
|
|
ldata, rdata = readBlockData(row[0], pn, self.m)
|
|
for i in range(self.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=self.n * self.m)
|
|
for i in range(self.n):
|
|
for j in range(self.m):
|
|
bind: Bind = self.binds[i][j]
|
|
bind.slash()
|
|
bind.get_line()
|
|
bind.get_kb(self.bias, self.delta_e)
|
|
pbar.update(1)
|
|
pbar.close()
|
|
|
|
# Fit
|
|
pbar = tqdm(desc="Bind Fit", total=self.n * self.m)
|
|
for i in range(self.n):
|
|
for j in range(self.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.
|
|
"""
|
|
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"))
|
|
|
|
def draw_check(self, path="Check.png"):
|
|
"""Draw check figure
|
|
|
|
Parameters
|
|
----------
|
|
path : str, optional
|
|
save path
|
|
"""
|
|
pbar = tqdm(desc="Draw Figure Check", total=self.n * self.m)
|
|
fig = plt.figure(figsize=(24, 15), dpi=200)
|
|
ax1 = fig.add_subplot(2, 1, 1)
|
|
ax2 = fig.add_subplot(2, 1, 2)
|
|
peaks = np.array([])
|
|
|
|
for i in range(self.n):
|
|
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, delta=0.5)
|
|
ax1.scatter(pX, eng, s=0.1, color="k")
|
|
ax2.scatter(center, count + 3000 * (7 - j), s=0.5, color="k")
|
|
|
|
pbar.update(1)
|
|
|
|
peaks = np.unique(peaks)
|
|
for x in peaks:
|
|
ax2.vlines(x, 0, 3000 * self.m, color="gray", linestyles="dashed")
|
|
for j in range(self.m):
|
|
ax2.hlines(
|
|
2500 * j,
|
|
(np.min(peaks) // 50) * 50,
|
|
(np.min(peaks) // 50 + 1) * 50,
|
|
color="r",
|
|
linestyles="dashdot",
|
|
)
|
|
|
|
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
|
|
)
|
|
)
|