NewCode2024/macros/split_tree.C

0001 void split_tree() {
0002
0003     string input_directory = "/sphenix/user/ChengWei/sPH_dNdeta/Run24AuAuMC/Sim_Ntuple_HIJING_ana443_20241102";
0004     string input_filename = "ntuple_merged.root";
0005     string TreeName = "EventTree";
0006     int nEvent_EachFile = 500;
0007
0008     // gSystem->Load("libg4dst");
0009
0010     // string input_directory = "/sphenix/tg/tg01/commissioning/INTT/work/cwshih/seflgendata/run_23911/INTTRawHit/completed";
0011     // string input_filename = "inttrawhit_00023911_00000_syncDST_600k.root";
0012     // string TreeName = "T";
0013     // int nEvent_EachFile = 1000;
0014
0015     string input_filename_NoSuffix = input_filename.substr(0, input_filename.find("."));
0016
0017     // note : Open the input ROOT file
0018     TFile *inputFile = TFile::Open((input_directory + "/" + input_filename).c_str(), "READ");
0019     if (!inputFile || inputFile->IsZombie()) {
0020         std::cerr << "Error: Could not open input file!" << std::endl;
0021         return;
0022     }
0023
0024     // note : Get the TTree from the file
0025     TTree *inputTree = dynamic_cast<TTree*>(inputFile->Get(TreeName.c_str())); // note : Replace "tree" with your TTree's name
0026     if (!inputTree) {
0027         std::cerr << "Error: TTree not found in the input file!" << std::endl;
0028         inputFile->Close();
0029         return;
0030     }
0031
0032     // note : Get the total number of events
0033     Long64_t nTotalEvents = inputTree->GetEntries();
0034
0035     // note : Calculate the number of full files
0036     int nFiles = nTotalEvents / nEvent_EachFile;
0037     // int remainingEvents = nTotalEvents % nEvent_EachFile;
0038
0039     std::cout << "Total events: " << nTotalEvents << std::endl;
0040     std::cout << "Events per file: " << nEvent_EachFile << std::endl;
0041     std::cout << "n output files: " << nFiles << std::endl;
0042     // std::cout << "Remaining events: " << remainingEvents << std::endl;
0043
0044     // note : Loop to create output files
0045     for (int i = 0; i < nFiles; ++i) {
0046         // note : Define the event range for this file
0047         Long64_t startEvent = i * nEvent_EachFile;
0048         Long64_t endEvent = (i == nFiles - 1) ? nTotalEvents - 1 : startEvent + nEvent_EachFile - 1;
0049
0050         std::cout<<"In core: "<<i<<" startEvent: "<<startEvent<<" endEvent: "<<endEvent<<std::endl;
0051
0052         // note : Skip the last file if it has no events
0053         if (startEvent >= nTotalEvents) break;
0054
0055         std::string job_index = std::to_string( i );
0056         int job_index_len = 5;
0057         job_index.insert(0, job_index_len - job_index.size(), '0');
0058
0059         // note : Create output file name
0060         string output_filename = input_filename_NoSuffix + "_" + job_index + ".root";
0061         TFile *outputFile = TFile::Open((input_directory+"/"+output_filename).c_str(), "RECREATE");
0062
0063         // note : Create a new TTree with the subset of events
0064         TTree *outputTree = inputTree->CopyTree(Form("Entry$ >= %lld && Entry$ <= %lld", startEvent, endEvent));
0065
0066         // note : Write the new TTree to the output file
0067         outputTree->Write();
0068         outputFile->Close();
0069         std::cout << "Created file: " << output_filename << " with events [" << startEvent << " - " << endEvent << "]" << std::endl;
0070     }
0071
0072     // note : Close the input file
0073     inputFile->Close();
0074 }