GeneratorFromFile.h

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.
 
 
#ifndef ALICEO2_GENERATORFROMFILE_H_
#define ALICEO2_GENERATORFROMFILE_H_
 
#include "FairGenerator.h"
#include "Generators/Generator.h"
#include "Generators/GeneratorFromO2KineParam.h"
#include "SimulationDataFormat/MCEventHeader.h"
#include <TRandom3.h>
#include <random>
 
class TBranch;
class TFile;
class TParticle;
class TGrid;
 
namespace o2
{
namespace eventgen
{
class GeneratorFromFile : public FairGenerator
{
 public:
  GeneratorFromFile() = default;
  GeneratorFromFile(const char* name);
 
  // the FairGenerator interface methods
 
  bool ReadEvent(FairPrimaryGenerator* primGen) override;
 
  // Set from which event to start
  void SetStartEvent(int start);
 
  void SetSkipNonTrackable(bool b) { mSkipNonTrackable = b; }
  void setFixOffShell(bool b) { mFixOffShell = b; }
  bool rejectOrFixKinematics(TParticle& p);
 
 private:
  TFile* mEventFile = nullptr; 
  int mEventCounter = 0;
  int mEventsAvailable = 0;
  bool mSkipNonTrackable = true; 
  bool mFixOffShell = true;      // fix particles with M_assigned != M_calculated
  ClassDefOverride(GeneratorFromFile, 1);
};
 
class GeneratorFromO2Kine : public o2::eventgen::Generator
{
 public:
  GeneratorFromO2Kine() = default;
  GeneratorFromO2Kine(const char* name);
  GeneratorFromO2Kine(O2KineGenConfig const& pars);
 
  bool Init() override;
 
  // the o2 Generator interface methods
  bool generateEvent() override
  { /* trivial - actual work in importParticles */
    return true;
  }
  bool importParticles() override;
 
  // Set from which event to start
  void SetStartEvent(int start);
 
  void setContinueMode(bool val) { mContinueMode = val; };
  void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override;
  const o2::dataformats::MCEventHeader* getOrigMCEventHeader() const { return mOrigMCEventHeader.get(); }
 
 private:
  TFile* mEventFile = nullptr;     
  TBranch* mEventBranch = nullptr; 
  TBranch* mMCHeaderBranch = nullptr; 
  int mEventCounter = 0;
  int mEventsAvailable = 0;
  bool mSkipNonTrackable = true; 
  bool mContinueMode = false;    
  bool mRoundRobin = false;      
  bool mRandomize = false;       
  unsigned int mRngSeed = 0;     
  bool mRandomPhi = false;       
  TGrid* mAlienInstance = nullptr; // a cached connection to TGrid (needed for Alien locations)
  std::unique_ptr<O2KineGenConfig> mConfig; 
 
  std::unique_ptr<o2::dataformats::MCEventHeader> mOrigMCEventHeader; 
 
  ClassDefOverride(GeneratorFromO2Kine, 2);
};
 
class GeneratorFromEventPool : public o2::eventgen::Generator
{
 public:
  constexpr static std::string_view eventpool_filename = "evtpool.root";
  constexpr static std::string_view alien_protocol_prefix = "alien://";
 
  GeneratorFromEventPool() = default; // mainly for ROOT IO
  GeneratorFromEventPool(EventPoolGenConfig const& pars);
 
  bool Init() override;
 
  // the o2 Generator interface methods
  bool generateEvent() override
  { /* trivial - actual work in importParticles */
    return mO2KineGenerator->generateEvent();
  }
  bool importParticles() override
  {
    mO2KineGenerator->clearParticles(); // clear old container before filling with new ones
    auto import_good = mO2KineGenerator->importParticles();
    // transfer the particles (could be avoided)
    mParticles = mO2KineGenerator->getParticles();
    auto original_header = mO2KineGenerator->getOrigMCEventHeader();
    // Workaround to fix vertex shifted particles from event pools (valid for builds released before 14 March 2026)
    if (original_header) {
      double vertex[3] = {original_header->GetX(), original_header->GetY(), original_header->GetZ()};
      if (vertex[0] != 0. || vertex[1] != 0. || vertex[2] != 0.) {
        LOG(debug) << "Subtracting shifted vertex from EventPool: (" << vertex[0] << ", " << vertex[1] << ", " << vertex[2] << ")";
        for (auto& p : mParticles) {
          p.SetProductionVertex(p.Vx() - vertex[0], p.Vy() - vertex[1], p.Vz() - vertex[2], p.T());
        }
      }
    }
    return import_good;
  }
 
  void updateHeader(o2::dataformats::MCEventHeader* eventHeader) override
  {
    // Copy current vertex position from the event header
    const double xyz[3] = {eventHeader->GetX(), eventHeader->GetY(), eventHeader->GetZ()};
    mO2KineGenerator->updateHeader(eventHeader);
    // Event pool uses vertex position from current simulation, only extKinO2 takes the one from the file instead
    eventHeader->SetVertex(xyz[0], xyz[1], xyz[2]);
  }
 
  // determine the collection of available files
  std::vector<std::string> setupFileUniverse(std::string const& path) const;
 
  std::vector<std::string> const& getFileUniverse() const { return mPoolFilesAvailable; }
 
 private:
  EventPoolGenConfig mConfig;                                                    
  std::unique_ptr<o2::eventgen::GeneratorFromO2Kine> mO2KineGenerator = nullptr; 
  std::vector<std::string> mPoolFilesAvailable;                                  
  std::string mFileChosen;                                                       
  // random number generator to determine a concrete file name
  std::mt19937 mRandomEngine; 
 
  ClassDefOverride(GeneratorFromEventPool, 1);
};
 
} // end namespace eventgen
} // end namespace o2
 
#endif