|
|
# flake8: noqaq
import osimport argparseimport ROOTfrom ROOT import TMVA, TList, TTree
def res_train_matching_ghost_mlp( input_file: str = "data/ghost_data_B.root", tree_name: str = "PrMatchNN.PrMCDebugMatchToolNN/MVAInputAndOutput", exclude_electrons: bool = False, only_electrons: bool = True, residuals: str = "None", n_train_signal: int = 2e3, # 50e3 n_train_bkg: int = 5e3, # 500e3 n_test_signal: int = 1e3, n_test_bkg: int = 2e3, prepare_data: bool = True, outdir: str = "nn_electron_training",): """Trains an MLP to classify the match between Velo and Seed track.
Args: input_file (str, optional): Defaults to "data/ghost_data.root". tree_name (str, optional): Defaults to "PrMatchNN.PrMCDebugMatchToolNN/Tuple". exclude_electrons (bool, optional): Defaults to False. only_electrons (bool, optional): Signal only of electrons, but bkg of all particles. Defaults to True. residuals (bool, optional): Signal only of mlp<0.215. Defaults to False. n_train_signal (int, optional): Number of true matches to train on. Defaults to 200e3. n_train_bkg (int, optional): Number of fake matches to train on. Defaults to 200e3. n_test_signal (int, optional): Number of true matches to test on. Defaults to 75e3. n_test_bkg (int, optional): Number of fake matches to test on. Defaults to 75e3. prepare_data (bool, optional): Split data into signal and background file. Defaults to False. """
if prepare_data and not residuals == "None": resrdf = ROOT.RDataFrame(residuals, input_file) rdf = ROOT.RDataFrame(tree_name, input_file) if exclude_electrons: rdf_signal = rdf.Filter( "quality == 1 && mlp<0.215", # -1 elec, 0 ghost, 1 all part wo elec "Signal is defined as one label (excluding electrons)", ) rdf_bkg = rdf.Filter( "quality == 0 && mlp<0.215", "Ghosts are defined as zero label", ) resrdf_signal = resrdf.Filter( "quality == 1", # -1 elec, 0 ghost, 1 all part wo elec "Signal is defined as one label (excluding electrons)", ) resrdf_bkg = resrdf.Filter( "quality == 0", "Ghosts are defined as zero label", ) else: if only_electrons: rdf_signal = rdf.Filter( "quality == -1 && mlp<0.215", # electron that is true match but mlp said no match "Signal is defined as one label (only electrons)", ) resrdf_signal = resrdf.Filter( "quality == -1", # electron that is true match but mlp said no match "Signal is defined as one label (only electrons)", ) else: rdf_signal = rdf.Filter( "abs(quality) > 0 && mlp<0.215", "Signal is defined as non-zero label", ) resrdf_signal = resrdf.Filter( "abs(quality) > 0", "Signal is defined as non-zero label", ) rdf_bkg = rdf.Filter( "quality == 0 && mlp<0.215", "Ghosts are defined as zero label", ) resrdf_bkg = resrdf.Filter( "quality == 0", "Ghosts are defined as zero label", )
rdf_signal.Snapshot( "Signal", outdir + "/" + input_file.strip(".root") + "_mlp_matching_signal.root", ) rdf_bkg.Snapshot( "Bkg", outdir + "/" + input_file.strip(".root") + "_mlp_matching_bkg.root", ) resrdf_signal.Snapshot( "Signal", outdir + "/" + input_file.strip(".root") + "_res_matching_signal.root", ) resrdf_bkg.Snapshot( "Bkg", outdir + "/" + input_file.strip(".root") + "_res_matching_bkg.root", )
mlp_signal_file = ROOT.TFile.Open( outdir + "/" + input_file.strip(".root") + "_mlp_matching_signal.root", "READ", ) mlp_signal_tree = mlp_signal_file.Get("Signal")
mlp_bkg_file = ROOT.TFile.Open( outdir + "/" + input_file.strip(".root") + "_mlp_matching_bkg.root", "READ", ) mlp_bkg_tree = mlp_bkg_file.Get("Bkg") ####
res_signal_file = ROOT.TFile.Open( outdir + "/" + input_file.strip(".root") + "_res_matching_signal.root", "READ", ) res_signal_tree = res_signal_file.Get("Signal")
outputsignalFile = ROOT.TFile( outdir + "/" + input_file.strip(".root") + "_merged_matching_signal.root", "RECREATE", ) signaltreeList = TList()
signaltreeList.Add(mlp_signal_tree) signaltreeList.Add(res_signal_tree)
mergedsignalTree = TTree.MergeTrees(signaltreeList) mergedsignalTree.Write() outputsignalFile.Close()
res_bkg_file = ROOT.TFile.Open( outdir + "/" + input_file.strip(".root") + "_res_matching_bkg.root", "READ", ) res_bkg_tree = res_bkg_file.Get("Bkg")
outputbkgFile = ROOT.TFile( outdir + "/" + input_file.strip(".root") + "_merged_matching_bkg.root", "RECREATE", ) bkgtreeList = TList()
bkgtreeList.Add(mlp_bkg_tree) bkgtreeList.Add(res_bkg_tree)
mergedbkgTree = TTree.MergeTrees(bkgtreeList) mergedbkgTree.Write() outputbkgFile.Close()
#####
signal_file = ROOT.TFile.Open( outdir + "/" + input_file.strip(".root") + "_merged_matching_signal.root", "READ", ) signal_tree = signal_file.Get("Signal")
bkg_file = ROOT.TFile.Open( outdir + "/" + input_file.strip(".root") + "_merged_matching_bkg.root" ) bkg_tree = bkg_file.Get("Bkg")
###
os.chdir(outdir + "/result") output = ROOT.TFile( "matching_ghost_mlp_training.root", "RECREATE", )
factory = TMVA.Factory( "TMVAClassification", output, "V:!Silent:Color:DrawProgressBar:AnalysisType=Classification", ) factory.SetVerbose(True) dataloader = TMVA.DataLoader("MatchNNDataSet")
dataloader.AddVariable("chi2", "F") dataloader.AddVariable("teta2", "F") dataloader.AddVariable("distX", "F") dataloader.AddVariable("distY", "F") dataloader.AddVariable("dSlope", "F") dataloader.AddVariable("dSlopeY", "F")
dataloader.AddSignalTree(signal_tree, 1.0) dataloader.AddBackgroundTree(bkg_tree, 1.0)
# these cuts are also applied in the algorithm preselectionCuts = ROOT.TCut( # "chi2<30 && distX<500 && distY<500 && dSlope<2.0 && dSlopeY<0.15", #### ganz raus für elektronen ) dataloader.PrepareTrainingAndTestTree( preselectionCuts, f"SplitMode=random:V:nTrain_Signal={n_train_signal}:nTrain_Background={n_train_bkg}:nTest_Signal={n_test_signal}:nTest_Background={n_test_bkg}", # normmode default is EqualNumEvents )
factory.BookMethod( dataloader, TMVA.Types.kMLP, "matching_mlp", "!H:V:TrainingMethod=BP:NeuronType=ReLU:EstimatorType=CE:VarTransform=Norm:NCycles=700:HiddenLayers=N+2,N:TestRate=50:Sampling=1.0:SamplingImportance=1.0:LearningRate=0.02:DecayRate=0.01:!UseRegulator", ) factory.TrainAllMethods() factory.TestAllMethods() factory.EvaluateAllMethods() output.Close()
if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--input-file", type=str, help="Path to the input file", required=False, ) parser.add_argument( "--exclude_electrons", action="store_true", help="Excludes electrons from training.", required=False, ) parser.add_argument( "--only_electrons", action="store_true", help="Only electrons for signal training.", required=False, ) parser.add_argument( "--n-train-signal", type=int, help="Number of training tracks for signal.", required=False, ) parser.add_argument( "--n-train-bkg", type=int, help="Number of training tracks for bkg.", required=False, ) parser.add_argument( "--n-test-signal", type=int, help="Number of testing tracks for signal.", required=False, ) parser.add_argument( "--n-test-bkg", type=int, help="Number of testing tracks for bkg.", required=False, )
args = parser.parse_args() args_dict = {arg: val for arg, val in vars(args).items() if val is not None}
res_train_matching_ghost_mlp(**args_dict)
|
