Q3D-Calibration/qdx/Bind.py

import numpy as np
import matplotlib.pyplot as plt

from .utils import get_hist, GMM_slash
from .model import Linear1D, FixedSlopeLine, pXLine
from .fit import fit_line, fit_hist_gaussian


class Bind(object):
    """Bind of a Block, using data under two energy conditions for fitting.
    The energy of an ADC device can be expressed as: $E = kx + b$.
    So the complete energy is represented by $E(x,y) = k_1x + b_1 + k_2y + b_2$.
    This formula can be converted to: $x = -\\frac{k_2}{k_1}y + \\frac{E-b_1-b_2}{k_1}$.
    Plotting the data in a 2D plane for linear regression, we can get $\\frac{k_2}{k_1}$ and $C_i$

    Using data of $E_1$ and $E_2$, now $k_1 = \\frac{E_1 - E_2}{C_1 - C_2}$.

    The incident position produces different responses in the left and right detectors, which can be expressed as:
    $Px = \\frac{k_2y - k_1x}{E}L+CL$.

    Attributes
    ----------
    L/C : float
        parameters in $Px = \\frac{k_2y - k_1x}{E}L+CL$.

    k1, k2: float
        parameters in $Px = \\frac{k_2y - k_1x}{E}L+CL$.

    b : float
        $b = b_1 + b_2$

    K1/K2/C1/C2: float
        $K_i = -\\frac{k_2}{k_1}$
        $C_i = \\frac{E_i-b}{k_1}$
    """

    def __init__(self, n, m):
        self.n, self.m = n, m
        self.px = [None, None]
        self.x, self.y = [None, None], [None, None]

        self.L, self.C = 40, 0
        self.k1, self.k2, self.b = 0, 0, 0
        self.K1, self.C1, self.K2, self.C2 = 0, 0, 0, 0

    def add_data(self, k, x, y, px):
        """k-th energy data

        Parameters
        ----------
        k : int
            k-th energy
        x : array
            left data
        y : array
            right data
        px : int
            px (set value)
        """
        self.x[k] = np.concatenate((self.x[k], x)) if self.x[k] is not None else x
        self.y[k] = np.concatenate((self.y[k], y)) if self.y[k] is not None else y
        self.px[k] = (
            np.concatenate((self.px[k], np.full(len(x), px)))
            if self.px[k] is not None
            else np.full(len(x), px)
        )

    def slash(self):
        """Using Gaussian Mixture Method (GMM) to decompose the data into noise and slashes"""
        data = GMM_slash(
            np.array(list(zip(self.x[0], self.y[0], self.px[0])), dtype=object)
        )
        self.x[0] = data[:, 0]
        self.y[0] = data[:, 1]
        self.px[0] = data[:, 2]

        data = GMM_slash(
            np.array(list(zip(self.x[1], self.y[1], self.px[1])), dtype=object)
        )
        self.x[1] = data[:, 0]
        self.y[1] = data[:, 1]
        self.px[1] = data[:, 2]

    def get_energy(self):
        """Get energy (in channel) by fit Gaussian Model with ldata + rdata"""
        eng = self.x[0] + self.y[0]
        fitted_model = fit_hist_gaussian(eng)
        self.E1 = fitted_model.mean.value

        eng = self.x[1] + self.y[1]
        fitted_model = fit_hist_gaussian(eng)
        self.E2 = fitted_model.mean.value

    def get_peak_center(self):
        """Get peak center (in channel) using Gaussian Model"""
        self.fx, self.fy, self.fz = [], [], []

        peaks = np.unique(self.px[0])
        for px in peaks:
            idx = np.where(self.px[0] == px)[0]
            x, y = self.x[0][idx], self.y[0][idx]

            if len(idx) < 400:
                continue

            fitted_model = fit_hist_gaussian(x)
            self.fx.append(fitted_model.mean.value)

            fitted_model = fit_hist_gaussian(y)
            self.fy.append(fitted_model.mean.value)

            self.fz.append(px)

        self.fx = np.array(self.fx)
        self.fy = np.array(self.fy)
        self.fz = np.array(self.fz)

    def get_line(self):
        """Fit data with $x = -\\frac{k_2}{k_1}y + \\frac{E-b_1-b_2}{k_1}$."""
        model = Linear1D()

        self.reg1 = fit_line(self.x[0], self.y[0], model)
        self.K1 = self.reg1.slope.value
        self.C1 = self.reg1.intercept.value

        self.reg2 = fit_line(self.x[1], self.y[1], model)
        self.K2 = self.reg2.slope.value
        self.C2 = self.reg2.intercept.value

    def get_kb(self):
        """Get $k_2$, $k_2$, $b$ from $K_i$, $C_i$
        Set slope equal average of $K_i$.
        """
        K = (self.K1 + self.K2) / 2
        self.K1 = self.K2 = K

        model = FixedSlopeLine(slope=K, intercept=self.C1)
        fitted_model = fit_line(self.x[0], self.y[0], model)
        self.C1 = fitted_model.intercept.value

        model = FixedSlopeLine(slope=K, intercept=self.C2)
        fitted_model = fit_line(self.x[1], self.y[1], model)
        self.C2 = fitted_model.intercept.value

        self.k1 = (self.E1 - self.E2) / (self.C1 - self.C2)
        self.k2 = -(self.K1 + self.K2) * self.k1 / 2
        self.b = (self.E1 - self.C1 * self.k1 + self.E2 - self.C2 * self.k1) / 2

    def fit_px(self):
        """
        Fit using $Px = \\frac{k_2y - k_1x}{E}L+CL$.
        $x + y = constant$, so change to $pX = kx + c$.
        """
        model = Linear1D()
        self.reg3 = fit_line(self.fx, self.fz, model)
        self.K3 = self.reg3.slope.value
        self.C3 = self.reg3.intercept.value

    def draw_fit_line(self, title):
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(1, 1, 1)
        ax.scatter(self.x[0], self.y[0], s=0.1, c="black", label=r"$E_1$")
        ax.scatter(self.x[1], self.y[1], s=0.1, c="dimgray", label=r"$E_2$")
        ax.plot(
            self.x[0],
            self.reg1(self.x[0]),
            c="red",
            label=r"$x_2={:.4f}x_1+{:.4f},\ R^2={:.5f}$".format(
                self.K1, self.C1, self.RSquare1
            ),
        )
        ax.plot(
            self.x[1],
            self.reg2(self.x[1]),
            c="orangered",
            label=r"$x_2={:.4f}x_1+{:.4f},\ R^2={:.5f}$".format(
                self.K2, self.C2, self.RSquare2
            ),
        )

        ax.legend()
        fig.savefig(title, facecolor="w", transparent=False)
        plt.close()

    def draw_peak(self, title):
        fig = plt.figure(figsize=(8, 8))

        peaks = np.unique(self.px[0])
        n, k = len(peaks), 1
        for px in peaks:
            ax = fig.add_subplot(n, 1, k)
            ax.set_title("pX = {:.2f}".format(px))

            idx = np.where(self.px[0] == px)[0]
            x, y = self.x[0][idx], self.y[0][idx]

            count1, center1 = get_hist(x)
            count2, center2 = get_hist(y)

            ax.scatter(center1, count1, s=0.5)
            ax.scatter(center2, count2, s=0.5)

            k += 1

        plt.tight_layout()
        fig.savefig(title, facecolor="w", transparent=False)
        plt.close()

    def draw_cluster(self, title):
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(1, 1, 1)

        peaks = np.unique(self.px[0])
        for px in peaks:
            idx = np.where(self.px[0] == px)[0]
            x, y = self.x[0][idx], self.y[0][idx]
            ax.scatter(x, y, s=0.1, label="pX={:.0f}".format(px))

        plt.legend()
        fig.savefig(title, facecolor="w", transparent=False)
        plt.close()

    @property
    def name(self):
        return "{:d}-{:d}-{:.1f}-{:.1f}".format(self.n, self.m, self.E1, self.E2)

    def _r_square(self, y, yp):
        mean = np.mean(y)
        SST = np.sum((y - mean) ** 2)
        SSR = np.sum((yp - mean) ** 2)
        return SSR / SST

    @property
    def RSquare1(self):
        return self._r_square(self.y[0], self.reg1(self.x[0]))

    @property
    def RSquare2(self):
        return self._r_square(self.y[1], self.reg2(self.x[1]))