Пример #1
0
//std::mutex mtx;
void ProcessFilePar(TString fileIn, TString fileOut, TString treename,  vector<TString> friends, vector<TString> branches, vector<TString> newbranches, unsigned jobid = 0, unsigned NPAR = 1)
{
  //mtx.lock(); //for threads
  TFile *fin = new TFile(fileIn);
  TTree *tjet = (TTree*)fin->Get(treename);
  //mtx.unlock();

  vector<TTree *> friendTrees;
  vector<bool> sameFileFriend;
  for (auto f:friends) {
    auto fr = tokenize(f,":");
    if (fr.size()==1) {tjet->AddFriend(fr[0]); sameFileFriend.push_back(true);}
    else if (fr.size()==2) {tjet->AddFriend(fr[1],fr[0]); sameFileFriend.push_back(false);}

    TTree *tfriend = (TTree*)fin->Get(f);
    friendTrees.push_back(tfriend);
  }

  AddBranchesByCounter(tjet, branches);
  for (unsigned i=0;i<friendTrees.size();i++) AddBranchesByCounter(friendTrees[i],branches);

  //sort branches into categories
  for (auto bName:branches) {
    TBranch *b=tjet->GetBranch(bName);
    if (b==0) cout <<"Branch "<<bName<<" doesn't exist!"<<endl;

    //parse in case there is a tree name in the branch
    auto branchtoken = tokenize(bName,".");
    auto leafname = branchtoken[branchtoken.size()-1];

    TObjArray *bl = b->GetListOfLeaves();
    if (bl->GetEntries()>1)
      cout <<" Branch "<<b->GetTitle()<<" has more than 1 leave. Taking first..."<<endl;
    if (bl->GetEntries()==0) {
      cout <<" Branch "<<b->GetTitle()<<" has no leaves! Skipping..."<<endl;
      continue;
    }

    TLeaf * l = (TLeaf *)(*bl)[0];
    //what's the type?
    TString type = l->GetTypeName();
    if (VERBOSE) cout<<l->GetTitle()<<endl;

    //array?
    bool array = l->GetLeafCount()!=0;
    TString counter;
    if (array) counter = l->GetLeafCount()->GetBranch()->GetName();

    if (type=="Float_t")
      {  if (array) {brVFloat.push_back(bName); brVFloatCounter.push_back(counter); }else brFloat.push_back(bName);}
    else if (type=="Int_t")
      {  if (array) {brVInt.push_back(bName); brVIntCounter.push_back(counter); }else brInt.push_back(bName);}
    else cout << "Unsupported branch type "<<type<<" for branch "<<bName<<". Skipping..."<<endl;
  }


  //treat counters as ints only
  AppendToListUniquely(brVIntCounter, brInt);
  AppendToListUniquely(brVFloatCounter, brInt);

  //too keep track of old branches, which cannot be read (todo: just check it...)
  int noldbrInt = brInt.size(), noldbrFloat = brFloat.size(), noldbrVInt = brVInt.size(), noldbrVIntCounter = brVIntCounter.size(), noldbrVFloat = brVFloat.size(), noldbrVFloatCounter = brVFloatCounter.size();
  //add new branches
  ParseNewBranches(newbranches, brInt, brFloat, brVInt, brVIntCounter, brVFloat, brVFloatCounter);

  //print for debugging
  if (VERBOSE) {
    cout<<"int       : "; for (auto s:brInt) cout <<s<<", "; cout<<endl;
    cout<<"float     : "; for (auto s:brFloat) cout <<s<<", "; cout<<endl;
    cout<<"Vint      : "; for (auto s:brVInt) cout <<s<<", "; cout<<endl;
    cout<<"Vfloat    : "; for (auto s:brVFloat) cout <<s<<", "; cout<<endl;
    cout<<"Vintcnt   : "; for (auto s:brVIntCounter) cout <<s<<", "; cout<<endl;
    cout<<"Vfloatcnt : "; for (auto s:brVFloatCounter) cout <<s<<", "; cout<<endl;
  }

  tjet->SetBranchStatus("*",1);
  for (auto b:brVFloat) if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,0);
  for (auto b:brFloat)  if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,0);
  for (auto b:brVInt)   if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,0);
  for (auto b:brInt)    if (tjet->GetBranch(b)!=0) {unsigned f=0; tjet->SetBranchStatus(b,0,&f); if (VERBOSE) cout<<"turning off "<<b<<", found = "<<f<<endl;}


  TFile *fout = new TFile(fileOut,"recreate");
  TTree *tjetout;


  //in case of one-to-many event - do not copy branches automatically!
  if (useOneToOne) tjetout = tjet->CloneTree(0);
  else tjetout = new TTree(tjet->GetName(),tjet->GetTitle());


  //think about memory tree
  // tjetout->SetDirectory(0);
  tjetout->SetName(tjet->GetName());
  //TTree *tjetout = new TTree("t","t");

  //to see what is copied...
  //tjetout->Print();

  for (auto b:brVFloat) if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,1);
  for (auto b:brFloat)  if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,1);
  for (auto b:brVInt)   if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,1);
  for (auto b:brInt)    if (tjet->GetBranch(b)!=0) tjet->SetBranchStatus(b,1);

  vector<int> valIntIn(brInt.size()), valIntOut(brInt.size());
  vector<float> valFloatIn(brFloat.size()), valFloatOut(brFloat.size());
  vector<vector<int> >valVIntIn (brVInt.size());  vector<vector<int> >valVIntOut (brVInt.size());
  vector<vector<float> >valVFloatIn (brVFloat.size());  vector<vector<float> >valVFloatOut (brVFloat.size());

  for (unsigned i=0;i<brInt.size();i++) {
    if (tjet->GetBranch(brInt[i])!=0)
      tjet->SetBranchAddress(brInt[i],&valIntIn[i]);

    if (tjetout->GetBranch(brInt[i])!=0) {//why would it?
      tjetout->SetBranchAddress(brInt[i],&valIntOut[i]); 
      cout<<"branch "<<brInt[i]<<" already exist for some reason..."<<endl; 
    }
    else //logical...
      if (NonFriendBranch(tjet, brInt[i])) 
	tjetout->Branch(brInt[i],&valIntOut[i],Form("%s/I",brInt[i].Data()));
  }

  for (unsigned i=0;i<brFloat.size();i++) {
    if (tjet->GetBranch(brFloat[i])!=0) 
      tjet->SetBranchAddress(brFloat[i],&valFloatIn[i]);
    
    if (NonFriendBranch(tjet, brFloat[i]))
      tjetout->Branch(brFloat[i],&valFloatOut[i],Form("%s/F",brFloat[i].Data()));
  }

  for (unsigned i=0;i<brVFloat.size();i++) {
    if (tjet->GetBranch(brVFloat[i])!=0) {
      valVFloatIn[i] = vector<float>(NMAX);
      tjet->SetBranchAddress(brVFloat[i],&valVFloatIn[i][0]);
    }
    

    valVFloatOut[i] = vector<float>(NMAX);
    if (NonFriendBranch(tjet, brVFloat[i]))
      tjetout->Branch(brVFloat[i],&valVFloatOut[i][0],Form("%s[%s]/F",brVFloat[i].Data(),brVFloatCounter[i].Data()));
  }

  for (unsigned i=0;i<brVInt.size();i++) {
    if (tjet->GetBranch(brVInt[i])) {
      valVIntIn[i] = vector<int>(NMAX);
      tjet->SetBranchAddress(brVInt[i],&valVIntIn[i][0]);
    }
    valVIntOut[i] = vector<int>(NMAX);
    if (NonFriendBranch(tjet, brVInt[i]))
      tjetout->Branch(brVInt[i],&valVIntOut[i][0],Form("%s[%s]/I",brVInt[i].Data(),brVIntCounter[i].Data()));
  }

  Long64_t nentries = tjet->GetEntries();
  int nentries1 = nentries/NPAR*jobid;
  int nentries2 = nentries/NPAR*(jobid+1);
  
  nentries = nentries2-nentries1; 
  int oneperc = nentries/100; if (oneperc==0) oneperc = 1;

  cout<<"Start processing..."<<endl;
  TStopwatch s;
  s.Start();
  TTimeStamp t0;
  double readTime = 0, processingTime = 0, copyToTime = 0, cloneTime=0, copyFromTime=0, fillTime = 0;
  

  for (Long64_t i=0; i<nentries;i++) {
    if (jobid==0 && i % oneperc == 0 && i>0) {
      std::cout << std::fixed;
      TTimeStamp t1; cout<<" \r"<<i/oneperc<<"%   "<<" est. time "<<setprecision(2) <<(t1-t0)*nentries/(i+.1)<<" s ";
      cout<<";Processing:"<<setprecision(2)<<processingTime/(t1-t0)*100<<" %";
      cout<<";Copy1:"<<setprecision(2) <<copyToTime/(t1-t0)*100<<" %";
      cout<<";Clone:"<<setprecision(2) <<cloneTime/(t1-t0)*100<<" %";
      cout<<";Copy2:"<<setprecision(2) <<copyFromTime/(t1-t0)*100<<" %";
      cout<<";Fill:"<<setprecision(2) <<fillTime/(t1-t0)*100<<" %";
      cout<<";Read:"<<setprecision(2) <<readTime/(t1-t0)*100<<" %";
      cout<<flush; 
    }
    
    TTimeStamp tsRead0;
    tjet->GetEntry(i+nentries1);
    TTimeStamp tsRead1;
    readTime+=tsRead1-tsRead0;

    Everything ev;
    TTimeStamp tsCpTo0;
    for (unsigned j=0;j<brInt.size();j++) ev.PutInt(brInt[j],valIntIn[j]);
    for (unsigned j=0;j<brFloat.size();j++) ev.PutFloat(brFloat[j],valFloatIn[j]);
    for (unsigned j=0;j<brVFloat.size();j++) ev.PutVFloat(brVFloat[j],brVFloatCounter[j],valVFloatIn[j]);
    for (unsigned j=0;j<brVInt.size();j++) ev.PutVInt(brVInt[j],brVIntCounter[j],valVIntIn[j]);
    TTimeStamp tsCpTo1;
    copyToTime+=tsCpTo1-tsCpTo0;
    

    TTimeStamp tsCl0;
    //think about: copy object (timing 10% ->3%) 
    //but it copies vectors, so push_back will add in the end...
    //    Everything evout = ev;
    //or even reference(in place?) (->0.2%)
    //Everything &evout = ev;
    Everything evout = ev.CloneStructure();
    TTimeStamp tsCl1;
    cloneTime+=tsCl1-tsCl0;

    bool exitEvent = false;
    int counter = 0;

    while (!exitEvent) {
    
    TTimeStamp tsPr0;
    if (useOneToOne) {
      fProcessOneToOne(ev, evout);
      evout.UpdateCounters();
      exitEvent = true;
    } else {
      exitEvent = fProcessOneToMany(ev, evout, counter); 
      //      if (!exitEvent) cout<<"event to write "<<counter<<endl;
      counter++;
    }

    TTimeStamp tsPr1;  
    processingTime+=tsPr1-tsPr0;
    
    
    //Everything evout = ev;
    TTimeStamp tsCpFrom0;
    for (unsigned j=0;j<brInt.size();j++) valIntOut[j] = evout.GetInt(brInt[j]);
    for (unsigned j=0;j<brFloat.size();j++) {valFloatOut[j] = evout.GetFloat(brFloat[j]); 
      //  cout<<brFloat[j]<<" "<<evout.GetFloat(brFloat[j])<<endl;
    }
    for (unsigned j=0;j<brVFloat.size();j++) valVFloatOut[j] = evout[brVFloat[j]];
    for (unsigned j=0;j<brVInt.size();j++) valVIntOut[j] = evout.GetVInt(brVInt[j]);
    TTimeStamp tsCpFrom1;
    copyFromTime+=tsCpFrom1-tsCpFrom0;

    TTimeStamp tsFill0;
    tjetout->Fill();
    TTimeStamp tsFill1;
    fillTime+=tsFill1-tsFill0;
  }

  }
  cout<<endl;
  s.Stop();
  cout<<"Done in ";s.Print();
  
  tjetout->FlushBaskets();
  tjetout->Write();
  
  cout<<"Copying other trees..."<<endl;

  for (unsigned i=0;i<friendTrees.size();i++) {
    TTree *t = friendTrees[i];
    if (sameFileFriend[i]) {
      //TTree *triendout = t->CloneTree(-1,"fast");
      TTree *triendout = t->CopyTree("","",nentries,nentries1);
      triendout->Write();
    }
  }
  


  fout->Close();
  friendTrees.clear();
}
Пример #2
0
bool TreeReader::Initialize(vector <string> br, string opt)
{
	if(!init)
	{
	    if( !fChain )
		{
		        cout << endl;
			cout << "No tree to initialize" << endl;
			cout << endl;
			return false;
		}

		TObjArray *fileElements = fChain->GetListOfFiles();
		if( !fileElements || ( fileElements->GetEntries() == 0 ))
		{
		        cout << endl;
			cout << "No file(s) to initialize" << endl;
			cout << endl;
			return false;
		}
	}

	varList.clear();

	TObjArray* branches = fChain->GetListOfBranches();
	int nBranches = branches->GetEntries();

	for (int i = 0; i < nBranches; ++i)
	{
		TBranch* branch = (TBranch*)branches->At(i);
		string brname = branch->GetName();
		TLeaf* leaf = branch->GetLeaf(branch->GetName());

		if ( leaf == 0 )  // leaf name is different from branch name
		{
			TObjArray* leafs = branch->GetListOfLeaves();
			leaf = (TLeaf*)leafs->At(0);
		}

		string curtype = leaf->GetTypeName();
        int id = TypeDB::getType(curtype.c_str());
		int arreysize = 1;
		string title = leaf->GetTitle();
        //cout << curtype << "   " << title << endl;

		// Find out whether we have array by inspecting leaf title
		if ( title.find("[")!=std::string::npos )
		{
			TLeaf * nelem = leaf->GetLeafCounter(arreysize);
			if(arreysize == 1 && nelem != NULL) arreysize = nelem->GetMaximum() + 1; //search for maximum value of the lenght
		}


		if(id >= 0)
		{
			bool addVar = true;
			if(br.size()>0)
			{
				addVar = false;
				for(unsigned b = 0; b < br.size(); b++)
				{
					if(opt == "names" || opt == "except")
					{
						if(br[b] == brname) { addVar = true; break;}
					}
					else if(opt.find("contains")!=string::npos)
					{
						if((string(brname)).find(br[b])!=string::npos) { addVar = true; break;}
					}
					else if(opt.find("except")==string::npos) cout << "Option " << opt << " not found" << endl;
				}

				if(opt.find("except")!=string::npos) addVar = !addVar;
			}

			if(addVar)
			{
				variable * tmpVar = new variable(id,arreysize);

				tmpVar->name = leaf->GetName();
				tmpVar->bname = branch->GetName();
				tmpVar->title = title;

				varList.push_back(tmpVar);
				fChain->SetBranchAddress(tmpVar->bname,tmpVar->value.address);
			}
		}
		else
		{
			cout << curtype << ": type not found" << endl;
			exit(1);
			return false;
		}
	}

	init = true;
	continueSorting = true;
	if(pmode=="v") cout << endl << "Set up " << varList.size() << " / " << nBranches << " branches" << endl;

	return true;
}