RecoParticleFlow/PFSimProducer/plugins/SimPFProducer.cc - master - githubmirror/cmssw

RecoParticleFlow/PFSimProducer/plugins/SimPFProducer.cc
445 строк20 KB
Andrea Bocci
Clean up unused variables
19 фев 2025, 00:29
19 фев 2025, 00:293ef1d8b
// This producer assigns vertex times (with a specified resolution) to tracks.
// The times are produced as valuemaps associated to tracks, so the track dataformat doesn't
// need to be modified.

#include "CLHEP/Units/SystemOfUnits.h"
#include "DataFormats/Common/interface/ValueMap.h"
#include "DataFormats/Common/interface/View.h"
#include "DataFormats/EgammaReco/interface/SuperCluster.h"
#include "DataFormats/EgammaReco/interface/SuperClusterFwd.h"
#include "DataFormats/GsfTrackReco/interface/GsfTrack.h"
#include "DataFormats/GsfTrackReco/interface/GsfTrackFwd.h"
#include "DataFormats/MuonReco/interface/Muon.h"
#include "DataFormats/MuonReco/interface/MuonFwd.h"
#include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"
#include "DataFormats/ParticleFlowCandidate/interface/PFCandidateFwd.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlock.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockFwd.h"
#include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
#include "DataFormats/TrackReco/interface/Track.h"
#include "DataFormats/TrackReco/interface/TrackFwd.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/global/EDProducer.h"
#include "FWCore/MessageLogger/interface/MessageLogger.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
#include "FWCore/Utilities/interface/isFinite.h"
#include "FWCore/Utilities/interface/transform.h"
#include "SimDataFormats/Associations/interface/TrackToTrackingParticleAssociator.h"
#include "SimDataFormats/CaloAnalysis/interface/CaloParticle.h"
#include "SimDataFormats/CaloAnalysis/interface/CaloParticleFwd.h"
#include "SimDataFormats/CaloAnalysis/interface/SimCluster.h"
#include "SimDataFormats/CaloAnalysis/interface/SimClusterFwd.h"
#include "SimDataFormats/Track/interface/SimTrackContainer.h"
#include "SimDataFormats/TrackingAnalysis/interface/TrackingParticle.h"
#include "SimDataFormats/TrackingAnalysis/interface/TrackingVertex.h"
#include "SimDataFormats/TrackingAnalysis/interface/TrackingVertexContainer.h"
#include "SimDataFormats/Vertex/interface/SimVertexContainer.h"
#include "SimTracker/TrackerHitAssociation/interface/TrackerHitAssociator.h"

#include <unordered_map>
#include <memory>

class SimPFProducer : public edm::global::EDProducer<> {
public:
  SimPFProducer(const edm::ParameterSet&);

  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);

  void produce(edm::StreamID, edm::Event&, const edm::EventSetup&) const override;

private:
  // parameters
  const double superClusterThreshold_, neutralEMThreshold_, neutralHADThreshold_;
  const bool useTiming_;
  const bool useTimingQuality_;
  const double timingQualityThreshold_;

  // inputs
  const edm::EDGetTokenT<edm::View<reco::PFRecTrack>> pfRecTracks_;
  const edm::EDGetTokenT<edm::View<reco::Track>> tracks_;
  const edm::EDGetTokenT<edm::View<reco::Track>> gsfTracks_;
  const edm::EDGetTokenT<reco::MuonCollection> muons_;
  const edm::EDGetTokenT<edm::ValueMap<float>> srcTrackTime_, srcTrackTimeError_, srcTrackTimeQuality_;
  const edm::EDGetTokenT<edm::ValueMap<float>> srcGsfTrackTime_, srcGsfTrackTimeError_, srcGsfTrackTimeQuality_;
  const edm::EDGetTokenT<TrackingParticleCollection> trackingParticles_;
  const edm::EDGetTokenT<SimClusterCollection> simClustersTruth_;
  const edm::EDGetTokenT<CaloParticleCollection> caloParticles_;
  const edm::EDGetTokenT<std::vector<reco::PFCluster>> simClusters_;
  // tracking particle associators by order of preference
  const std::vector<edm::EDGetTokenT<reco::TrackToTrackingParticleAssociator>> associators_;
};

#include "FWCore/Framework/interface/MakerMacros.h"
DEFINE_FWK_MODULE(SimPFProducer);

void SimPFProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  // simPFProducer
  edm::ParameterSetDescription desc;
  desc.add<edm::InputTag>("simClustersSrc", {"particleFlowClusterHGCalFromSimCl"});
  desc.add<edm::InputTag>("trackSrc", {"generalTracks"});
  desc.add<std::vector<edm::InputTag>>("associators",
                                       {
                                           {"quickTrackAssociatorByHits"},
                                       });
  desc.add<edm::InputTag>("pfRecTrackSrc", {"hgcalTrackCollection", "TracksInHGCal"});
  desc.add<edm::InputTag>("trackingParticleSrc", {"mix", "MergedTrackTruth"});
  desc.add<double>("neutralEMThreshold", 0.25);
  desc.add<edm::InputTag>("caloParticlesSrc", {"mix", "MergedCaloTruth"});
  desc.add<double>("superClusterThreshold", 4.0);
  desc.add<edm::InputTag>("simClusterTruthSrc", {"mix", "MergedCaloTruth"});
  desc.add<edm::InputTag>("muonSrc", {"muons1stStep"});
  desc.add<double>("neutralHADThreshold", 0.25);
  desc.add<edm::InputTag>("gsfTrackSrc", {"electronGsfTracks"});

  // if useTiming_
  desc.add<bool>("useTiming", false);
  desc.add<bool>("useTimingQuality", false);
  desc.add<edm::InputTag>("trackTimeValueMap", {});
  desc.add<edm::InputTag>("trackTimeErrorMap", {});
  desc.add<edm::InputTag>("trackTimeQualityMap", {});
  desc.add<double>("timingQualityThreshold", 0);
  desc.add<edm::InputTag>("gsfTrackTimeValueMap", {});
  desc.add<edm::InputTag>("gsfTrackTimeErrorMap", {});
  desc.add<edm::InputTag>("gsfTrackTimeQualityMap", {});

  descriptions.addWithDefaultLabel(desc);
}

namespace {

  template <typename T>
  uint64_t hashSimInfo(const T& simTruth, size_t i = 0) {
    uint64_t evtid = simTruth.eventId().rawId();
    uint64_t trackid = simTruth.g4Tracks()[i].trackId();
    return ((evtid << 3) + 23401923) ^ trackid;
  };
}  // namespace

SimPFProducer::SimPFProducer(const edm::ParameterSet& conf)
    : superClusterThreshold_(conf.getParameter<double>("superClusterThreshold")),
      neutralEMThreshold_(conf.getParameter<double>("neutralEMThreshold")),
      neutralHADThreshold_(conf.getParameter<double>("neutralHADThreshold")),
      useTiming_(conf.getParameter<bool>("useTiming")),
      useTimingQuality_(conf.getParameter<bool>("useTimingQuality")),
      timingQualityThreshold_(useTimingQuality_ ? conf.getParameter<double>("timingQualityThreshold") : -99.),
      pfRecTracks_(consumes<edm::View<reco::PFRecTrack>>(conf.getParameter<edm::InputTag>("pfRecTrackSrc"))),
      tracks_(consumes<edm::View<reco::Track>>(conf.getParameter<edm::InputTag>("trackSrc"))),
      gsfTracks_(consumes<edm::View<reco::Track>>(conf.getParameter<edm::InputTag>("gsfTrackSrc"))),
      muons_(consumes<reco::MuonCollection>(conf.getParameter<edm::InputTag>("muonSrc"))),
      srcTrackTime_(useTiming_ ? consumes<edm::ValueMap<float>>(conf.getParameter<edm::InputTag>("trackTimeValueMap"))
                               : edm::EDGetTokenT<edm::ValueMap<float>>()),
      srcTrackTimeError_(useTiming_
                             ? consumes<edm::ValueMap<float>>(conf.getParameter<edm::InputTag>("trackTimeErrorMap"))
                             : edm::EDGetTokenT<edm::ValueMap<float>>()),
      srcTrackTimeQuality_(useTimingQuality_
                               ? consumes<edm::ValueMap<float>>(conf.getParameter<edm::InputTag>("trackTimeQualityMap"))
                               : edm::EDGetTokenT<edm::ValueMap<float>>()),
      srcGsfTrackTime_(useTiming_
                           ? consumes<edm::ValueMap<float>>(conf.getParameter<edm::InputTag>("gsfTrackTimeValueMap"))
                           : edm::EDGetTokenT<edm::ValueMap<float>>()),
      srcGsfTrackTimeError_(
          useTiming_ ? consumes<edm::ValueMap<float>>(conf.getParameter<edm::InputTag>("gsfTrackTimeErrorMap"))
                     : edm::EDGetTokenT<edm::ValueMap<float>>()),
      srcGsfTrackTimeQuality_(
          useTimingQuality_ ? consumes<edm::ValueMap<float>>(conf.getParameter<edm::InputTag>("gsfTrackTimeQualityMap"))
                            : edm::EDGetTokenT<edm::ValueMap<float>>()),
      trackingParticles_(consumes<TrackingParticleCollection>(conf.getParameter<edm::InputTag>("trackingParticleSrc"))),
      simClustersTruth_(consumes<SimClusterCollection>(conf.getParameter<edm::InputTag>("simClusterTruthSrc"))),
      caloParticles_(consumes<CaloParticleCollection>(conf.getParameter<edm::InputTag>("caloParticlesSrc"))),
      simClusters_(consumes<std::vector<reco::PFCluster>>(conf.getParameter<edm::InputTag>("simClustersSrc"))),
      associators_(edm::vector_transform(
          conf.getParameter<std::vector<edm::InputTag>>("associators"),
          [this](const edm::InputTag& tag) { return this->consumes<reco::TrackToTrackingParticleAssociator>(tag); })) {
  produces<reco::PFBlockCollection>();
  produces<reco::SuperClusterCollection>("perfect");
  produces<reco::PFCandidateCollection>();
}

void SimPFProducer::produce(edm::StreamID, edm::Event& evt, const edm::EventSetup& es) const {
  //get associators
  std::vector<edm::Handle<reco::TrackToTrackingParticleAssociator>> associators;
  for (const auto& token : associators_) {
    associators.emplace_back();
    auto& back = associators.back();
    evt.getByToken(token, back);
  }

  //get PFRecTrack
  edm::Handle<edm::View<reco::PFRecTrack>> PFTrackCollectionH;
  evt.getByToken(pfRecTracks_, PFTrackCollectionH);
  const edm::View<reco::PFRecTrack> PFTrackCollection = *PFTrackCollectionH;
  std::unordered_set<unsigned> PFTrackToGeneralTrack;
  for (unsigned i = 0; i < PFTrackCollection.size(); ++i) {
    const auto ptr = PFTrackCollection.ptrAt(i);
    PFTrackToGeneralTrack.insert(ptr->trackRef().key());
  }

  //get track collections
  edm::Handle<edm::View<reco::Track>> TrackCollectionH;
  evt.getByToken(tracks_, TrackCollectionH);
  const edm::View<reco::Track>& TrackCollection = *TrackCollectionH;

  edm::Handle<reco::MuonCollection> muons;
  evt.getByToken(muons_, muons);
  std::unordered_set<unsigned> MuonTrackToGeneralTrack;
  for (auto const& mu : *muons.product()) {
    reco::TrackRef muTrkRef = mu.track();
    if (muTrkRef.isNonnull())
      MuonTrackToGeneralTrack.insert(muTrkRef.key());
  }

  // get timing, if enabled
  edm::Handle<edm::ValueMap<float>> trackTimeH, trackTimeErrH, trackTimeQualH, gsfTrackTimeH, gsfTrackTimeErrH,
      gsfTrackTimeQualH;
  if (useTiming_) {
    evt.getByToken(srcTrackTime_, trackTimeH);
    evt.getByToken(srcTrackTimeError_, trackTimeErrH);
    evt.getByToken(srcGsfTrackTime_, gsfTrackTimeH);
    evt.getByToken(srcGsfTrackTimeError_, gsfTrackTimeErrH);

    if (useTimingQuality_) {
      evt.getByToken(srcTrackTimeQuality_, trackTimeQualH);
      evt.getByToken(srcGsfTrackTimeQuality_, gsfTrackTimeQualH);
    }
  }

  //get tracking particle collections
  edm::Handle<TrackingParticleCollection> TPCollectionH;
  evt.getByToken(trackingParticles_, TPCollectionH);
  //const TrackingParticleCollection&  TPCollection = *TPCollectionH;

  // grab phony clustering information
  edm::Handle<SimClusterCollection> SimClustersTruthH;
  evt.getByToken(simClustersTruth_, SimClustersTruthH);
  const SimClusterCollection& SimClustersTruth = *SimClustersTruthH;

  edm::Handle<CaloParticleCollection> CaloParticlesH;
  evt.getByToken(caloParticles_, CaloParticlesH);
  const CaloParticleCollection& CaloParticles = *CaloParticlesH;

  edm::Handle<std::vector<reco::PFCluster>> SimClustersH;
  evt.getByToken(simClusters_, SimClustersH);
  const std::vector<reco::PFCluster>& SimClusters = *SimClustersH;

  std::unordered_map<uint64_t, size_t> hashToSimCluster;

  for (unsigned i = 0; i < SimClustersTruth.size(); ++i) {
    const auto& simTruth = SimClustersTruth[i];
    hashToSimCluster[hashSimInfo(simTruth)] = i;
  }

  // associate the reco tracks / gsf Tracks
  std::vector<reco::RecoToSimCollection> associatedTracks, associatedTracksGsf;
  for (const auto& associator : associators) {
    associatedTracks.emplace_back(associator->associateRecoToSim(TrackCollectionH, TPCollectionH));
    //associatedTracksGsf.emplace_back(associator->associateRecoToSim(GsfTrackCollectionH, TPCollectionH));
  }

  // make blocks out of calo particles so we can have cluster references
  // likewise fill out superclusters
  auto superclusters = std::make_unique<reco::SuperClusterCollection>();
  auto blocks = std::make_unique<reco::PFBlockCollection>();
  std::unordered_map<size_t, size_t> simCluster2Block;
  std::unordered_map<size_t, size_t> simCluster2BlockIndex;
  std::unordered_multimap<size_t, size_t> caloParticle2SimCluster;
  std::vector<int> caloParticle2SuperCluster;
  for (unsigned icp = 0; icp < CaloParticles.size(); ++icp) {
    blocks->emplace_back();
    auto& block = blocks->back();
    const auto& simclusters = CaloParticles[icp].simClusters();
    double pttot = 0.0;
    double etot = 0.0;
    std::vector<size_t> good_simclusters;
    for (unsigned isc = 0; isc < simclusters.size(); ++isc) {
      auto simc = simclusters[isc];
      auto pdgId = std::abs(simc->pdgId());
      edm::Ref<std::vector<reco::PFCluster>> clusterRef(SimClustersH, simc.key());
      if (((pdgId == 22 || pdgId == 11) && clusterRef->energy() > neutralEMThreshold_) ||
          clusterRef->energy() > neutralHADThreshold_) {
        good_simclusters.push_back(isc);
        etot += clusterRef->energy();
        pttot += clusterRef->pt();
        auto bec = std::make_unique<reco::PFBlockElementCluster>(clusterRef, reco::PFBlockElement::HGCAL);
        block.addElement(bec.get());
        simCluster2Block[simc.key()] = icp;
        simCluster2BlockIndex[simc.key()] = bec->index();
        caloParticle2SimCluster.emplace(CaloParticles[icp].g4Tracks()[0].trackId(), simc.key());
      }
    }

    auto pdgId = std::abs(CaloParticles[icp].pdgId());

    caloParticle2SuperCluster.push_back(-1);
    if ((pdgId == 22 || pdgId == 11) && pttot > superClusterThreshold_) {
      caloParticle2SuperCluster[icp] = superclusters->size();

      math::XYZPoint seedpos;  // take seed pos as supercluster point
      reco::CaloClusterPtr seed;
      reco::CaloClusterPtrVector clusters;
      for (auto idx : good_simclusters) {
        edm::Ptr<reco::PFCluster> ptr(SimClustersH, simclusters[idx].key());
        clusters.push_back(ptr);
        if (seed.isNull() || seed->energy() < ptr->energy()) {
          seed = ptr;
          seedpos = ptr->position();
        }
      }
      superclusters->emplace_back(etot, seedpos, seed, clusters);
    }
  }

  auto blocksHandle = evt.put(std::move(blocks));
  evt.put(std::move(superclusters), "perfect");

  // list tracks so we can mark them as used and/or fight over them
  std::vector<bool> usedTrack(TrackCollection.size(), false),
      //usedGsfTrack(GsfTrackCollection.size(),false),
      usedSimCluster(SimClusters.size(), false);

  auto candidates = std::make_unique<reco::PFCandidateCollection>();
  // in good particle flow fashion, start from the tracks and go out
  for (unsigned itk = 0; itk < TrackCollection.size(); ++itk) {
    auto tkRef = TrackCollection.refAt(itk);
    // skip tracks not selected by PF
    if (PFTrackToGeneralTrack.count(itk) == 0)
      continue;
    reco::RecoToSimCollection::const_iterator assoc_tps = associatedTracks.back().end();
    for (const auto& association : associatedTracks) {
      assoc_tps = association.find(tkRef);
      if (assoc_tps != association.end())
        break;
    }
    if (assoc_tps == associatedTracks.back().end())
      continue;
    // assured now that we are matched to a set of tracks
    const auto& matches = assoc_tps->val;

    const auto absPdgId = std::abs(matches[0].first->pdgId());
    const auto charge = tkRef->charge();
    const auto three_mom = tkRef->momentum();
    constexpr double mpion2 = 0.13957 * 0.13957;
    double energy = std::sqrt(three_mom.mag2() + mpion2);
    math::XYZTLorentzVector trk_p4(three_mom.x(), three_mom.y(), three_mom.z(), energy);

    reco::PFCandidate::ParticleType part_type;

    switch (absPdgId) {
      case 11:
        part_type = reco::PFCandidate::e;
        break;
      case 13:
        part_type = reco::PFCandidate::mu;
        break;
      default:
        part_type = reco::PFCandidate::h;
    }

    candidates->emplace_back(charge, trk_p4, part_type);
    auto& candidate = candidates->back();

    candidate.setTrackRef(tkRef.castTo<reco::TrackRef>());

    if (useTiming_) {
      // check if track-mtd match is of sufficient quality
      const bool assocQuality = useTimingQuality_ ? (*trackTimeQualH)[tkRef] > timingQualityThreshold_ : true;
      if (assocQuality) {
        candidate.setTime((*trackTimeH)[tkRef], (*trackTimeErrH)[tkRef]);
      } else {
        candidate.setTime(0., -1.);
      }
    }

    // bind to cluster if there is one and try to gather conversions, etc
    for (const auto& match : matches) {
      uint64_t hash = hashSimInfo(*(match.first));
      if (hashToSimCluster.count(hash)) {
        auto simcHash = hashToSimCluster[hash];

        if (!usedSimCluster[simcHash]) {
          if (simCluster2Block.count(simcHash) && simCluster2BlockIndex.count(simcHash)) {
            size_t block = simCluster2Block.find(simcHash)->second;
            size_t blockIdx = simCluster2BlockIndex.find(simcHash)->second;
            edm::Ref<reco::PFBlockCollection> blockRef(blocksHandle, block);
            candidate.addElementInBlock(blockRef, blockIdx);
            usedSimCluster[simcHash] = true;
          }
        }
        if (absPdgId == 11) {  // collect brems/conv. brems
          if (simCluster2Block.count(simcHash)) {
            auto block_index = simCluster2Block.find(simcHash)->second;
            auto supercluster_index = caloParticle2SuperCluster[block_index];
            if (supercluster_index != -1) {
              edm::Ref<reco::PFBlockCollection> blockRef(blocksHandle, block_index);
              for (const auto& elem : blockRef->elements()) {
                const auto& ref = elem.clusterRef();
                if (!usedSimCluster[ref.key()]) {
                  candidate.addElementInBlock(blockRef, elem.index());
                  usedSimCluster[ref.key()] = true;
                }
              }

              //*TODO* cluster time is not reliable at the moment, so just keep time from the track if available
            }
          }
        }
      }
      // Now try to include also electrons that have been reconstructed using
      // the GraphCaloParticles. In particular, recover the cases in which the
      // tracking particle associated to the CaloParticle has not left any hits
      // in the calorimeters or, if it had, the cluster has been skipped due to
      // threshold requirements.
      if (caloParticle2SimCluster.count(match.first->g4Tracks()[0].trackId())) {
        auto range = caloParticle2SimCluster.equal_range(match.first->g4Tracks()[0].trackId());
        for (auto it = range.first; it != range.second; ++it) {
          if (!usedSimCluster[it->second]) {
            usedSimCluster[it->second] = true;
            if (simCluster2Block.find(it->second) != simCluster2Block.end()) {
              size_t block = simCluster2Block.find(it->second)->second;
              size_t blockIdx = simCluster2BlockIndex.find(it->second)->second;
              edm::Ref<reco::PFBlockCollection> blockRef(blocksHandle, block);
              candidate.addElementInBlock(blockRef, blockIdx);
            }
          }
        }
      }
    }
    usedTrack[tkRef.key()] = true;
    // remove tracks already used by muons
    if (MuonTrackToGeneralTrack.count(itk) || absPdgId == 13)
      candidates->pop_back();
  }

  // now loop over the non-collected clusters in blocks
  // and turn them into neutral hadrons or photons
  const auto& theblocks = *blocksHandle;
  for (unsigned ibl = 0; ibl < theblocks.size(); ++ibl) {
    reco::PFBlockRef blref(blocksHandle, ibl);
    const auto& elements = theblocks[ibl].elements();
    for (const auto& elem : elements) {
      const auto& ref = elem.clusterRef();
      const auto& simtruth = SimClustersTruth[ref.key()];
      reco::PFCandidate::ParticleType part_type;
      if (!usedSimCluster[ref.key()]) {
        auto absPdgId = std::abs(simtruth.pdgId());
        switch (absPdgId) {
          case 11:
          case 22:
            part_type = reco::PFCandidate::gamma;
            break;
          default:
            part_type = reco::PFCandidate::h0;
        }
        const auto three_mom = (ref->position() - math::XYZPoint(0, 0, 0)).unit() * ref->correctedEnergy();
        math::XYZTLorentzVector clu_p4(three_mom.x(), three_mom.y(), three_mom.z(), ref->correctedEnergy());
        candidates->emplace_back(0, clu_p4, part_type);
        auto& candidate = candidates->back();
        candidate.addElementInBlock(blref, elem.index());
      }
    }
  }

  evt.put(std::move(candidates));
}