prepareAsicCrosstalkSimDB.py

 
'''
   Prepare CDC x-talk simulation  DB object, usage:
 
   python3 prepareAsicCrosstalkSimDB.py <input_asic_root> <output_path> <exp_num> <data_type>
 
     input_asic_root:   input path and files containing asic information,
                        produced by AsicBackgroundLibraryCreator.py
     output_path: output path
     exp_num:   experiment number of raw data or mc samples
     data_type: data or mc
 
'''
 
import uproot
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import UnivariateSpline
 
import basf2
from ROOT.Belle2 import FileSystem
from ROOT.Belle2 import CDCDatabaseImporter
from ROOT import TH1F, TFile
import sys
 
 
def getEff(var, cut, n=40, limits=(0., 2500.)):
    ''' Simple efficiency estimator
       var: pandas series/np array of variable to study
       cut: bool series/np array pass/fail
       n  : number of bins
       limits : histogram limits
    '''
    a = np.histogram(var, n, limits)
    b = np.histogram(var[cut], n, limits)
    eff = b[0]/a[0]
    # Simple binomial formula:
    effErr = 1/a[0]*np.sqrt(a[0]*eff*(1-eff))
    x = 0.5*(a[1][1:]+a[1][:-1])
    return x, eff, effErr
 
 
# Problematic boards ID in each run
bad_boards = {'22': [24, 196],
              '24': [0, 15, 24, 40, 117, 175, 196, 202, 86, 89],
              '25': [0, 24, 15],
              '26': [24, 210, 62],
              '27': [92, 24, 51, 106, 53],
              }
# Input arguments
if len(sys.argv) != 5:
    sys.exit("Four arguments are required: input_root, output_path, experiemnt numebr and data_type")
InputFile = sys.argv[1]
OutputPath = sys.argv[2]
exp = sys.argv[3]
data_type = sys.argv[4]
 
# Dataframe, containing relevant variables from the root file
df = uproot.open(InputFile, flatten=True)["ASIC"].arrays(["Channel", "ADC", "Board", "Nhit", "Asic"], library="pd")
df.columns = ['_'.join(col) if col[1] != '' else col[0] for col in df.columns.values]
# Defining asic
df['asic'] = df.Channel//8
# Remove the problematic boards
mask = df['Board'].isin([m for m in bad_boards[f'{exp}']])
df = df[~mask]
 
nhits = 3
 
u1 = getEff(df[(df.asic % 3 == 1)].ADC_ADC_Sig, df.Nhit > f'{nhits}', 128, (0, 1024.))
 
u2 = getEff(df[(df.asic % 3 == 1)].ADC_ADC_Sig, df.Nhit > f'{nhits}', 8, (1024, 4096.))
 
u = np.append(u1, u2, axis=1)
 
plt.figure(figsize=(10, 4))
 
x = np.nan_to_num(u[0])
 
e = np.nan_to_num(u[1])
 
ee = np.where(np.nan_to_num(u[2]) == 0, 1000., u[2])
 
 
f = UnivariateSpline(x, e, 1/ee)
 
xp = np.arange(-0.5, 4096.5, 1)
 
plt.subplot(121)
plt.xlim(0, 4096.)
plt.ylim(0., 1.02)
plt.errorbar(x, e, u[2], fmt='.')
plt.ylabel(f'Fraction of Nhit>{nhits}')
plt.title(f'exp{exp} {data_type}')
plt.plot(xp, f(xp))
 
plt.subplot(122)
plt.xlim(0, 1000.)
plt.ylim(0., 0.5)
plt.errorbar(x, e, u[2], fmt='.', label='eff: asic=1,4; Nhit>3')
plt.plot(xp, f(xp), label='spline fit')
plt.xlabel('ADC')
plt.legend()
plt.savefig(f'{OutputPath}/xTalkProb_{exp}_{data_type}.pdf')
 
# Write out root file
 
names = ["Board", "Channel", 'Nhit', 'asic']
for i in range(8):
    names += [f'Asic_ADC{i:d}', f'Asic_TDC{i:d}', f'Asic_TOT{i:d}']
 
with uproot.recreate(f'{OutputPath}/xTalkProb_{exp}_{data_type}.root') as output:
    df_tmp = df.query('nhit>3 & asic==1')
    output['ASIC'] = uproot.newtree({name: df_tmp[name].dtype for name in names})
    output['ASIC'].extend({name: df_tmp[name].values for name in names})
 
 
fi = TFile(f'{OutputPath}/xTalkProb_{exp}_{data_type}.root', "update")
 
t = TH1F("ProbXTalk", "Prob xTalk", 4096, 0, 4096)
t.SetContent(f(xp))
fi.Write()
fi.Close()
 
 
INPUT = FileSystem.findFile(f'{OutputPath}/xTalkProb_{exp}_{data_type}.root')
 
# Specify the exp and run where iov is valid.
# N.B. -1 means unbound.
 
expFirst = 0
 
expLast = -1
 
runFirst = 0
 
runLast = -1
# basf2.use_local_database("localdb/database.txt", "localdb")
basf2.conditions.testing_payloads = [f'localdb/exp{exp}/database_{exp}.txt']
 
 
main = basf2.create_path()
 
 
eventinfosetter = basf2.register_module('EventInfoSetter')
main.add_module(eventinfosetter)
 
# process single event
basf2.process(main)
 
 
dbImporter = CDCDatabaseImporter(expFirst, runFirst, expLast, runLast)
dbImporter.importCDCCrossTalkLibrary(INPUT)
 
dbImporter.printCDCCrossTalkLibrary()
dbImporter.testCDCCrossTalkLibrary()