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_line(self):
        """Fit data with $x = -\\frac{k_2}{k_1}y + \\frac{E-b_1-b_2}{k_1}$."""
        model = Linear1D()

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

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

    def get_kb(self, bias, deltaE=4):
        """Get $k_2$, $k_2$, $b$ from $K_i$, $C_i$
        Set slope equal average of $K_i$.

        Parameters
        ----------
        bias : float
            bias $b = b_1 + b_2$
        deltaE : float, optional
            delta energy between two beams
        """
        K = (self.K1 + self.K2) / 2
        self.K1 = self.K2 = K

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

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

        self.k2 = deltaE / abs(self.C1 - self.C2)
        self.k1 = -self.k2 * K

        eng = self.k1 * self.x[0] + self.k2 * self.y[0]
        fitted_model = fit_hist_gaussian(eng, delta=0.01)
        self.E1 = fitted_model.mean.value

        eng = self.k1 * self.x[1] + self.k2 * self.y[1]
        fitted_model = fit_hist_gaussian(eng, delta=0.01)
        self.E2 = fitted_model.mean.value

        self.b = bias - self.C1 * self.k2

        self.E1 += self.b
        self.E2 += self.b

    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 fit_px(self):
        """Fit using $Px = \\frac{k_2y - k_1x}{E}L+CL$."""
        model = Linear1D()
        reg = fit_line(model, self.fx, self.fz)

        E = np.mean(self.k1 * self.fx + self.k2 * self.fy + self.b)
        L = reg.slope.value * E / (self.k2 * self.K1 - self.k1)
        C = (reg.intercept.value - self.k2 * self.C1 * L / E) / L

        model = pXLine(L=L, C=C, k1=self.k1, k2=self.k2, b=self.b)
        reg = fit_line(
            model, np.array(list(zip(self.fx, self.fy)), dtype=object), self.fz
        )
        self.L = reg.L.value
        self.C = reg.C.value

    def predict_energy(self, x, y):
        """Use $E(x,y) = k_1x + b_1 + k_2y + b_2$. to calculate energy.

        Parameters
        ----------
        x/y : array
            data
        """
        return self.k1 * x + self.k2 * y + self.b

    def predict_px(self, x, y):
        """Use $Px = \\frac{k_2y - k_1x}{E}L+CL$ to calculate pX.

        Parameters
        ----------
        x/y : array
            data
        """
        eng = self.predict_energy(x, y)
        return (self.k2 * y - self.k1 * x) / eng * self.L + self.C * self.L

    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.K1 * self.x[0] + self.C1,
            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.K2 * self.x[1] + self.C2,
            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}-{:.2f}-{:.2f}".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.K1 * self.x[0] + self.C1)

    @property
    def RSquare2(self):
        return self._r_square(self.y[1], self.K2 * self.x[1] + self.C2)

    def __call__(self, data):
        """Data is read to complete initialization"""
        self.k1 = data[0]
        self.k2 = data[1]
        self.b = data[2]
        self.L = data[3]
        self.C = data[4]