NewCode2024/macros/split_tree_TChain.C

0001 #include <iostream>
0002 #include <fstream>
0003 #include <vector>
0004 #include <string>
0005
0006 std::vector<std::string> readFileToVector(const std::string& filePath) {
0007     std::vector<std::string> lines; // Vector to store lines
0008     std::ifstream inputFile(filePath); // Open the file
0009
0010     if (!inputFile) {
0011         std::cerr << "Error: Unable to open file at " << filePath << std::endl;
0012         return lines; // Return an empty vector
0013     }
0014
0015     std::string line;
0016     while (std::getline(inputFile, line)) { // Read file line by line
0017         lines.push_back(line); // Add each line to the vector
0018     }
0019
0020     inputFile.close(); // Close the file
0021
0022     for (const auto& string_line : lines) {
0023         std::cout << string_line << std::endl; // Print each line
0024     }
0025
0026     return lines; // Return the vector of lines
0027 }
0028
0029
0030
0031 void split_tree_TChain(std::string generator_name, int FileList_index = 1, int File_shift_index = 0) {
0032
0033     std::map<std::string,std::string> sample_directory ={
0034         {"HIJING", "Sim_HIJING_MDC2_ana472_20250307"}, // note : new
0035         {"AMPT", "Sim_AMPT_MDC2_ana472_20250310"}, // note : new
0036         {"EPOS", "Sim_EPOS_MDC2_ana472_20250310"}, // note : new
0037         {"HStrange", "Sim_HIJING_strangeness_MDC2_ana472_20250310"} // note : new
0038     };
0039
0040     string input_directory = Form("/sphenix/user/ChengWei/sPH_dNdeta/Run24AuAuMC/%s", sample_directory[generator_name].c_str());
0041     string output_directory = Form("/sphenix/user/ChengWei/sPH_dNdeta/Run24AuAuMC/%s/per5k", sample_directory[generator_name].c_str()); // note : auto
0042
0043     std::cout << "Input directory: " << input_directory << std::endl;
0044     std::cout << "Output directory: " << output_directory << std::endl;
0045
0046     system (Form("mkdir -p %s", output_directory.c_str()));
0047
0048     string input_filelist = Form("FileList_00%d.txt",FileList_index);
0049     string TreeName = "EventTree";
0050     int nEvent_EachFile = 5000;
0051     string input_filename_NoSuffix = "ntuple_per5k";
0052
0053     std::vector<std::string> file_list = readFileToVector(input_directory + "/" + input_filelist);
0054     // std::string first_filename = file_list[0].substr(file_list[0].find_last_of("/")+1);
0055     // string input_filename_NoSuffix = first_filename.substr(0, first_filename.find("."));
0056
0057     TChain * inputTree = new TChain(TreeName.c_str());
0058     for (const auto& file : file_list) {
0059         inputTree->Add((file).c_str());
0060     }
0061
0062     // note : Get the total number of events
0063     Long64_t nTotalEvents = inputTree->GetEntries();
0064
0065     // note : Calculate the number of full files
0066     int nFiles = nTotalEvents / nEvent_EachFile;
0067     // int remainingEvents = nTotalEvents % nEvent_EachFile;
0068
0069     std::cout << "Total events: " << nTotalEvents << std::endl;
0070     std::cout << "Events per file: " << nEvent_EachFile << std::endl;
0071     std::cout << "n output files: " << nFiles << std::endl;
0072     // std::cout << "Remaining events: " << remainingEvents << std::endl;
0073
0074     // note : Loop to create output files
0075     for (int i = 0; i < nFiles; ++i) {
0076         // note : Define the event range for this file
0077         Long64_t startEvent = i * nEvent_EachFile;
0078         Long64_t endEvent = (i == nFiles - 1) ? nTotalEvents - 1 : startEvent + nEvent_EachFile - 1;
0079
0080         std::cout<<"In core: "<<i<<" startEvent: "<<startEvent<<" endEvent: "<<endEvent<<std::endl;
0081
0082         // note : Skip the last file if it has no events
0083         if (startEvent >= nTotalEvents) break;
0084
0085         std::string job_index = std::to_string( i + File_shift_index );
0086         int job_index_len = 5;
0087         job_index.insert(0, job_index_len - job_index.size(), '0');
0088
0089         // note : Create output file name
0090         string output_filename = input_filename_NoSuffix + "_" + job_index + ".root";
0091         TFile *outputFile = TFile::Open((output_directory+"/"+output_filename).c_str(), "RECREATE");
0092
0093         // note : Create a new TTree with the subset of events
0094         TTree *outputTree = inputTree->CopyTree(Form("Entry$ >= %lld && Entry$ <= %lld", startEvent, endEvent));
0095
0096         // note : Write the new TTree to the output file
0097         outputTree->Write();
0098         outputFile->Close();
0099         std::cout << "Created file: " << output_filename << " with events [" << startEvent << " - " << endEvent << "]" << std::endl;
0100     }
0101
0102     // note : Close the input file
0103     // inputFile->Close();
0104 }