PrimaryGenerator.cxx

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#include "Generators/PrimaryGenerator.h"
#include <Generators/PrimaryGeneratorParam.h>
#include "Generators/Generator.h"
#include "SimConfig/InteractionDiamondParam.h"
#include "SimulationDataFormat/MCEventHeader.h"
#include "SimulationDataFormat/MCGenProperties.h"
#include "DataFormatsCalibration/MeanVertexObject.h"
#include "DetectorsBase/Stack.h"
#include <fairlogger/Logger.h>
 
#include "FairGenericStack.h"
#include "TFile.h"
#include "TTree.h"
 
#include "TDatabasePDG.h"
#include "TVirtualMC.h"
 
using o2::dataformats::MCEventHeader;
 
namespace o2
{
namespace eventgen
{
 
/*****************************************************************/
 
PrimaryGenerator::~PrimaryGenerator()
{
  LOG(info) << "Destructing PrimaryGenerator";
  if (mEmbedFile && mEmbedFile->IsOpen()) {
    mEmbedFile->Close();
    delete mEmbedFile;
  }
  if (mEmbedEvent) {
    delete mEmbedEvent;
  }
}
 
/*****************************************************************/
 
Bool_t PrimaryGenerator::Init()
{
  LOG(info) << "Initialising primary generator";
 
  // set generator ID and description
  auto& params = PrimaryGeneratorParam::Instance();
  setGeneratorId(params.id);
  setGeneratorDescription(params.description);
 
  if (mEmbedTree) {
    LOG(info) << "Embedding into: " << mEmbedFile->GetName()
              << " (" << mEmbedEntries << " events)";
    return FairPrimaryGenerator::Init();
  }
 
  return FairPrimaryGenerator::Init();
}
 
/*****************************************************************/
 
Bool_t PrimaryGenerator::GenerateEvent(FairGenericStack* pStack)
{
  if (!mEmbedTree || mEmbedIndex < 0) {
    fixInteractionVertex(); // <-- always fixes vertex outside of FairROOT
    auto ret = FairPrimaryGenerator::GenerateEvent(pStack);
    if (ret) {
      setGeneratorInformation();
    }
    return ret;
  }
 
  if (mEmbedIndex >= 0) {
    mEmbedTree->GetEntry(mEmbedIndex);
    setInteractionVertex(mEmbedEvent);
 
    auto genList = GetListOfGenerators();
    for (int igen = 0; igen < genList->GetEntries(); ++igen) {
      auto o2gen = dynamic_cast<Generator*>(genList->At(igen));
      if (o2gen) {
        o2gen->notifyEmbedding(mEmbedEvent);
      }
    }
  }
 
  if (!FairPrimaryGenerator::GenerateEvent(pStack)) {
    return kFALSE;
  }
 
  auto o2event = dynamic_cast<MCEventHeader*>(fEvent);
  if (o2event) {
    o2event->setEmbeddingFileName(mEmbedFile->GetName());
    o2event->setEmbeddingEventIndex(mEmbedIndex);
  }
  setGeneratorInformation();
 
  mEmbedIndex++;
  mEmbedIndex %= mEmbedEntries;
 
  return kTRUE;
}
 
/*****************************************************************/
 
void PrimaryGenerator::AddTrack(Int_t pdgid, Double_t px, Double_t py, Double_t pz,
                                Double_t vx, Double_t vy, Double_t vz,
                                Int_t mother1, Int_t mother2,
                                Int_t daughter1, Int_t daughter2,
                                Bool_t wanttracking,
                                Double_t e, Double_t tof,
                                Double_t weight, TMCProcess proc, Int_t generatorStatus)
{
  // check the status encoding
  if (!mcgenstatus::isEncoded(generatorStatus) && proc == TMCProcess::kPPrimary) {
    LOG(fatal) << "Generator status " << generatorStatus << " of particle is not encoded properly.";
  }
 
  vx += fVertex.X();
  vy += fVertex.Y();
  vz += fVertex.Z();
 
  auto particlePDG = TDatabasePDG::Instance()->GetParticle(pdgid);
  if (wanttracking && !particlePDG) {
    LOG(warn) << "Particle to be tracked is not defined in PDG: pdg = " << pdgid << " (disabling tracking)";
    wanttracking = false;
  }
 
  Int_t doTracking = 0; // Go to tracking
  if (fdoTracking && wanttracking) {
    doTracking = 1;
  }
  Int_t dummyparent = -1; // Primary particle (now the value is -1 by default)
  Double_t polx = 0.;     // Polarisation
  Double_t poly = 0.;
  Double_t polz = 0.;
  Int_t ntr = 0;                  // Track number; to be filled by the stack
  Int_t status = generatorStatus; // Generation status
 
  // correct for tracks which are in list before generator is called
  if (mother1 != -1) {
    mother1 += fMCIndexOffset;
  }
  if (mother2 != -1) {
    mother2 += fMCIndexOffset;
  }
  if (daughter1 != -1) {
    daughter1 += fMCIndexOffset;
  }
  if (daughter2 != -1) {
    daughter2 += fMCIndexOffset;
  }
 
  if (abs(pdgid) == 311 && doTracking) {
    LOG(warn) << "K0/antiK0 requested for tracking: converting into K0s/K0L";
    pdgid = gRandom->Uniform() < 0.5 ? 310 : 130;
  }
 
  if (e < 0) {
    double mass = particlePDG ? particlePDG->Mass() : 0.;
    e = std::sqrt(mass * mass + px * px + py * py + pz * pz);
  }
 
  auto stack = dynamic_cast<o2::data::Stack*>(fStack);
  if (!stack) {
    LOG(fatal) << "Stack must be an o2::data:Stack";
    return; // must be the o2 stack
  }
  stack->PushTrack(doTracking, mother1, pdgid, px, py, pz,
                   e, vx, vy, vz, tof, polx, poly, polz, TMCProcess::kPPrimary, ntr,
                   weight, status, mother2, daughter1, daughter2, proc);
 
  fNTracks++;
}
 
/*****************************************************************/
 
void PrimaryGenerator::AddTrack(Int_t pdgid, Double_t px, Double_t py,
                                Double_t pz, Double_t vx, Double_t vy,
                                Double_t vz, Int_t parent, Bool_t wanttracking,
                                Double_t e, Double_t tof, Double_t weight, TMCProcess proc)
{
  // Do this to encode status code correctly. In FairRoot's PrimaryGenerator, this is simply one number that is set to 0.
  // So we basically do the same.
  // Assuming that this is treated as the HepMC status code down the line (as it used to be).
  AddTrack(pdgid, px, py, pz, vx, vy, vz, parent, -1, -1, -1, wanttracking,
           e, tof, weight, proc, mcgenstatus::MCGenStatusEncoding(0, 0).fullEncoding);
}
 
/*****************************************************************/
 
void PrimaryGenerator::setInteractionVertex(const MCEventHeader* event)
{
  if (!mApplyVertex) {
    return;
  }
  Double_t xyz[3] = {event->GetX(), event->GetY(), event->GetZ()};
  SetBeam(xyz[0], xyz[1], 0., 0.);
  SetTarget(xyz[2], 0.);
  SmearVertexXY(false);
  SmearVertexZ(false);
  SmearGausVertexXY(false);
  SmearGausVertexZ(false);
}
 
/*****************************************************************/
void PrimaryGenerator::setExternalVertexForNextEvent(double x, double y, double z)
{
  if (!mApplyVertex) {
    return;
  }
  mExternalVertexX = x;
  mExternalVertexY = y;
  mExternalVertexZ = z;
  mHaveExternalVertex = true;
}
 
/*****************************************************************/
 
void PrimaryGenerator::setVertexMode(o2::conf::VertexMode const& mode, o2::dataformats::MeanVertexObject const* v)
{
  mVertexMode = mode;
  if (mode == o2::conf::VertexMode::kCCDB) {
    if (!v) {
      LOG(fatal) << "A valid MeanVertexObject needs to be passed with option o2::conf::VertexMode::kCCDB";
    }
    mMeanVertex = std::move(std::unique_ptr<o2::dataformats::MeanVertexObject>(new o2::dataformats::MeanVertexObject(*v)));
    LOG(info) << "The mean vertex is set to :";
    mMeanVertex->print();
  }
  if (mVertexMode == o2::conf::VertexMode::kNoVertex) {
    setApplyVertex(false);
    LOG(info) << "Disabling vertexing";
    mMeanVertex = std::move(std::unique_ptr<o2::dataformats::MeanVertexObject>(new o2::dataformats::MeanVertexObject(0, 0, 0, 0, 0, 0, 0, 0)));
    LOG(info) << "The mean vertex is set to :";
    mMeanVertex->print();
  }
}
 
/*****************************************************************/
 
void PrimaryGenerator::fixInteractionVertex()
{
  if (!mApplyVertex) {
    SetBeam(0., 0., 0., 0.);
    SetTarget(0., 0.);
    return;
  }
 
  // if someone gave vertex from outside; we will take it
  if (mHaveExternalVertex) {
    SetBeam(mExternalVertexX, mExternalVertexY, 0., 0.);
    SetTarget(mExternalVertexZ, 0.);
    mHaveExternalVertex = false; // the vertex is now consumed
    return;
  }
 
  // sampling a vertex and fixing for next event; no smearing will be done
  // inside FairPrimaryGenerator;
  SmearVertexXY(false);
  SmearVertexZ(false);
  SmearGausVertexXY(false);
  SmearGausVertexZ(false);
 
  // we use the mMeanVertexObject if initialized (initialize first)
  if (mMeanVertex.get() == nullptr) {
    if (mVertexMode == o2::conf::VertexMode::kDiamondParam) {
      auto const& param = InteractionDiamondParam::Instance();
      const auto& xyz = param.position;
      const auto& sigma = param.width;
      mMeanVertex = std::move(std::unique_ptr<o2::dataformats::MeanVertexObject>(new o2::dataformats::MeanVertexObject(xyz[0], xyz[1], xyz[2], sigma[0], sigma[1], sigma[2], param.slopeX, param.slopeY)));
    }
    if (mVertexMode == o2::conf::VertexMode::kNoVertex) {
      mMeanVertex = std::move(std::unique_ptr<o2::dataformats::MeanVertexObject>(new o2::dataformats::MeanVertexObject(0, 0, 0, 0, 0, 0, 0, 0)));
    }
    if (mVertexMode == o2::conf::VertexMode::kCCDB) {
      // fatal.. then the object should have been passed with setting
      LOG(fatal) << "MeanVertexObject is null ... but mode is kCCDB. Please inject the valid CCDB object via setVertexMode";
    }
  }
  auto sampledvertex = mMeanVertex->sample();
 
  if (PrimaryGeneratorParam::Instance().verbose) {
    LOG(info) << "Sampled interacting vertex " << sampledvertex;
  }
  SetBeam(sampledvertex.X(), sampledvertex.Y(), 0., 0.);
  SetTarget(sampledvertex.Z(), 0.);
}
 
/*****************************************************************/
 
Bool_t PrimaryGenerator::embedInto(TString fname)
{
  if (mEmbedFile && mEmbedFile->IsOpen()) {
    LOG(error) << "Another embedding file is currently open";
    return kFALSE;
  }
 
  mEmbedFile = TFile::Open(fname);
  if (!mEmbedFile || !mEmbedFile->IsOpen()) {
    LOG(error) << "Cannot open file for embedding: " << fname;
    return kFALSE;
  }
 
  mEmbedTree = (TTree*)mEmbedFile->Get("o2sim");
  if (!mEmbedTree) {
    LOG(error) << R"(Cannot find "o2sim" tree for embedding in )" << fname;
    return kFALSE;
  }
 
  mEmbedEntries = mEmbedTree->GetEntries();
  if (mEmbedEntries <= 0) {
    LOG(error) << "Invalid number of entries found in tree for embedding: " << mEmbedEntries;
    return kFALSE;
  }
 
  mEmbedEvent = new MCEventHeader;
  mEmbedTree->SetBranchAddress("MCEventHeader.", &mEmbedEvent);
 
  return kTRUE;
}
 
/*****************************************************************/
 
void PrimaryGenerator::setGeneratorInformation()
{
  auto o2event = dynamic_cast<MCEventHeader*>(fEvent);
  if (o2event) {
    o2event->putInfo<int>(o2::mcgenid::GeneratorProperty::GENERATORID, mGeneratorId);
    o2event->putInfo<std::string>(o2::mcgenid::GeneratorProperty::GENERATORDESCRIPTION, mGeneratorDescription);
  }
}
 
/*****************************************************************/
/*****************************************************************/
 
} /* namespace eventgen */
} /* namespace o2 */
 
ClassImp(o2::eventgen::PrimaryGenerator);