analysis.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy import odr


class Analysis:
    """
    Analysis class:
        This class computes mean and error on the detector output.
        It then fits the exact equations of motions for a particle
        in a magnetic field to the input data.
        The results can then be visualized with plots.
    """

    def __init__(self):
        ...

    def fillFromFile(self):
        """
        Fill the detector data from a file.
        For exact format consider detector documentation.
        """
        bounds = np.open('detectorOUTPUT.npy')
        self.fill(bounds)


    def fill(self, bounds):
        """
        Fill the detector data to dataframes.
        The mean is just the arithmetic average of the pixel bounds.
        The standard deviation is the differece in the bounds divided
        by sqrt(12) assuming a flat distribution.
        The time error is assigned as 1ns.
        For input format consult the detector output documentation.
        """
        self.results = pd.DataFrame({'w'    : np.zeros(2),
                                    'vx'    : np.zeros(2),
                                    'vy'    : np.zeros(2),
                                    'vz'    : np.zeros(2),
									'm'		: np.zeros(2)})

        #Determine mean and std of input

        self.xPoints = pd.DataFrame({'Mean': (bounds[1:,0,0] + bounds[1:,0,1])/2.,
                                    'Error': (bounds[1:,0,0] - bounds[1:,0,1])/np.sqrt(12.)})
        self.yPoints = pd.DataFrame({'Mean': (bounds[1:,1,0] + bounds[1:,1,1])/2.,
                                    'Error': (bounds[1:,1,0] - bounds[1:,1,1])/np.sqrt(12.)})
        self.zPoints = pd.DataFrame({'Mean': (bounds[1:,2,0] + bounds[1:,2,1])/2.,
                                    'Error': (bounds[1:,2,0] - bounds[1:,2,1])/np.sqrt(12.)})
        self.tPoints = pd.DataFrame({'Mean': bounds[1:,3,0],
                                     'Error': 1e-9})
        """
        print(self.xPoints)
        print(self.yPoints)
        print(self.zPoints)
        print(self.tPoints)
        """

    """----------------------------------------------------------------------------
    FITTING

    Here we fit the data. As a fitting procedure we use ORTHOGONAL DISTANCE REGRESSION.
    We fit all three spacial dimensions with respect to time, which is the same as fitting
    path length for non-relativistic paricles.

    The fitting functions / equations of motion are xEOM, yEOM, zEOM.

    The actual fits are performed by the functions xFit, yFit, zFit.

    The fitting functions are called by the function Fit. It also collects
    all the results and fills them into the results dataframe.
    """
    
    def Fit(self):
        """
        This function performs all the fits. The most stable are xFit and
        zFit, so we use the results of those.
        As the fitting functions are non-trivial and the signs of the
        parameters can get mixed up, the correct signs are also determined
        within this function.
        """
        xResults = self.xFit()
        yResults = self.yFit()
        zResults = self.zFit()
        """
        print(xResults.beta)
        print(yResults.beta)
        print(zResults.beta)
        """
        #Determine correct signs
        wSgn  = 1
        vySgn = 1

        if yResults.beta[0] < 0 or yResults.beta[2] < 0:
            wSgn  = -1
        if yResults.beta[1] < 0:
            vySgn = -1
        
        #Extract correct results from fitting and fill to results container
        self.results['vx'] = np.array([xResults.beta[0], xResults.sd_beta[0]])

        w = wSgn*np.abs(zResults.beta[0])
        wStd = zResults.sd_beta[0]
        self.results['w'] = np.array([w , wStd])

        vy = vySgn*np.abs(zResults.beta[1])
        vyStd = zResults.sd_beta[1]
        self.results['vy'] = np.array([vy,vyStd])

        vz = np.abs(zResults.beta[2])
        vzStd = zResults.sd_beta[2]
        self.results['vz'] = np.array([vz,vzStd])

        m = 0.511/np.abs(w)
        mStd = m*wStd/np.abs(w)
        self.results['m'] = np.array([m, mStd])


    """
    Equations of motion:
        w = qB/m, where q is the charge [C], B the magnetic field [T] and 
                m the particle mass.
        Vx = initial x velocity of the particle
        Vy = initial y velocity of the particle
        Vz = initial z velocity of the particle

    """
    def xEOM(self, B, t):
        """
        B = [Vx]
        """
        x = B[0]*t
        return x

    def yEOM(self, B, t):
        """
        B = [w, Vy, Vz]
        """
        #y = B[1] + np.sqrt(B[0]**2 - (z-B[2])**2)
        y = (1/B[0])*(B[2]*(1-np.cos(t*B[0])) + B[1]*np.sin(t*B[0]))
        return y

    def zEOM(self, B, t):
        """
            B= [w, Vy, Vz]
        """
        z = (1/B[0])*(-B[1]*(1-np.cos(t*B[0])) + B[2]*np.sin(t*B[0]))
        return z

    def pol(self, B, x):
        return B[0] + B[1]*x + B[2]*(x**2) + B[3]*(x**3)


    """
    Functions running the ODR fits:
        We fit the collision points of the particle with the detector plate
        with respect to time.
        The fits are performed by the scipy.odr package. Initial guesses for
        parameters are specified by beta[0]. See fitting functions for a
        description of the parameters.

    """
    def xFit(self):
        xModel  = odr.Model(self.xEOM)
        xData   = odr.RealData(self.tPoints['Mean'],
                               self.xPoints['Mean'],
                               sx = self.tPoints['Error'],
                               sy = self.xPoints['Error'])
        xODR = odr.ODR(xData, xModel,beta0=[1])# beta0=[1, 100, 100, 100])
        xOUT = xODR.run()
        #xOUT.pprint()
        print('The x fit was performend with chi2 = ', xOUT.res_var)
        #print(xOUT.sd_beta)
        return xOUT
    
    def yFit(self):
        yModel  = odr.Model(self.yEOM)
        yData   = odr.RealData(self.tPoints['Mean'],
                               self.yPoints['Mean'],
                               sx = self.tPoints['Error'],
                               sy = self.yPoints['Error'])
        yODR = odr.ODR(yData, yModel, beta0=[1, 10, 100])
        yOUT = yODR.run()
        print('The y fit was performend with chi2 = ', yOUT.res_var)
        #yOUT.pprint()
        return yOUT

    def zFit(self):
        zModel  = odr.Model(self.zEOM)
        zData   = odr.RealData(self.tPoints['Mean'],
                               self.zPoints['Mean'],
                               sx = self.tPoints['Error'],
                               sy = self.zPoints['Error'])
        zODR = odr.ODR(zData, zModel, beta0=[1, 1, 100])
        zOUT = zODR.run()
        print('The z fit was performend with chi2 = ', zOUT.res_var)
        #zOUT.pprint()
        #print('z :', zOUT.res_var)
        return zOUT
    
    def output(self):
        """
        A function to print results to terminal.
        """
        print(self.results)

    def save(self, filename = 'results'):
        """
        A function to save the analysis results to a .csv file.
        The Format is:
            w   vx  vy  vz
            [...values...]
            [...errors...]
        The name of the file is set by 
            'filename'.
        Default, 'results'
        """
        self.results.to_csv(path_or_buf = filename+'.csv')
        
    def plot(self):

        times = pd.Series(np.linspace(0, self.tPoints["Mean"][len(self.tPoints["Mean"]) - 1], 1000))
        x_eom = np.array(self.xEOM([self.results["vx"][0]], times))
        y_eom = np.array(self.yEOM([self.results["w"][0], self.results["vy"][0], self.results["vz"][0]], times), dtype=pd.Series)
        z_eom = np.array(self.zEOM([self.results["w"][0], self.results["vy"][0], self.results["vz"][0]], times), dtype=pd.Series)
        spaces = np.zeros(5)

        fig = plt.figure()

        # Make a 3D plot
        ax1 = fig.add_subplot(1, 2, 1, projection="3d")
        ax1.plot(x_eom, y_eom, z_eom)
        ax1.plot(self.xPoints["Mean"], self.yPoints["Mean"], 'o')
        ax1.plot(self.xPoints["Mean"], spaces, self.zPoints["Mean"], 'o')
        ax1.plot(spaces , self.yPoints["Mean"], self.zPoints["Mean"], 'o')
        ax1.set_xlabel('x')
        ax1.set_ylabel('y')
        ax1.set_zlabel('z')
        ax1.set_xlim([np.min(x_eom), np.max(x_eom)])
        ax1.set_ylim([np.min(y_eom), np.max(y_eom)])
        ax1.set_zlim([np.min(z_eom), np.max(z_eom)])

        # Plot the x-z plane
        ax2 = fig.add_subplot(2, 2, 2)
        ax2.set_xlabel('z')
        ax2.set_ylabel('x', rotation=0)
        # ax2.set_xlimit(np.min(x_eom))
        # ax2.set_ylimit(np.min(z_eom))

        ax2.plot(z_eom, x_eom)
        ax2.plot(self.zPoints['Mean'], self.xPoints['Mean'], 'o')
        # ax2.savefig('x.pdf', format='pdf')
        # ax2.gcf().clear()

        # Plot the y-z plane
        ax3 = fig.add_subplot(2, 2, 4)
        ax3.set_xlabel('z')
        ax3.set_ylabel('y', rotation=0)
        # ax3.set_xlimit(np.min(y_eom))
        # ax3.set_ylimit(np.min(z_eom))

        ax3.plot(z_eom, y_eom)
        ax3.plot(self.zPoints["Mean"], self.yPoints['Mean'], 'o')
        # ax3.plt.savefig('y.pdf', format='pdf')
        # ax3.plt.gcf().clear()

        plt.subplots_adjust(left=0, right=0.99)
        fig.tight_layout()

        plt.savefig('3D.png', format='png')