reconstruct.py

import pandas as pd
import numpy as np

import Clustering as trk
import Fitting as fit

N = 1000

if __name__ == '__main__':
    print("Reading file ...")

    fitter = fit.TrackFitter(B=1.)

    data_track = pd.DataFrame({'event':[0],'track':[0],'pt':[0.], 'phi':[0.], 'xVtx':[0.], 'yVtx':[0.]})
    data_track = data_track.drop(data_track.index[[0]])

    df = pd.read_csv("hits_"+str(N)+".csv")
    y_df = df[['particle']]
    #    X_df = df.drop(['hit','layer','particle','event'], axis=1)
    X_df = df.drop(['particle','layer','iphi'], axis=1)

#replace particle with 100000000*event+particle

    y_train = df['particle'].values + df['event'].values * 1000
    tracker = trk.ClusterDBSCAN(eps=0.004, rscale=0.001)
    
    tracker.fit(X_df.values[:1000], y_train[:1000])

    df_result = df.drop(['particle'], axis=1)
    y_predicted = tracker.predict(df_result.values)
    
    
    df_result=pd.concat([df_result,pd.DataFrame({'track':y_predicted})],axis=1)
    for col in ['event','track']:
        df_result[col] = df_result[col].astype('int32')
    df_result_truth=pd.concat([df_result,pd.DataFrame({'particle':y_train})],axis=1)
    df_result_truth['particle'] = df_result_truth['particle'].astype('int32')

    df_result.to_csv("result_"+str(N)+".csv",header=True,cols=['event','track', 'x', 'y'], engine='python')
    df_result_truth.to_csv("result_truth_"+str(N)+".csv",header=True,cols=['event','particle','track', 'x', 'y'], engine='python')
    
    
    # event loop for fitting
    for ievent in []: # np.unique(df['event'].values):
        if ievent % 100 == 0 : print "processing event ", ievent
        event_indices=(df['event']==ievent).values
        event_df = df.loc[event_indices]
 
        y_event_predicted = y_predicted[event_indices]
        #       y_event_predicted = y_event_predicted + [ievent*1000]*len(y_event_predicted)
        y_predicted = np.append(y_predicted,y_event_predicted)

        for itrack in np.unique(y_event_predicted):
            xhit = X_df['x'].values[y_event_predicted[:] == itrack]
            yhit = X_df['y'].values[y_event_predicted[:] == itrack]
    #         pt,phi,vx,vy,chg=fitter.fit(xhit,yhit)