[u/mrichter/AliRoot.git] / PWG / DevNanoAOD / AliNanoAODReplicator.cxx

/**************************************************************************
* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
*                                                                        *
* Author: The ALICE Off-line Project.                                    *
* Contributors are mentioned in the code where appropriate.              *
*                                                                        *
* Permission to use, copy, modify and distribute this software and its   *
* documentation strictly for non-commercial purposes is hereby granted   *
* without fee, provided that the above copyright notice appears in all   *
* copies and that both the copyright notice and this permission notice   *
* appear in the supporting documentation. The authors make no claims     *
* about the suitability of this software for any purpose. It is          *
* provided "as is" without express or implied warranty.                  *
**************************************************************************/

// $Id: AliNanoAODReplicator.cxx 56492 2012-05-15 18:42:47Z pcrochet $

// Implementation of a branch replicator 
// to produce nanoAODs.
//
//
// This replicator is in charge of replicating the tracks,vertices,headers
// branches of a standard AOD or ESD file into a nanoAODs 
// (AliAOD.Special.root)
// 
// The class was inspired by AliAODMuonReplicator
// 
// Author: Michele Floris, michele.floris@cern.ch


class AliESDv0;
class AliESDVertex;
class AliAODVertex;
class AliAODRecoDecay;

#include "AliAODDimuon.h"
#include "AliAODEvent.h"
#include "AliAODMCHeader.h"
#include "AliAODMCParticle.h"
#include "AliAODTZERO.h"
#include "AliAODTrack.h"
#include "AliAODVZERO.h"
#include "AliAnalysisCuts.h"
#include "TF1.h"
#include "AliExternalTrackParam.h"
#include "AliESDv0.h"
#include "AliAODv0.h"
#include "AliPIDResponse.h"
#include <iostream>
#include <cassert>
#include "AliESDtrack.h"
#include "TObjArray.h"
#include "AliAnalysisFilter.h"
#include "AliNanoAODTrack.h"

#include <TFile.h>
#include <TDatabasePDG.h>
#include <TString.h>
#include <TList.h>
#include "AliLog.h"
#include "AliVEvent.h"
#include "AliVVertex.h"
#include "AliVTrack.h"
#include "AliVertexerTracks.h"
#include "AliKFVertex.h"
#include "AliESDEvent.h"
#include "AliESDVertex.h"
#include "AliESDtrackCuts.h"
#include "AliAODEvent.h"
#include "AliAnalysisFilter.h"

#include "AliNanoAODReplicator.h"
#include "TH1.h"
#include "TCanvas.h"
#include "AliNanoAODHeader.h"
#include "AliNanoAODCustomSetter.h"

using std::cout;
using std::endl;

ClassImp(AliNanoAODReplicator)

//_____________________________________________________________________________
AliNanoAODReplicator::AliNanoAODReplicator() :
AliAODBranchReplicator(), 
  fTrackCut(0), fTracks(0x0), fHeader(0x0), fNTracksVariables(0), // FIXME: Start using cuts, and check if fNTracksVariables is needed
  fVertices(0x0), 
  fList(0x0),
  fMCParticles(0x0),
  fMCHeader(0x0),
  fMCMode(0),
  fLabelMap(),
  fParticleSelected(),
  fVarList(""),
  fVarListHeader(""),
  fCustomSetter(0){
  // Default ctor. we need it to avoid instantiating a wrong mapping when reading from file 
  }

AliNanoAODReplicator::AliNanoAODReplicator(const char* name, const char* title,
					   const char * varlist,
					   AliAnalysisCuts* trackCut,
					   Int_t mcMode
					     ) :
  AliAODBranchReplicator(name,title), 

  fTrackCut(trackCut), fTracks(0x0), fHeader(0x0), fNTracksVariables(0), // FIXME: Start using cuts, and check if fNTracksVariables is needed
  fVertices(0x0), 
  fList(0x0),
  fMCParticles(0x0),
  fMCHeader(0x0),
  fMCMode(mcMode),
  fLabelMap(),
  fParticleSelected(),
  fVarList(varlist),
  fVarListHeader(""),// FIXME: this should be set to a meaningful value: add an arg to the constructor
  fCustomSetter(0)
{
  // default ctor
  AliNanoAODTrackMapping * tm =new AliNanoAODTrackMapping(fVarList);
  fNTracksVariables = tm->GetSize();  
  //  tm->Print();
}

//_____________________________________________________________________________
AliNanoAODReplicator::~AliNanoAODReplicator()
{
  // dtor
  delete fTrackCut;
  delete fList;
}

//_____________________________________________________________________________
void AliNanoAODReplicator::SelectParticle(Int_t i)
{
  // taking the absolute values here, need to take care 
  // of negative daughter and mother
  // IDs when setting!
  
  if (!IsParticleSelected(TMath::Abs(i)))
  {
    fParticleSelected.Add(TMath::Abs(i),1);    
  }
}

//_____________________________________________________________________________
Bool_t AliNanoAODReplicator::IsParticleSelected(Int_t i)  
{
  // taking the absolute values here, need to take 
  // care with negative daughter and mother
  // IDs when setting!
  return (fParticleSelected.GetValue(TMath::Abs(i))==1);
}


//_____________________________________________________________________________
void AliNanoAODReplicator::CreateLabelMap(const AliAODEvent& source)
{  
  //
  // this should be called once all selections are done 
  // This method associates to the original index of the mc particle
  // (i) the new one (j). J runs only over particles which are
  // actually kept.
  //
  
  fLabelMap.Delete();
  
  TClonesArray* mcParticles = static_cast<TClonesArray*>(source.FindListObject(AliAODMCParticle::StdBranchName()));
  
  Int_t j(0);
  Int_t i(0); // We need i, we cannot rely on part->GetLabel, because some of the original mc particles are not kept in the stack, apparently
  
  TIter next(mcParticles);

  while ( next() )
  {
    if (IsParticleSelected(i))
    {
      fLabelMap.Add(i,j++);
      //      std::cout << i <<  "->" << j-1 << std::endl;
    }
    ++i;
  }  


}

//_____________________________________________________________________________
Int_t AliNanoAODReplicator::GetNewLabel(Int_t i) 
{
  // Gets the label from the new created Map
  // Call CreatLabelMap before
  // otherwise only 0 returned
  return fLabelMap.GetValue(TMath::Abs(i));
}

//_____________________________________________________________________________
void AliNanoAODReplicator::FilterMC(const AliAODEvent& source)
{
  // Filter MC information


  AliAODMCHeader* mcHeader(0x0);
  TClonesArray* mcParticles(0x0);
  
  fParticleSelected.Delete();

  //  std::cout << "MC Mode: " << fMCMode << ", Tracks " << fTracks->GetEntries() << std::endl;
  
  if ( fMCMode>=2 && !fTracks->GetEntries() ) {
    return;
  }
  // for fMCMode==1 we only copy MC information for events where there's at least one muon track
    
  mcHeader = static_cast<AliAODMCHeader*>(source.FindListObject(AliAODMCHeader::StdBranchName()));
  
  if ( mcHeader ) 
    {
      *fMCHeader = *mcHeader;
    }
  
  
  mcParticles = static_cast<TClonesArray*>(source.FindListObject(AliAODMCParticle::StdBranchName()));
  
  if ( mcParticles && fMCMode>=2 )
    {
      // keep all primaries
      TIter nextPart(mcParticles);
      static Int_t iev = -1; // FIXME: remove this (debug)
      iev ++;
      AliAODMCParticle * prim = 0;
      Int_t iprim = 0;  // We need iprim, we cannot rely on part->GetLabel, because some of the original mc particles are not kept in the stack, apparently
      // also select all charged primaries 
      while ((prim = (AliAODMCParticle*) nextPart())) {
	if(prim->IsPhysicalPrimary() && prim->Charge()) SelectParticle(iprim);
	// FIXME DEBUG
	if(iev == 2009) {
	  // std::cout << "IEV " << iev << std::endl;
	  // std::cout << " PART " << iprim << " " << prim->IsPhysicalPrimary() <<","<<prim->Charge() << "=" << IsParticleSelected(iprim) <<  std::endl;
	  if(iprim == 15) {
	    prim->Print();
	  }
	  
	}
	iprim++;
      } 

      // loop on (kept) tracks to find their ancestors
      TIter nextTRACK(fTracks);
      AliNanoAODTrack* track;
    
      while ( ( track = static_cast<AliNanoAODTrack*>(nextTRACK()) ) )
	{
	  Int_t label = TMath::Abs(track->GetLabel()); 
      
	  while ( label >= 0 ) 
	    {
	      SelectParticle(label);
	      AliAODMCParticle* mother = static_cast<AliAODMCParticle*>(mcParticles->UncheckedAt(label));
	      if (!mother)
		{
		  AliError(Form("Got a null mother ! Check that ! (label %d",label)); // FIXME: I think this error is not needed
		  label = -1;
		}
	      else
		{
		  label = mother->GetMother();// do not only keep particles which created a track, but all their mothers
		}
	    }
	}
    
      CreateLabelMap(source);
    
      // Actual filtering and label remapping (shamelessly taken for the implementation of AliAODHandler::StoreMCParticles)
      TIter nextMC(mcParticles);
      AliAODMCParticle* p;
      Int_t nmc(0);  // We need nmc, we cannot rely on part->GetLabel, because some of the original mc particles are not kept in the stack, apparently
      Int_t nmcout(0);
    
      while ( ( p = static_cast<AliAODMCParticle*>(nextMC()) ) )
	{
	  AliAODMCParticle c(*p);

	  if ( IsParticleSelected(nmc) )
	    {
	      // 
	      Int_t d0 =  p->GetDaughter(0);
	      Int_t d1 =  p->GetDaughter(1);
	      Int_t m =   p->GetMother();
        
	      // other than for the track labels, negative values mean
	      // no daughter/mother so preserve it
        
	      if(d0<0 && d1<0)
		{
		  // no first daughter -> no second daughter
		  // nothing to be done
		  // second condition not needed just for sanity check at the end
		  c.SetDaughter(0,d0);
		  c.SetDaughter(1,d1);
		} else if(d1 < 0 && d0 >= 0) 
		{
		  // Only one daughter
		  // second condition not needed just for sanity check at the end
		  if(IsParticleSelected(d0))
		    {
		      c.SetDaughter(0,GetNewLabel(d0));
		    } else 
		    {
		      c.SetDaughter(0,-1);
		    }
		  c.SetDaughter(1,d1);
		}
	      else if (d0 > 0 && d1 > 0 )
		{
		  // we have two or more daughters loop on the stack to see if they are
		  // selected
		  Int_t d0tmp = -1;
		  Int_t d1tmp = -1;
		  for (int id = d0; id<=d1;++id)
		    {
		      if (IsParticleSelected(id))
			{
			  if(d0tmp==-1)
			    {
			      // first time
			      d0tmp = GetNewLabel(id);
			      d1tmp = d0tmp; // this is to have the same schema as on the stack i.e. with one daugther d0 and d1 are the same 
			    }
			  else d1tmp = GetNewLabel(id);
			}
		    }
		  c.SetDaughter(0,d0tmp);
		  c.SetDaughter(1,d1tmp);
		} else 
		{
		  AliFatal(Form("Unxpected indices %d %d",d0,d1));
		}
        
	      if ( m < 0 )
		{
		  c.SetMother(m);
		} else 
		{
		  if (IsParticleSelected(m)) 
		    {
		      c.SetMother(GetNewLabel(m));              
		    }
		  // else // FIXME: re-enable this checj. Sometimes it gets here. Still to be understood why
		  //   {
		  //     //		      AliError(Form("PROBLEM Mother not selected %d", m));              
		  //   }
		}
        
	      new ((*fMCParticles)[nmcout++]) AliAODMCParticle(c);
	    }
      
	  ++nmc;        
	} //closes loop over MC particles
    
      // now remap the tracks...
    
      TIter nextTrack(fTracks);
      AliNanoAODTrack* t;
      //      std::cout << "Remapping tracks" << std::endl;
    
      while ( ( t = dynamic_cast<AliNanoAODTrack*>(nextTrack()) ) )
	{
	  
	  t->SetLabel(GetNewLabel(t->GetLabel()));
	}
    
    } // closes fMCMode == 1
  else if ( mcParticles ) 
    {
      // simple copy of input MC particles to ouput MC particles
      TIter nextMC(mcParticles);
      AliAODMCParticle* p;
      Int_t nmcout(0);
    
      while ( ( p = static_cast<AliAODMCParticle*>(nextMC()) ) )
	{
	  new ((*fMCParticles)[nmcout++]) AliAODMCParticle(*p);
	}
    }
  
  AliDebug(1,Form("input mc %d output mc %d",
                  mcParticles ? mcParticles->GetEntries() : 0,
                  fMCParticles ? fMCParticles->GetEntries() : 0));
  
  Printf("input mc %d output mc %d",
                  mcParticles ? mcParticles->GetEntries() : 0,
                  fMCParticles ? fMCParticles->GetEntries() : 0);
  

}

// //_____________________________________________________________________________
TList* AliNanoAODReplicator::GetList() const
{
  // return (and build if not already done) our internal list of managed objects
  
  if (!fList)
    {
      fList = new TList;
      fList->SetOwner(kTRUE);

      fTracks = new TClonesArray("AliNanoAODTrack");      
      fTracks->SetName("tracks"); // TODO: consider the possibility to use a different name to distinguish in AliAODEvent
      fList->Add(fTracks);    

      fHeader = new AliNanoAODHeader(3);// TODO: to be customized
      fHeader->SetName("header"); // TODO: consider the possibility to use a different name to distinguish in AliAODEvent
      fList->Add(fHeader);    


      fVertices = new TClonesArray("AliAODVertex",2);
      fVertices->SetName("vertices");    
    
        
      fList->Add(fVertices);
    
      if ( fMCMode > 0 )
	{
	  fMCHeader = new AliAODMCHeader;    
	  fMCParticles = new TClonesArray("AliAODMCParticle",1000);
	  fMCParticles->SetName(AliAODMCParticle::StdBranchName());
	  fList->Add(fMCHeader);
	  fList->Add(fMCParticles);
	}
    }
  return fList;
}

//_____________________________________________________________________________
void AliNanoAODReplicator::ReplicateAndFilter(const AliAODEvent& source)
//void AliNanoAODReplicator::ReplicateAndFilter(AliAODEvent *source)
{
  // Replicate (and filter if filters are there) the relevant parts we're interested in AODEvent
  

  // assert(fTracks!=0x0);
  
  //*fTZERO = *(source.GetTZEROData());
  
  
  fTracks->Clear("C");			
  assert(fVertices!=0x0);
  fVertices->Clear("C");
  if (fMCMode > 0){
    if(!fMCHeader) {
      AliFatal(Form("fMCMode = %d, but MC header not found", fMCMode));
    }
    fMCHeader->Reset();
    if(!fMCParticles){
      AliFatal(Form("fMCMode = %d, but MC particles not found", fMCMode));
    }
    fMCParticles->Clear("C");
  }
  Int_t ntracks(0);
  Int_t input(0);

  AliAODVertex *vtx = source.GetPrimaryVertex();

  // TODO: implement header!
  //  *fHeader = *source.GetHeader();
  if(fCustomSetter){
    // Set custom variables in the header if the callback is set
    fCustomSetter->SetNanoAODHeader(&source, fHeader);
  }


  //  Int_t nindices=0;
  const Int_t entries = source.GetNumberOfTracks();

  Double_t pos[3],cov[6];
  vtx->GetXYZ(pos);
  vtx->GetCovarianceMatrix(cov);
  
  if(entries<=0) return;

  if(vtx->GetNContributors()<1) {
    // SPD vertex cut
    vtx =source.GetPrimaryVertexSPD();    
    if(vtx->GetNContributors()<1) { // FIXME: Keep this cut?
      Info("AliNanoAODReplicator","No good vertex, skip event");
      return; // NO GOOD VERTEX, SKIP EVENT 
    }
  }
  
  
  for(Int_t j=0; j<entries; j++){
    
    AliVTrack *track = (AliVTrack*)source.GetTrack(j);
    
    AliAODTrack *aodtrack =(AliAODTrack*)track;// FIXME DYNAMIC CAST?
    if(fTrackCut && !fTrackCut->IsSelected(aodtrack)) continue;

    AliNanoAODTrack * special = new((*fTracks)[ntracks++]) AliNanoAODTrack (aodtrack, fVarList);
    
    if(fCustomSetter) fCustomSetter->SetNanoAODTrack(aodtrack, special);
  }  
  //----------------------------------------------------------
  
  TIter nextV(source.GetVertices());
  AliAODVertex* v;
  Int_t nvertices(0);
  
  while ( ( v = static_cast<AliAODVertex*>(nextV()) ) )
    {
      AliAODVertex* tmp = v->CloneWithoutRefs();
      AliAODVertex* copiedVertex = new((*fVertices)[nvertices++]) AliAODVertex(*tmp);
      
      copiedVertex->SetNContributors(v->GetNContributors()); 
      
      delete tmp;
    }
  
  
  AliDebug(1,Form("input mu tracks=%d tracks=%d vertices=%d",
                  input,fTracks->GetEntries(),fVertices->GetEntries())); 
  
  
  // Finally, deal with MC information, if needed
  
  if ( fMCMode > 0 ) {
    FilterMC(source);      
  }
  

}


//-----------------------------------------------------------------------------

//----------------------------------------------------------------------------
	

//-----------------------------------------------------------------------------

// AliAODVertex* AliNanoAODReplicator::PrimaryVertex(const TObjArray *trkArray,
// 						   AliAODEvent &event) const
// {
//   // Returns primary vertex to be used for this candidate
//   //AliCodeTimerAuto("",0);

//   AliESDVertex *vertexESD = 0;
//   AliAODVertex *vertexAOD = 0;


//   if(!fRecoPrimVtxSkippingTrks && !fRmTrksFromPrimVtx) { 
//     // primary vertex from the input event
    
//     vertexESD = new AliESDVertex(*fV1);

//   } else {
//     // primary vertex specific to this candidate

//     Int_t nTrks = trkArray->GetEntriesFast();
//     AliVertexerTracks *vertexer = new AliVertexerTracks(event.GetMagneticField());

//     if(fRecoPrimVtxSkippingTrks) { 
//       // recalculating the vertex
      
//       if(strstr(fV1->GetTitle(),"VertexerTracksWithConstraint")) {
// 	Float_t diamondcovxy[3];
// 	event.GetDiamondCovXY(diamondcovxy);
// 	Double_t pos[3]={event.GetDiamondX(),event.GetDiamondY(),0.};
// 	Double_t cov[6]={diamondcovxy[0],diamondcovxy[1],diamondcovxy[2],0.,0.,10.*10.};
// 	AliESDVertex *diamond = new AliESDVertex(pos,cov,1.,1);
// 	vertexer->SetVtxStart(diamond);
// 	delete diamond; diamond=NULL;
// 	if(strstr(fV1->GetTitle(),"VertexerTracksWithConstraintOnlyFitter")) 
// 	  vertexer->SetOnlyFitter();
//       }
//       Int_t skipped[1000];
//       Int_t nTrksToSkip=0,id;
//       AliExternalTrackParam *t = 0;
//       for(Int_t i=0; i<nTrks; i++) {
// 	t = (AliExternalTrackParam*)trkArray->UncheckedAt(i);
// 	id = (Int_t)t->GetID();
// 	if(id<0) continue;
// 	skipped[nTrksToSkip++] = id;
//       }
//       // TEMPORARY FIX
//       // For AOD, skip also tracks without covariance matrix
//       if(fInputAOD) {
// 	Double_t covtest[21];
// 	for(Int_t j=0; j<event.GetNumberOfTracks(); j++) {
// 	  AliVTrack *vtrack = (AliVTrack*)event.GetTrack(j);
// 	  if(!vtrack->GetCovarianceXYZPxPyPz(covtest)) {
// 	    id = (Int_t)vtrack->GetID();
// 	    if(id<0) continue;
// 	    skipped[nTrksToSkip++] = id;
// 	  }
// 	}
//       }
//       for(Int_t ijk=nTrksToSkip; ijk<1000; ijk++) skipped[ijk]=-1;
//       //
//       vertexer->SetSkipTracks(nTrksToSkip,skipped);
//       vertexESD = (AliESDVertex*)vertexer->FindPrimaryVertex(event); 
      
//     } else if(fRmTrksFromPrimVtx && nTrks>0) { 
//       // removing the prongs tracks
      
//       TObjArray rmArray(nTrks);
//       UShort_t *rmId = new UShort_t[nTrks];
//       AliESDtrack *esdTrack = 0;
//       AliESDtrack *t = 0;
//       for(Int_t i=0; i<nTrks; i++) {
// 	t = (AliESDtrack*)trkArray->UncheckedAt(i);
// 	esdTrack = new AliESDtrack(*t);
// 	rmArray.AddLast(esdTrack);
// 	if(esdTrack->GetID()>=0) {
// 	  rmId[i]=(UShort_t)esdTrack->GetID();
// 	} else {
// 	  rmId[i]=9999;
// 	}
//       }
//       Float_t diamondxy[2]={event.GetDiamondX(),event.GetDiamondY()};
//       vertexESD = vertexer->RemoveTracksFromVertex(fV1,&rmArray,rmId,diamondxy);
//       delete [] rmId; rmId=NULL;
//       rmArray.Delete();
      
//     }

//     if(!vertexESD) return vertexAOD;
//     if(vertexESD->GetNContributors()<=0) { 
//       //AliDebug(2,"vertexing failed"); 
//       delete vertexESD; vertexESD=NULL;
//       return vertexAOD;
//     }

//     delete vertexer; vertexer=NULL;

//   }

//   // convert to AliAODVertex
//   Double_t pos[3],cov[6],chi2perNDF;
//   vertexESD->GetXYZ(pos); // position
//   vertexESD->GetCovMatrix(cov); //covariance matrix
//   chi2perNDF = vertexESD->GetChi2toNDF();
//   delete vertexESD; vertexESD=NULL;

//   vertexAOD = new AliAODVertex(pos,cov,chi2perNDF);

//   return vertexAOD;
// }

//_____________________________________________________________________________


// //---------------------------------------------------------------------------

void AliNanoAODReplicator::Terminate(){

}
Commit	Line	Data
5229b088	1	/**************************************************************************
	2	* Copyright(c) 1998-1999, ALICE Experiment at CERN, All rights reserved. *
	3	* *
	4	* Author: The ALICE Off-line Project. *
	5	* Contributors are mentioned in the code where appropriate. *
	6	* *
	7	* Permission to use, copy, modify and distribute this software and its *
	8	* documentation strictly for non-commercial purposes is hereby granted *
	9	* without fee, provided that the above copyright notice appears in all *
	10	* copies and that both the copyright notice and this permission notice *
	11	* appear in the supporting documentation. The authors make no claims *
	12	* about the suitability of this software for any purpose. It is *
	13	* provided "as is" without express or implied warranty. *
	14	**************************************************************************/
	15
	16	// $Id: AliNanoAODReplicator.cxx 56492 2012-05-15 18:42:47Z pcrochet $
	17
	18	// Implementation of a branch replicator
	19	// to produce nanoAODs.
	20	//
	21	//
	22	// This replicator is in charge of replicating the tracks,vertices,headers
	23	// branches of a standard AOD or ESD file into a nanoAODs
	24	// (AliAOD.Special.root)
	25	//
	26	// The class was inspired by AliAODMuonReplicator
	27	//
	28	// Author: Michele Floris, michele.floris@cern.ch
	29
	30
	31	class AliESDv0;
	32	class AliESDVertex;
	33	class AliAODVertex;
	34	class AliAODRecoDecay;
	35
	36	#include "AliAODDimuon.h"
	37	#include "AliAODEvent.h"
	38	#include "AliAODMCHeader.h"
	39	#include "AliAODMCParticle.h"
	40	#include "AliAODTZERO.h"
	41	#include "AliAODTrack.h"
	42	#include "AliAODVZERO.h"
	43	#include "AliAnalysisCuts.h"
	44	#include "TF1.h"
	45	#include "AliExternalTrackParam.h"
	46	#include "AliESDv0.h"
	47	#include "AliAODv0.h"
	48	#include "AliPIDResponse.h"
	49	#include <iostream>
	50	#include <cassert>
	51	#include "AliESDtrack.h"
	52	#include "TObjArray.h"
	53	#include "AliAnalysisFilter.h"
	54	#include "AliNanoAODTrack.h"
	55
	56	#include <TFile.h>
	57	#include <TDatabasePDG.h>
	58	#include <TString.h>
	59	#include <TList.h>
	60	#include "AliLog.h"
	61	#include "AliVEvent.h"
	62	#include "AliVVertex.h"
	63	#include "AliVTrack.h"
	64	#include "AliVertexerTracks.h"
65	#include "AliKFVertex.h"
66	#include "AliESDEvent.h"
67	#include "AliESDVertex.h"
68	#include "AliESDtrackCuts.h"
69	#include "AliAODEvent.h"
70	#include "AliAnalysisFilter.h"
71
72	#include "AliNanoAODReplicator.h"
73	#include "TH1.h"
74	#include "TCanvas.h"
75	#include "AliNanoAODHeader.h"
76	#include "AliNanoAODCustomSetter.h"
77
78	using std::cout;
79	using std::endl;
80
81	ClassImp(AliNanoAODReplicator)
82
83	//_____________________________________________________________________________
84	AliNanoAODReplicator::AliNanoAODReplicator() :
85	AliAODBranchReplicator(),
86	fTrackCut(0), fTracks(0x0), fHeader(0x0), fNTracksVariables(0), // FIXME: Start using cuts, and check if fNTracksVariables is needed
87	fVertices(0x0),
88	fList(0x0),
89	fMCParticles(0x0),
90	fMCHeader(0x0),
91	fMCMode(0),
92	fLabelMap(),
93	fParticleSelected(),
94	fVarList(""),
95	fVarListHeader(""),
96	fCustomSetter(0){
97	// Default ctor. we need it to avoid instantiating a wrong mapping when reading from file
98	}
99
100	AliNanoAODReplicator::AliNanoAODReplicator(const char* name, const char* title,
101	const char * varlist,
102	AliAnalysisCuts* trackCut,
103	Int_t mcMode
104	) :
105	AliAODBranchReplicator(name,title),
106
107	fTrackCut(trackCut), fTracks(0x0), fHeader(0x0), fNTracksVariables(0), // FIXME: Start using cuts, and check if fNTracksVariables is needed
108	fVertices(0x0),
109	fList(0x0),
110	fMCParticles(0x0),
111	fMCHeader(0x0),
112	fMCMode(mcMode),
113	fLabelMap(),
114	fParticleSelected(),
115	fVarList(varlist),
116	fVarListHeader(""),// FIXME: this should be set to a meaningful value: add an arg to the constructor
117	fCustomSetter(0)
118	{
119	// default ctor
120	AliNanoAODTrackMapping * tm =new AliNanoAODTrackMapping(fVarList);
121	fNTracksVariables = tm->GetSize();
122	// tm->Print();
123	}
124
125	//_____________________________________________________________________________
126	AliNanoAODReplicator::~AliNanoAODReplicator()
127	{
128	// dtor
129	delete fTrackCut;
130	delete fList;
131	}
132
133	//_____________________________________________________________________________
134	void AliNanoAODReplicator::SelectParticle(Int_t i)
135	{
136	// taking the absolute values here, need to take care
137	// of negative daughter and mother
138	// IDs when setting!
139
140	if (!IsParticleSelected(TMath::Abs(i)))
141	{
142	fParticleSelected.Add(TMath::Abs(i),1);
143	}
144	}
145
146	//_____________________________________________________________________________
147	Bool_t AliNanoAODReplicator::IsParticleSelected(Int_t i)
148	{
149	// taking the absolute values here, need to take
150	// care with negative daughter and mother
151	// IDs when setting!
152	return (fParticleSelected.GetValue(TMath::Abs(i))==1);
153	}
154
155
156	//_____________________________________________________________________________
157	void AliNanoAODReplicator::CreateLabelMap(const AliAODEvent& source)
158	{
159	//
160	// this should be called once all selections are done
161	// This method associates to the original index of the mc particle
162	// (i) the new one (j). J runs only over particles which are
163	// actually kept.
164	//
165
166	fLabelMap.Delete();
167
168	TClonesArray* mcParticles = static_cast<TClonesArray*>(source.FindListObject(AliAODMCParticle::StdBranchName()));
169
170	Int_t j(0);
171	Int_t i(0); // We need i, we cannot rely on part->GetLabel, because some of the original mc particles are not kept in the stack, apparently
172
173	TIter next(mcParticles);
174
175	while ( next() )
176	{
177	if (IsParticleSelected(i))
178	{
179	fLabelMap.Add(i,j++);
180	// std::cout << i << "->" << j-1 << std::endl;
181	}
182	++i;
183	}
184
185
186	}
187
188	//_____________________________________________________________________________
189	Int_t AliNanoAODReplicator::GetNewLabel(Int_t i)
190	{
191	// Gets the label from the new created Map
192	// Call CreatLabelMap before
193	// otherwise only 0 returned
194	return fLabelMap.GetValue(TMath::Abs(i));
195	}
196
197	//_____________________________________________________________________________
198	void AliNanoAODReplicator::FilterMC(const AliAODEvent& source)
199	{
200	// Filter MC information
201
202
203	AliAODMCHeader* mcHeader(0x0);
204	TClonesArray* mcParticles(0x0);
205
206	fParticleSelected.Delete();
207
208	// std::cout << "MC Mode: " << fMCMode << ", Tracks " << fTracks->GetEntries() << std::endl;
209
210	if ( fMCMode>=2 && !fTracks->GetEntries() ) {
211	return;
212	}
213	// for fMCMode==1 we only copy MC information for events where there's at least one muon track
214
215	mcHeader = static_cast<AliAODMCHeader*>(source.FindListObject(AliAODMCHeader::StdBranchName()));
216
217	if ( mcHeader )
218	{
219	fMCHeader = mcHeader;
220	}
221
222
223	mcParticles = static_cast<TClonesArray*>(source.FindListObject(AliAODMCParticle::StdBranchName()));
224
225	if ( mcParticles && fMCMode>=2 )
226	{
227	// keep all primaries
228	TIter nextPart(mcParticles);
229	static Int_t iev = -1; // FIXME: remove this (debug)
230	iev ++;
231	AliAODMCParticle * prim = 0;
232	Int_t iprim = 0; // We need iprim, we cannot rely on part->GetLabel, because some of the original mc particles are not kept in the stack, apparently
233	// also select all charged primaries
234	while ((prim = (AliAODMCParticle*) nextPart())) {
235	if(prim->IsPhysicalPrimary() && prim->Charge()) SelectParticle(iprim);
236	// FIXME DEBUG
237	if(iev == 2009) {
238	// std::cout << "IEV " << iev << std::endl;
239	// std::cout << " PART " << iprim << " " << prim->IsPhysicalPrimary() <<","<<prim->Charge() << "=" << IsParticleSelected(iprim) << std::endl;
240	if(iprim == 15) {
241	prim->Print();
242	}
243
244	}
245	iprim++;
246	}
247
248	// loop on (kept) tracks to find their ancestors
249	TIter nextTRACK(fTracks);
250	AliNanoAODTrack* track;
251
252	while ( ( track = static_cast<AliNanoAODTrack*>(nextTRACK()) ) )
253	{
254	Int_t label = TMath::Abs(track->GetLabel());
255
256	while ( label >= 0 )
257	{
258	SelectParticle(label);
259	AliAODMCParticle* mother = static_cast<AliAODMCParticle*>(mcParticles->UncheckedAt(label));
260	if (!mother)
261	{
262	AliError(Form("Got a null mother ! Check that ! (label %d",label)); // FIXME: I think this error is not needed
263	label = -1;
264	}
265	else
266	{
267	label = mother->GetMother();// do not only keep particles which created a track, but all their mothers
268	}
269	}
270	}
271
272	CreateLabelMap(source);
273
274	// Actual filtering and label remapping (shamelessly taken for the implementation of AliAODHandler::StoreMCParticles)
275	TIter nextMC(mcParticles);
276	AliAODMCParticle* p;
277	Int_t nmc(0); // We need nmc, we cannot rely on part->GetLabel, because some of the original mc particles are not kept in the stack, apparently
278	Int_t nmcout(0);
279
280	while ( ( p = static_cast<AliAODMCParticle*>(nextMC()) ) )
281	{
282	AliAODMCParticle c(*p);
283
284	if ( IsParticleSelected(nmc) )
285	{
286	//
287	Int_t d0 = p->GetDaughter(0);
288	Int_t d1 = p->GetDaughter(1);
289	Int_t m = p->GetMother();
290
291	// other than for the track labels, negative values mean
292	// no daughter/mother so preserve it
293
294	if(d0<0 && d1<0)
295	{
296	// no first daughter -> no second daughter
297	// nothing to be done
298	// second condition not needed just for sanity check at the end
299	c.SetDaughter(0,d0);
300	c.SetDaughter(1,d1);
301	} else if(d1 < 0 && d0 >= 0)
302	{
303	// Only one daughter
304	// second condition not needed just for sanity check at the end
305	if(IsParticleSelected(d0))
306	{
307	c.SetDaughter(0,GetNewLabel(d0));
308	} else
309	{
310	c.SetDaughter(0,-1);
311	}
312	c.SetDaughter(1,d1);
313	}
314	else if (d0 > 0 && d1 > 0 )
315	{
316	// we have two or more daughters loop on the stack to see if they are
317	// selected
318	Int_t d0tmp = -1;
319	Int_t d1tmp = -1;
320	for (int id = d0; id<=d1;++id)
321	{
322	if (IsParticleSelected(id))
323	{
324	if(d0tmp==-1)
325	{
326	// first time
327	d0tmp = GetNewLabel(id);
328	d1tmp = d0tmp; // this is to have the same schema as on the stack i.e. with one daugther d0 and d1 are the same
329	}
330	else d1tmp = GetNewLabel(id);
331	}
332	}
333	c.SetDaughter(0,d0tmp);
334	c.SetDaughter(1,d1tmp);
335	} else
336	{
337	AliFatal(Form("Unxpected indices %d %d",d0,d1));
338	}
339
340	if ( m < 0 )
341	{
342	c.SetMother(m);
343	} else
344	{
345	if (IsParticleSelected(m))
346	{
347	c.SetMother(GetNewLabel(m));
348	}
349	// else // FIXME: re-enable this checj. Sometimes it gets here. Still to be understood why
350	// {
351	// // AliError(Form("PROBLEM Mother not selected %d", m));
352	// }
353	}
354
355	new ((*fMCParticles)[nmcout++]) AliAODMCParticle(c);
356	}
357
358	++nmc;
359	} //closes loop over MC particles
360
361	// now remap the tracks...
362
363	TIter nextTrack(fTracks);
364	AliNanoAODTrack* t;
365	// std::cout << "Remapping tracks" << std::endl;
366
367	while ( ( t = dynamic_cast<AliNanoAODTrack*>(nextTrack()) ) )
368	{
369
370	t->SetLabel(GetNewLabel(t->GetLabel()));
371	}
372
373	} // closes fMCMode == 1
374	else if ( mcParticles )
375	{
376	// simple copy of input MC particles to ouput MC particles
377	TIter nextMC(mcParticles);
378	AliAODMCParticle* p;
379	Int_t nmcout(0);
380
381	while ( ( p = static_cast<AliAODMCParticle*>(nextMC()) ) )
382	{
383	new ((fMCParticles)[nmcout++]) AliAODMCParticle(p);
384	}
385	}
386
387	AliDebug(1,Form("input mc %d output mc %d",
388	mcParticles ? mcParticles->GetEntries() : 0,
389	fMCParticles ? fMCParticles->GetEntries() : 0));
390
391	Printf("input mc %d output mc %d",
392	mcParticles ? mcParticles->GetEntries() : 0,
393	fMCParticles ? fMCParticles->GetEntries() : 0);
394
395
396	}
397
398	// //_____________________________________________________________________________
399	TList* AliNanoAODReplicator::GetList() const
400	{
401	// return (and build if not already done) our internal list of managed objects
402
403	if (!fList)
404	{
405	fList = new TList;
406	fList->SetOwner(kTRUE);
407
408	fTracks = new TClonesArray("AliNanoAODTrack");
409	fTracks->SetName("tracks"); // TODO: consider the possibility to use a different name to distinguish in AliAODEvent
410	fList->Add(fTracks);
411
201dae36	412	fHeader = new AliNanoAODHeader(3);// TODO: to be customized
5229b088	413	fHeader->SetName("header"); // TODO: consider the possibility to use a different name to distinguish in AliAODEvent
	414	fList->Add(fHeader);
	415
	416
	417	fVertices = new TClonesArray("AliAODVertex",2);
	418	fVertices->SetName("vertices");
	419
	420
	421	fList->Add(fVertices);
	422
	423	if ( fMCMode > 0 )
	424	{
	425	fMCHeader = new AliAODMCHeader;
	426	fMCParticles = new TClonesArray("AliAODMCParticle",1000);
	427	fMCParticles->SetName(AliAODMCParticle::StdBranchName());
	428	fList->Add(fMCHeader);
	429	fList->Add(fMCParticles);
	430	}
	431	}
	432	return fList;
	433	}
	434
	435	//_____________________________________________________________________________
	436	void AliNanoAODReplicator::ReplicateAndFilter(const AliAODEvent& source)
	437	//void AliNanoAODReplicator::ReplicateAndFilter(AliAODEvent *source)
	438	{
	439	// Replicate (and filter if filters are there) the relevant parts we're interested in AODEvent
	440
	441
	442	// assert(fTracks!=0x0);
	443
	444	//fTZERO = (source.GetTZEROData());
	445
	446
	447
	448	fTracks->Clear("C");
	449	assert(fVertices!=0x0);
	450	fVertices->Clear("C");
	451	if (fMCMode > 0){
	452	if(!fMCHeader) {
	453	AliFatal(Form("fMCMode = %d, but MC header not found", fMCMode));
	454	}
	455	fMCHeader->Reset();
	456	if(!fMCParticles){
	457	AliFatal(Form("fMCMode = %d, but MC particles not found", fMCMode));
	458	}
	459	fMCParticles->Clear("C");
	460	}
	461	Int_t ntracks(0);
	462	Int_t input(0);
	463
	464	AliAODVertex *vtx = source.GetPrimaryVertex();
	465
	466	// TODO: implement header!
	467	// fHeader = source.GetHeader();
	468	if(fCustomSetter){
	469	// Set custom variables in the header if the callback is set
	470	fCustomSetter->SetNanoAODHeader(&source, fHeader);
	471	}
	472
	473
	474	// Int_t nindices=0;
	475	const Int_t entries = source.GetNumberOfTracks();
	476
477	Double_t pos[3],cov[6];
478	vtx->GetXYZ(pos);
479	vtx->GetCovarianceMatrix(cov);
480
481	if(entries<=0) return;
482
483	if(vtx->GetNContributors()<1) {
484	// SPD vertex cut
485	vtx =source.GetPrimaryVertexSPD();
486	if(vtx->GetNContributors()<1) { // FIXME: Keep this cut?
487	Info("AliNanoAODReplicator","No good vertex, skip event");
488	return; // NO GOOD VERTEX, SKIP EVENT
489	}
490	}
491
492
493	for(Int_t j=0; j<entries; j++){
494
495	AliVTrack track = (AliVTrack)source.GetTrack(j);
496
497	AliAODTrack aodtrack =(AliAODTrack)track;// FIXME DYNAMIC CAST?
105e7451	498	if(fTrackCut && !fTrackCut->IsSelected(aodtrack)) continue;
5229b088	499
0df1f11b	500	AliNanoAODTrack * special = new((*fTracks)[ntracks++]) AliNanoAODTrack (aodtrack, fVarList);
5229b088	501
5229b088	502	if(fCustomSetter) fCustomSetter->SetNanoAODTrack(aodtrack, special);
5229b088	503	}
	504	//----------------------------------------------------------
	505
	506	TIter nextV(source.GetVertices());
	507	AliAODVertex* v;
	508	Int_t nvertices(0);
	509
	510	while ( ( v = static_cast<AliAODVertex*>(nextV()) ) )
	511	{
	512	AliAODVertex* tmp = v->CloneWithoutRefs();
	513	AliAODVertex* copiedVertex = new((fVertices)[nvertices++]) AliAODVertex(tmp);
	514
5229b088	515	copiedVertex->SetNContributors(v->GetNContributors());
5229b088	516
5229b088	517	delete tmp;
	518	}
	519
	520
	521	AliDebug(1,Form("input mu tracks=%d tracks=%d vertices=%d",
	522	input,fTracks->GetEntries(),fVertices->GetEntries()));
	523
	524
	525	// Finally, deal with MC information, if needed
	526
	527	if ( fMCMode > 0 ) {
	528	FilterMC(source);
	529	}
	530
	531
	532	}
	533
	534
	535
	536	//-----------------------------------------------------------------------------
	537
	538	//----------------------------------------------------------------------------
	539
	540
	541	//-----------------------------------------------------------------------------
	542
	543	// AliAODVertex* AliNanoAODReplicator::PrimaryVertex(const TObjArray *trkArray,
	544	// AliAODEvent &event) const
	545	// {
	546	// // Returns primary vertex to be used for this candidate
	547	// //AliCodeTimerAuto("",0);
	548
	549	// AliESDVertex *vertexESD = 0;
	550	// AliAODVertex *vertexAOD = 0;
	551
	552
	553	// if(!fRecoPrimVtxSkippingTrks && !fRmTrksFromPrimVtx) {
	554	// // primary vertex from the input event
	555
	556	// vertexESD = new AliESDVertex(*fV1);
	557
	558	// } else {
	559	// // primary vertex specific to this candidate
	560
	561	// Int_t nTrks = trkArray->GetEntriesFast();
	562	// AliVertexerTracks *vertexer = new AliVertexerTracks(event.GetMagneticField());
	563
	564	// if(fRecoPrimVtxSkippingTrks) {
	565	// // recalculating the vertex
	566
	567	// if(strstr(fV1->GetTitle(),"VertexerTracksWithConstraint")) {
	568	// Float_t diamondcovxy[3];
	569	// event.GetDiamondCovXY(diamondcovxy);
	570	// Double_t pos[3]={event.GetDiamondX(),event.GetDiamondY(),0.};
	571	// Double_t cov[6]={diamondcovxy[0],diamondcovxy[1],diamondcovxy[2],0.,0.,10.*10.};
	572	// AliESDVertex *diamond = new AliESDVertex(pos,cov,1.,1);
	573	// vertexer->SetVtxStart(diamond);
	574	// delete diamond; diamond=NULL;
	575	// if(strstr(fV1->GetTitle(),"VertexerTracksWithConstraintOnlyFitter"))
	576	// vertexer->SetOnlyFitter();
	577	// }
	578	// Int_t skipped[1000];
	579	// Int_t nTrksToSkip=0,id;
	580	// AliExternalTrackParam *t = 0;
581	// for(Int_t i=0; i<nTrks; i++) {
582	// t = (AliExternalTrackParam*)trkArray->UncheckedAt(i);
583	// id = (Int_t)t->GetID();
584	// if(id<0) continue;
585	// skipped[nTrksToSkip++] = id;
586	// }
587	// // TEMPORARY FIX
588	// // For AOD, skip also tracks without covariance matrix
589	// if(fInputAOD) {
590	// Double_t covtest[21];
591	// for(Int_t j=0; j<event.GetNumberOfTracks(); j++) {
592	// AliVTrack vtrack = (AliVTrack)event.GetTrack(j);
593	// if(!vtrack->GetCovarianceXYZPxPyPz(covtest)) {
594	// id = (Int_t)vtrack->GetID();
595	// if(id<0) continue;
596	// skipped[nTrksToSkip++] = id;
597	// }
598	// }
599	// }
600	// for(Int_t ijk=nTrksToSkip; ijk<1000; ijk++) skipped[ijk]=-1;
601	// //
602	// vertexer->SetSkipTracks(nTrksToSkip,skipped);
603	// vertexESD = (AliESDVertex*)vertexer->FindPrimaryVertex(event);
604
605	// } else if(fRmTrksFromPrimVtx && nTrks>0) {
606	// // removing the prongs tracks
607
608	// TObjArray rmArray(nTrks);
609	// UShort_t *rmId = new UShort_t[nTrks];
610	// AliESDtrack *esdTrack = 0;
611	// AliESDtrack *t = 0;
612	// for(Int_t i=0; i<nTrks; i++) {
613	// t = (AliESDtrack*)trkArray->UncheckedAt(i);
614	// esdTrack = new AliESDtrack(*t);
615	// rmArray.AddLast(esdTrack);
616	// if(esdTrack->GetID()>=0) {
617	// rmId[i]=(UShort_t)esdTrack->GetID();
618	// } else {
619	// rmId[i]=9999;
620	// }
621	// }
622	// Float_t diamondxy[2]={event.GetDiamondX(),event.GetDiamondY()};
623	// vertexESD = vertexer->RemoveTracksFromVertex(fV1,&rmArray,rmId,diamondxy);
624	// delete [] rmId; rmId=NULL;
625	// rmArray.Delete();
626
627	// }
628
629	// if(!vertexESD) return vertexAOD;
630	// if(vertexESD->GetNContributors()<=0) {
631	// //AliDebug(2,"vertexing failed");
632	// delete vertexESD; vertexESD=NULL;
633	// return vertexAOD;
634	// }
635
636	// delete vertexer; vertexer=NULL;
637
638	// }
639
640	// // convert to AliAODVertex
641	// Double_t pos[3],cov[6],chi2perNDF;
642	// vertexESD->GetXYZ(pos); // position
643	// vertexESD->GetCovMatrix(cov); //covariance matrix
644	// chi2perNDF = vertexESD->GetChi2toNDF();
645	// delete vertexESD; vertexESD=NULL;
646
647	// vertexAOD = new AliAODVertex(pos,cov,chi2perNDF);
648
649	// return vertexAOD;
650	// }
651
652	//_____________________________________________________________________________
653
654
655
656	// //---------------------------------------------------------------------------
657
658	void AliNanoAODReplicator::Terminate(){
659
660	}