MCKinematicsReader.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 MC_KINEMATICS_READER_H
#define MC_KINEMATICS_READER_H
 
#include "SimulationDataFormat/DigitizationContext.h"
#include "SimulationDataFormat/MCTrack.h"
#include "SimulationDataFormat/MCCompLabel.h"
#include "SimulationDataFormat/MCEventHeader.h"
#include "SimulationDataFormat/TrackReference.h"
#include "SimulationDataFormat/MCTruthContainer.h"
#include <vector>
 
class TChain;
 
namespace o2
{
 
namespace steer
{
 
class MCKinematicsReader
{
 public:
  enum class Mode {
    kDigiContext,
    kMCKine
  };
 
  MCKinematicsReader() = default;
 
  ~MCKinematicsReader();
 
  MCKinematicsReader(std::string_view name, Mode mode = Mode::kDigiContext)
  {
    if (mode == Mode::kMCKine) {
      initFromKinematics(name);
    } else if (mode == Mode::kDigiContext) {
      initFromDigitContext(name);
    }
  }
 
  MCKinematicsReader(o2::steer::DigitizationContext const* context)
  {
    initFromDigitContext(context);
  }
 
  bool initFromDigitContext(std::string_view filename);
 
  bool initFromDigitContext(o2::steer::DigitizationContext const* digicontext);
 
  bool initFromKinematics(std::string_view filename);
 
  bool isInitialized() const { return mInitialized; }
 
  MCTrack const* getTrack(o2::MCCompLabel const&) const;
 
  MCTrack const* getTrack(int source, int event, int track) const;
 
  MCTrack const* getTrack(int event, int track) const;
 
  std::vector<MCTrack> const& getTracks(int source, int event) const;
 
  void releaseTracksForSourceAndEvent(int source, int event);
 
  std::vector<MCTrack> const& getTracks(int event) const;
 
 
 
 
  gsl::span<o2::TrackReference> getTrackRefs(int source, int event, int track) const;
  const std::vector<o2::TrackReference>& getTrackRefsByEvent(int source, int event) const;
  gsl::span<o2::TrackReference> getTrackRefs(int event, int track) const;
 
  o2::dataformats::MCEventHeader const& getMCEventHeader(int source, int event) const;
 
  size_t getNSources() const;
 
  size_t getNEvents(int source) const;
 
  DigitizationContext const* getDigitizationContext() const
  {
    return mDigitizationContext;
  }
 
 private:
  void ensureTracksForSourceAndEvent(int source, int event) const;
  [[noreturn]] static void reportMissingSource(int source, size_t available);
  [[noreturn]] static void reportMissingEvent(const char* what, int source, int event, size_t available);
 
  void initTracksForSource(int source) const;
  void loadTracksForSourceAndEvent(int source, int eventID) const;
  void loadHeadersForSource(int source) const;
  void loadTrackRefsForSource(int source) const;
  void initIndexedTrackRefs(std::vector<o2::TrackReference>& refs, o2::dataformats::MCTruthContainer<o2::TrackReference>& indexedrefs) const;
 
  DigitizationContext const* mDigitizationContext = nullptr;
  bool mOwningDigiContext = false;
 
  // chains for each source
  std::vector<TChain*> mInputChains;
 
  // a vector of tracks foreach source and each collision
  mutable std::vector<std::vector<std::vector<o2::MCTrack>*>> mTracks;                                       // the in-memory track container
  mutable std::vector<std::vector<o2::dataformats::MCEventHeader>> mHeaders;                                 // the in-memory header container
  mutable std::vector<std::vector<o2::dataformats::MCTruthContainer<o2::TrackReference>>> mIndexedTrackRefs; // the in-memory track ref container
 
  bool mInitialized = false; // whether initialized
};
 
inline MCTrack const* MCKinematicsReader::getTrack(o2::MCCompLabel const& label) const
{
  const auto source = label.getSourceID();
  const auto event = label.getEventID();
  const auto track = label.getTrackID();
  return getTrack(source, event, track);
}
 
inline MCTrack const* MCKinematicsReader::getTrack(int source, int event, int track) const
{
  auto const& tracks = getTracks(source, event);
  // one comparison, and it covers the negative track ID of a hit whose track was not kept
  return static_cast<size_t>(track) < tracks.size() ? &tracks[track] : nullptr;
}
 
inline MCTrack const* MCKinematicsReader::getTrack(int event, int track) const
{
  return getTrack(0, event, track);
}
 
inline std::vector<MCTrack> const& MCKinematicsReader::getTracks(int source, int event) const
{
  if (static_cast<size_t>(source) >= mTracks.size()) {
    reportMissingSource(source, mTracks.size());
  }
  auto& perEvent = mTracks[source];
  if (perEvent.size() == 0) {
    initTracksForSource(source);
  }
  // the event range shares the branch of the lazy load, so the fast path grows by one comparison
  if (static_cast<size_t>(event) >= perEvent.size() || perEvent[event] == nullptr) {
    ensureTracksForSourceAndEvent(source, event);
  }
  return *perEvent[event];
}
 
inline std::vector<MCTrack> const& MCKinematicsReader::getTracks(int event) const
{
  return getTracks(0, event);
}
 
inline o2::dataformats::MCEventHeader const& MCKinematicsReader::getMCEventHeader(int source, int event) const
{
  auto const& headers = mHeaders.at(source);
  if (headers.size() == 0) {
    loadHeadersForSource(source);
  }
  if (static_cast<size_t>(event) >= headers.size()) {
    reportMissingEvent("event headers", source, event, headers.size());
  }
  return headers[event];
}
 
inline gsl::span<o2::TrackReference> MCKinematicsReader::getTrackRefs(int source, int event, int track) const
{
  if (static_cast<size_t>(source) >= mIndexedTrackRefs.size()) {
    return {};
  }
  auto& perEvent = mIndexedTrackRefs[source];
  if (perEvent.size() == 0) {
    loadTrackRefsForSource(source);
  }
  if (static_cast<size_t>(event) >= perEvent.size()) {
    return {};
  }
  return perEvent[event].getLabels(track);
}
 
inline const std::vector<o2::TrackReference>& MCKinematicsReader::getTrackRefsByEvent(int source, int event) const
{
  auto const& perEvent = mIndexedTrackRefs.at(source);
  if (perEvent.size() == 0) {
    loadTrackRefsForSource(source);
  }
  if (static_cast<size_t>(event) >= perEvent.size()) {
    reportMissingEvent("events of track references", source, event, perEvent.size());
  }
  return perEvent[event].getTruthArray();
}
 
inline gsl::span<o2::TrackReference> MCKinematicsReader::getTrackRefs(int event, int track) const
{
  return getTrackRefs(0, event, track);
}
 
inline size_t MCKinematicsReader::getNSources() const
{
  return mTracks.size();
}
 
inline size_t MCKinematicsReader::getNEvents(int source) const
{
  if (mTracks[source].size() == 0) {
    initTracksForSource(source);
  }
  return mTracks[source].size();
}
 
} // namespace steer
} // namespace o2
 
#endif