update of scripts+macros for non-prompt J/psi simulations

ffionda · gconesab · commit acc8c670b741 · 2021-09-03T10:36:19.000+02:00
diff --git a/MC/config/PWGDQ/external/generator/GeneratorBeautyToJpsi_EvtGen.C b/MC/config/PWGDQ/external/generator/GeneratorBeautyToJpsi_EvtGen.C
@@ -1,7 +1,7 @@
 // usage (fwdy) :
-//o2-sim -j 4 -n 10 -g external -m "PIPE ITS TPC" -o sgn --configKeyValues "GeneratorExternal.fileName=GeneratorBeautyToJpsi_EvtGen.C;GeneratorExternal.funcName=GeneratorBeautyToJpsi_EvtGenFwdY()" --configFile GeneratorHF_bbbar_fwdy.ini 
+//o2-sim -j 4 -n 10 -g external -t external -m "PIPE ITS TPC" -o sgn --configFile GeneratorHF_bbbar_fwdy.ini 
 // usage (midy) :
-//o2-sim -j 4 -n 10 -g external -m "PIPE ITS TPC" -o sgn --configKeyValues "GeneratorExternal.fileName=GeneratorBeautyToJpsi_EvtGen.C;GeneratorExternal.funcName=GeneratorBeautyToJpsi_EvtGenMidY()" --configFile GeneratorHF_bbbar_midy.ini 
+//o2-sim -j 4 -n 10 -g external -t external -m "PIPE ITS TPC" -o sgn --configFile GeneratorHF_bbbar_midy.ini 
 //
 //
 R__ADD_INCLUDE_PATH($O2DPG_ROOT/MC/config/PWGDQ/EvtGen)
@@ -10,14 +10,15 @@ R__ADD_INCLUDE_PATH($O2DPG_ROOT/MC/config/PWGHF/external/generator)
 #include "GeneratorHF.C"
 
 FairGenerator*
-GeneratorBeautyToJpsi_EvtGenMidY(double rapidityMin = -1.5, double rapidityMax = 1.5, bool verbose = false, TString pdgs = "511;521;531;5112;5122;5232;5132")
+GeneratorBeautyToJpsi_EvtGenMidY(double rapidityMin = -1.5, double rapidityMax = 1.5, bool ispp = true, bool verbose = false, TString pdgs = "511;521;531;5112;5122;5232;5132")
 {
   auto gen = new o2::eventgen::GeneratorEvtGen<o2::eventgen::GeneratorHF>(); 
   gen->setRapidity(rapidityMin,rapidityMax);
   gen->setPDG(5);
 
   gen->setVerbose(verbose);
-  gen->setFormula("max(1.,120.*(x<5.)+80.*(1.-x/20.)*(x>5.)*(x<11.)+240.*(1.-x/13.)*(x>11.))");
+  if(ispp) gen->setFormula("1");
+  else gen->setFormula("max(1.,120.*(x<5.)+80.*(1.-x/20.)*(x>5.)*(x<11.)+240.*(1.-x/13.)*(x>11.))");
   std::string spdg;
   TObjArray *obj = pdgs.Tokenize(";");
   gen->SetSizePdg(obj->GetEntriesFast());
@@ -34,14 +35,15 @@ GeneratorBeautyToJpsi_EvtGenMidY(double rapidityMin = -1.5, double rapidityMax =
 }
 
 FairGenerator*
-GeneratorBeautyToJpsi_EvtGenFwdY(double rapidityMin = -4.3, double rapidityMax = -2.2, bool verbose = false, TString pdgs = "511;521;531;5112;5122;5232;5132")
+GeneratorBeautyToJpsi_EvtGenFwdY(double rapidityMin = -4.3, double rapidityMax = -2.2, bool ispp = true, bool verbose = false, TString pdgs = "511;521;531;5112;5122;5232;5132")
 {
   auto gen = new o2::eventgen::GeneratorEvtGen<o2::eventgen::GeneratorHF>();
   gen->setRapidity(rapidityMin,rapidityMax);
   gen->setPDG(5);
 
   gen->setVerbose(verbose);
-  gen->setFormula("max(1.,120.*(x<5.)+80.*(1.-x/20.)*(x>5.)*(x<11.)+240.*(1.-x/13.)*(x>11.))");
+  if(ispp) gen->setFormula("1");
+  else gen->setFormula("max(1.,120.*(x<5.)+80.*(1.-x/20.)*(x>5.)*(x<11.)+240.*(1.-x/13.)*(x>11.))");
   std::string spdg;
   TObjArray *obj = pdgs.Tokenize(";");
   gen->SetSizePdg(obj->GetEntriesFast());
diff --git a/MC/config/PWGDQ/ini/GeneratorHF_bbbar_fwdy.ini b/MC/config/PWGDQ/ini/GeneratorHF_bbbar_fwdy.ini
@@ -2,8 +2,8 @@
 ### This part sets the path of the file and the function call to retrieve it
 
 [GeneratorExternal]
-fileName = ${O2DPG_ROOT}/MC/config/PWGHF/external/generator/GeneratorHF.C
-funcName = GeneratorHF_bbbar()
+fileName = ${O2DPG_ROOT}/MC/config/PWGDQ/external/generator/GeneratorBeautyToJpsi_EvtGen.C
+funcName = GeneratorBeautyToJpsi_EvtGenFwdY()
 
 ### The external generator derives from GeneratorPythia8.
 ### This part configures the bits of the interface: configuration and user hooks
@@ -13,6 +13,14 @@ config = ${O2DPG_ROOT}/MC/config/common/pythia8/generator/pythia8_hf.cfg
 hooksFileName = ${O2DPG_ROOT}/MC/config/PWGHF/pythia8/hooks/pythia8_userhooks_qqbar.C
 hooksFuncName = pythia8_userhooks_bbbar(-4.3,-2.3)
 
+### The setup uses an external even generator trigger which is
+### defined in the following file and it is retrieved and configured
+### according to the specified function call
+
+[TriggerExternal]
+fileName = ${O2DPG_ROOT}/MC/config/PWGDQ/trigger/selectDaughterFromHFwithinAcc.C
+funcName = selectDaughterFromHFwithinAcc(443,kTRUE,-4.3,-2.3)
+
 ### The setup inhibits transport of primary particles which are produce at forward rapidity.
 ### The settings below only transports particles in the barrel, which is currently defined as |eta| < 2
 
diff --git a/MC/config/PWGDQ/ini/GeneratorHF_bbbar_midy.ini b/MC/config/PWGDQ/ini/GeneratorHF_bbbar_midy.ini
@@ -2,8 +2,8 @@
 ### This part sets the path of the file and the function call to retrieve it
 
 [GeneratorExternal]
-fileName = ${O2DPG_ROOT}/MC/config/PWGHF/external/generator/GeneratorHF.C
-funcName = GeneratorHF_bbbar()
+fileName = ${O2DPG_ROOT}/MC/config/PWGDQ/external/generator/GeneratorBeautyToJpsi_EvtGen.C
+funcName = GeneratorBeautyToJpsi_EvtGenMidY()
 
 ### The external generator derives from GeneratorPythia8.
 ### This part configures the bits of the interface: configuration and user hooks
@@ -13,6 +13,14 @@ config = ${O2DPG_ROOT}/MC/config/common/pythia8/generator/pythia8_hf.cfg
 hooksFileName = ${O2DPG_ROOT}/MC/config/PWGHF/pythia8/hooks/pythia8_userhooks_qqbar.C
 hooksFuncName = pythia8_userhooks_bbbar(-1.5,1.5)
 
+### The setup uses an external even generator trigger which is
+### defined in the following file and it is retrieved and configured
+### according to the specified function call
+
+[TriggerExternal]
+fileName = ${O2DPG_ROOT}/MC/config/PWGDQ/trigger/selectDaughterFromHFwithinAcc.C
+funcName = selectDaughterFromHFwithinAcc(443,kTRUE,-1.5,1.5)
+
 ### The setup inhibits transport of primary particles which are produce at forward rapidity.
 ### The settings below only transports particles in the barrel, which is currently defined as |eta| < 2
 
diff --git a/MC/config/PWGDQ/trigger/selectDaughterFromHFwithinAcc.C b/MC/config/PWGDQ/trigger/selectDaughterFromHFwithinAcc.C
@@ -0,0 +1,76 @@
+R__ADD_INCLUDE_PATH($O2DPG_ROOT)
+#include <TParticle.h>
+#include "Generators/Trigger.h"
+
+/// =================================================================================================================================
+/// Select daughters from HF particles produced in a given rapidity window
+/// pdgPartForAccCut: pdg of the particle (coming from c / b) requested within the rapidity window [rapidityMin, rapidityMax]
+/// cutOnSingleChild: if true the condition on the rapidity is required for only one of the child particles (e.g. bb -> J/psi J/psi, bb -> ee,...)
+/// Tested for: 
+///  	- non-prompt J/psi / Psi(2S) 
+///  	- dielectron / dimuon pairs from cc and bb
+///	- single electrons / muons from b and b -> c -> e 
+/// =================================================================================================================================
+Int_t GetFlavour(Int_t pdgCode);
+
+o2::eventgen::Trigger selectDaughterFromHFwithinAcc(Int_t pdgPartForAccCut=443, Bool_t cutOnSingleChild = kTRUE, double rapidityMin = -1., double rapidityMax = -1.)
+{
+  return [pdgPartForAccCut,cutOnSingleChild,rapidityMin,rapidityMax](const std::vector<TParticle>& particles) -> bool {
+  
+  int nsig = 0; int mpdg = -1; int mpdgUpperFamily = -1; double rapidity = -999.; 
+  Bool_t isSelectedY = kFALSE; if(!cutOnSingleChild) isSelectedY = kTRUE; 
+  Bool_t isHF = kFALSE;
+  for (const auto& particle : particles) {
+	  if(cutOnSingleChild && TMath::Abs(particle.GetPdgCode()) == pdgPartForAccCut){
+	  Int_t mi = particle.GetMother(0); 
+	  if(mi<0) continue;
+          TParticle mother = particles.at(mi);
+          mpdg=TMath::Abs(mother.GetPdgCode());
+          mpdgUpperFamily=(mpdg>1000 ? mpdg+1000 : mpdg+100); 
+          if(GetFlavour(mpdg) == 5 || GetFlavour(mpdgUpperFamily) == 5){ // keep particles from (b->) c 
+	     rapidity = particle.Y();
+	     if( (rapidity > rapidityMin) && (rapidity < rapidityMax) ) isSelectedY = kTRUE; 
+	    }
+	  }
+         ///////
+	 if(!cutOnSingleChild && TMath::Abs(particle.GetPdgCode()) == pdgPartForAccCut){
+          Int_t mi = particle.GetMother(0);
+          if(mi<0) continue;
+          TParticle mother = particles.at(mi);
+          mpdg=TMath::Abs(mother.GetPdgCode());
+          if( (GetFlavour(mpdg) == 5) || (GetFlavour(mpdg) == 4)){
+	     isHF = kTRUE; 
+             rapidity = particle.Y();
+             if( (rapidity < rapidityMin) || (rapidity > rapidityMax) ) isSelectedY = kFALSE;
+            }
+          }
+    }
+    //
+    if(cutOnSingleChild && !isSelectedY) return kFALSE; 
+    if(!cutOnSingleChild && !(isHF && isSelectedY)) return kFALSE; 
+    return kTRUE;  
+  };
+
+}
+
+Int_t GetFlavour(Int_t pdgCode)
+  {
+  //
+  // return the flavour of a particle
+  // input: pdg code of the particle
+  // output: Int_t
+  //         3 in case of strange (open and hidden)
+  //         4 in case of charm (")
+  //         5 in case of beauty (")
+  //
+  Int_t pdg = TMath::Abs(pdgCode);
+  //Resonance
+  if (pdg > 100000) pdg %= 100000;
+  if(pdg > 10000)  pdg %= 10000;
+  // meson ?
+  if(pdg > 10) pdg/=100;
+  // baryon ?
+  if(pdg > 10) pdg/=10;
+  return pdg;
+  }
+
diff --git a/MC/run/PWGDQ/runBeautyToJpsi_fwdy_pp.sh b/MC/run/PWGDQ/runBeautyToJpsi_fwdy_pp.sh
@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+
+# make sure O2DPG + O2 is loaded
+[ ! "${O2DPG_ROOT}" ] && echo "Error: This needs O2DPG loaded" && exit 1
+[ ! "${O2_ROOT}" ] && echo "Error: This needs O2 loaded" && exit 1
+
+
+# ----------- LOAD UTILITY FUNCTIONS --------------------------
+. ${O2_ROOT}/share/scripts/jobutils.sh 
+
+RNDSEED=${RNDSEED:-0}
+NSIGEVENTS=${NSIGEVENTS:-1}
+NBKGEVENTS=${NBKGEVENTS:-1}
+NWORKERS=${NWORKERS:-8} 
+NTIMEFRAMES=${NTIMEFRAMES:-1}
+
+${O2DPG_ROOT}/MC/bin/o2dpg_sim_workflow.py -eCM 900 -gen external -j ${NWORKERS} -ns ${NSIGEVENTS} -tf ${NTIMEFRAMES} -e TGeant4 -mod "--skipModules ZDC" \
+	-trigger "external" -ini $O2DPG_ROOT/MC/config/PWGDQ/ini/GeneratorHF_bbbar_fwdy.ini  \
+        -genBkg pythia8 -procBkg inel -colBkg pp --embedding -nb ${NBKGEVENTS}
+
+# run workflow
+${O2DPG_ROOT}/MC/bin/o2_dpg_workflow_runner.py -f workflow.json
diff --git a/MC/run/PWGDQ/runBeautyToJpsi_midy_pp.sh b/MC/run/PWGDQ/runBeautyToJpsi_midy_pp.sh
@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+
+# make sure O2DPG + O2 is loaded
+[ ! "${O2DPG_ROOT}" ] && echo "Error: This needs O2DPG loaded" && exit 1
+[ ! "${O2_ROOT}" ] && echo "Error: This needs O2 loaded" && exit 1
+
+
+# ----------- LOAD UTILITY FUNCTIONS --------------------------
+. ${O2_ROOT}/share/scripts/jobutils.sh 
+
+RNDSEED=${RNDSEED:-0}
+NSIGEVENTS=${NSIGEVENTS:-1}
+NBKGEVENTS=${NBKGEVENTS:-1}
+NWORKERS=${NWORKERS:-8}
+NTIMEFRAMES=${NTIMEFRAMES:-1}
+
+${O2DPG_ROOT}/MC/bin/o2dpg_sim_workflow.py -eCM 900 -gen external -j ${NWORKERS} -ns ${NSIGEVENTS} -tf ${NTIMEFRAMES} -e TGeant4 -mod "--skipModules ZDC" \
+	-trigger "external" -ini $O2DPG_ROOT/MC/config/PWGDQ/ini/GeneratorHF_bbbar_midy.ini  \
+	-genBkg pythia8 -procBkg inel -colBkg pp --embedding -nb ${NBKGEVENTS} 
+
+# run workflow
+${O2DPG_ROOT}/MC/bin/o2_dpg_workflow_runner.py -f workflow.json
