Q3D-Calibration/qdx/calibration.py

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
                    )
                )