Commit 467a96b4 authored by Dan Guest's avatar Dan Guest Committed by Alexander Froch
Browse files

Consolidate dump-single-btag and SingleBTagAlg

parent 769a5314
......@@ -59,6 +59,7 @@ atlas_add_library(dataset-dumper
src/errorLogger.cxx
src/streamers.cxx
src/trackSort.cxx
src/processSingleBTagEvent.cxx
PUBLIC_HEADERS src
LINK_LIBRARIES ${LINK_LIBRARIES}
)
......
#include "SingleBTagAlg.h"
#include "SingleBTagTools.hh"
#include "SafeShallowCopy.hh"
#include "TruthTools.hh"
#include "xAODBTagging/BTaggingUtilities.h"
#include "SingleBTagConfig.hh"
#include "processSingleBTagEvent.hh"
#include "H5Cpp.h"
#include "AsgDataHandles/ReadHandle.h"
namespace {
void check_rc(StatusCode code) {
if (!code.isSuccess()) throw std::runtime_error("bad return code");
}
constexpr float operator "" _GeV(unsigned long long d) { return d*1e3; }
constexpr float operator "" _GeV(long double d) { return d*1e3; }
const xAOD::Vertex* primary(const xAOD::VertexContainer& vertices) {
if (vertices.size() == 0) {
throw std::runtime_error("no primary vertices");
}
for ( const xAOD::Vertex *vertex : vertices ) {
if ( vertex->vertexType() == xAOD::VxType::PriVtx ) {
return vertex;
}
}
// if we find nothing else this should be the beam spot
return vertices.front();
}
}
SingleBTagAlg::SingleBTagAlg(const std::string& name,
ISvcLocator* pSvcLocator):
EL::AnaAlgorithm(name, pSvcLocator),
......@@ -72,142 +43,15 @@ StatusCode SingleBTagAlg::execute () {
const auto& jobcfg = *m_config;
auto& tools = *m_tools;
const xAOD::JetContainer *raw_jets = nullptr;
check_rc( event.retrieve(raw_jets, jobcfg.jet_collection) );
ATH_MSG_DEBUG("Processing " + std::to_string(raw_jets->size()) + " jets");
// first loop: add decorations to jets. These have to be done
// before calibration to be consistent with reconstruction
auto [jets, jetaux] = safeShallowCopyContainer(*raw_jets);
for (const xAOD::Jet* uncalib_jet: *jets) {
// this is more important stuff
const auto* btag = xAOD::BTaggingUtilities::getBTagging(*uncalib_jet);
if (!btag) throw std::runtime_error("no btaggingLink");
if (jobcfg.decorate.jet_aug) {
tools.jet_augmenter.augment(*btag, *btag);
}
if (jobcfg.decorate.soft_muon) {
tools.muon_augmenter.augment(*btag);
}
if (jobcfg.decorate.btag_jes){
tools.jet_augmenter.augmentBtagJes(*btag, *btag);
}
for (const auto& dl2: tools.dl2s) {
dl2.decorate(*btag);
}
}
if (jobcfg.do_calibration) {
check_rc(tools.calibration_tool.applyCalibration(*jets));
}
// sort jets by descending pt
//
// we make a new container first to preserve the indices
std::vector<const xAOD::Jet*> sorted_jets(jets->begin(), jets->end());
std::sort(sorted_jets.begin(), sorted_jets.end(),
[](const auto* j1, const auto* j2) {
return j1->pt() > j2->pt();
});
// second loop: select calibrated jets and write out to HDF5
//
// we'll need truth particles and event info
//
const xAOD::EventInfo *event_info = nullptr;
check_rc( event.retrieve(event_info, "EventInfo") );
std::vector<const xAOD::TruthParticle*> truth_leptons;
try {
const xAOD::TruthParticleContainer *tpc = nullptr;
check_rc( event.retrieve(tpc, "TruthParticles") );
std::vector<const xAOD::TruthParticle*> truth_particles(
tpc->begin(), tpc->end());
truth_leptons = truth::getLeptonsFromWZ(truth_particles);
m_counts.n_successful_truth_reads++;
} catch (truth::TruthRecordError& err) {
ATH_MSG_ERROR(getEventInfoString(err, *event_info));
m_counts.n_failed_truth_reads++;
return StatusCode::SUCCESS;
}
// save some primary vertex information on eventinfo
const xAOD::VertexContainer* primary_vertices = nullptr;
check_rc( event.retrieve(primary_vertices, jobcfg.vertex_collection));
tools.dec.n_primary_vertices(*event_info) = primary_vertices->size();
const xAOD::TruthEventContainer* truthEventContainer = nullptr;
const xAOD::TruthVertex* truth_PV = nullptr;
if ( jobcfg.decorate.track_truth_info ) {
check_rc( event.retrieve(truthEventContainer, "TruthEvents"));
truth_PV = truthEventContainer->at(0)->truthVertex(0); // truthEventContainer always has size == 1?
}
tools.dec.primary_vertex_detector_z(*event_info) =
primary(*primary_vertices)->z();
std::vector<const xAOD::Jet*> jets_to_write;
unsigned int rank = 0;
for (const xAOD::Jet* calib_jet: sorted_jets) {
if (truth::is_overlaping_lepton(*calib_jet, truth_leptons, 0.3)) {
continue;
}
tools.dec.jet_rank(*calib_jet) = rank++;
if (jobcfg.do_calibration){
// if we're calibrating jets we need to check the JVT again
float updated_jvt_value= tools.jvttool.updateJvt(*calib_jet);
tools.dec.jvt(*calib_jet) = updated_jvt_value;
bool fail_jvt = (
calib_jet->pt() > 20_GeV &&
calib_jet->pt() < 60_GeV &&
std::abs(calib_jet->eta()) < 2.4 &&
updated_jvt_value < jobcfg.jvt_cut );
if (fail_jvt) continue;
if ( ! tools.cleaning_tool.keep(*calib_jet)) continue;
}
else {
tools.dec.jvt(*calib_jet) = NAN;
}
if (calib_jet->pt() < jobcfg.pt_cut || std::abs(calib_jet->eta()) > 2.5) {
continue;
}
jets_to_write.push_back(calib_jet);
for (auto& tracktool: tools.tracks ) {
const xAOD::Jet* uncalib_jet = raw_jets->at(calib_jet->index());
auto tracks = tracktool.selector.get_tracks(*uncalib_jet);
for (const auto& track: tracks) {
tracktool.track_accessor.augment_with_ip(*track, *uncalib_jet);
if ( jobcfg.decorate.track_sv_info ) {
m_trkVertexDecorator.decorate(*track, *uncalib_jet);
}
}
if ( jobcfg.decorate.track_truth_info ) {
m_trkTruthDecorator.decorateAll(tracks, truth_PV);
}
sort(tracks.begin(), tracks.end(), tracktool.sort);
tracktool.writer.write(tracks, *uncalib_jet);
}
}
if (!jets_to_write.empty()) {
tools.jet_writer.write(jets_to_write, event_info);
}
// most of the real work happens in this function
processSingleBTagEvent(event, jobcfg, tools);
return StatusCode::SUCCESS;
}
StatusCode SingleBTagAlg::finalize () {
if (m_metadata_file_name.size() > 0) {
writeTruthCorruptionCounts(m_counts, m_metadata_file_name);
writeTruthCorruptionCounts(m_tools->truth_counts, m_metadata_file_name);
}
return StatusCode::SUCCESS;
......
#include "SingleBTagConfig.hh"
#include "TruthCorruptionCounter.hh"
#include "TrackTruthDecorator.hh"
#include "TrackVertexDecorator.hh"
#include "xAODBTagging/BTaggingContainer.h"
#include "xAODJet/JetContainer.h"
......@@ -43,9 +37,4 @@ private:
SG::ReadHandleKey<xAOD::JetContainer> m_jetKey{this, "jetKey", "", ""};
std::vector<std::unique_ptr<SG::ReadDecorHandleKey<xAOD::BTagging>>> m_bDec;
TruthCorruptionCounter m_counts;
//MT: add the decorators taking inspiration from dump_single_btag.cxx
TrackTruthDecorator m_trkTruthDecorator;
TrackVertexDecorator m_trkVertexDecorator;
};
......@@ -5,6 +5,10 @@
#include "BTagTrackWriter.hh"
#include "TrackSelector.hh"
#include "trackSort.hh"
#include "DecoratorExample.hh"
#include "TrackTruthDecorator.hh"
#include "TrackVertexDecorator.hh"
#include "TruthCorruptionCounter.hh"
#include "FlavorTagDiscriminants/BTagJetAugmenter.h"
#include "FlavorTagDiscriminants/BTagTrackIpAccessor.h"
......@@ -60,6 +64,13 @@ struct SingleBTagTools {
std::vector<TrackTools> tracks;
Decorators dec;
TrackTruthDecorator trkTruthDecorator;
TrackVertexDecorator trkVertexDecorator;
DecoratorExample example_decorator;
TruthCorruptionCounter truth_counts;
};
#endif
#include "processSingleBTagEvent.hh"
// the implementation is the same for every signature, we just have to
// recompile it a few times below. We do it here to avoid having the
// template headers interfere with anything that calls this function.
#include "SingleBTagTools.hh"
#include "SingleBTagConfig.hh"
#include "SafeShallowCopy.hh"
#include "TruthTools.hh"
#include "xAODRootAccess/TEvent.h"
#include "xAODCore/ShallowCopy.h"
#include "xAODTruth/TruthParticleContainer.h"
#include "xAODJet/JetContainer.h"
#include "xAODTracking/TrackParticleContainer.h"
#include "xAODBTagging/BTaggingUtilities.h"
#include "PathResolver/PathResolver.h"
namespace {
void check_rc(StatusCode code) {
if (!code.isSuccess()) throw std::runtime_error("bad return code");
}
constexpr float operator "" _GeV(unsigned long long d) { return d*1e3; }
constexpr float operator "" _GeV(long double d) { return d*1e3; }
const xAOD::Vertex* primary(const xAOD::VertexContainer& vertices) {
if (vertices.size() == 0) {
throw std::runtime_error("no primary vertices");
}
for ( const xAOD::Vertex *vertex : vertices ) {
if ( vertex->vertexType() == xAOD::VxType::PriVtx ) {
return vertex;
}
}
// if we find nothing else this should be the beam spot
return vertices.front();
}
}
// this is the templated code, the concrete instances are below
template <typename Event>
void processSingleBTagEventImpl(Event& event,
const SingleBTagConfig& jobcfg,
SingleBTagTools& tools) {
const xAOD::JetContainer *raw_jets = nullptr;
check_rc( event.retrieve(raw_jets, jobcfg.jet_collection) );
// first loop: add decorations to jets. These have to be done
// before calibration to be consistent with reconstruction
auto [jets, jetaux] = safeShallowCopyContainer(*raw_jets);
for (const xAOD::Jet* uncalib_jet: *jets) {
// this is more important stuff
const auto* btag = xAOD::BTaggingUtilities::getBTagging(*uncalib_jet);
if (!btag) throw std::runtime_error("no btaggingLink");
if (jobcfg.decorate.jet_aug) {
tools.jet_augmenter.augment(*btag, *btag);
}
if (jobcfg.decorate.soft_muon) {
tools.muon_augmenter.augment(*btag);
}
if (jobcfg.decorate.btag_jes){
tools.jet_augmenter.augmentBtagJes(*btag, *btag);
}
for (const auto& dl2: tools.dl2s) {
dl2.decorate(*btag);
}
}
if (jobcfg.do_calibration) {
check_rc(tools.calibration_tool.applyCalibration(*jets));
}
// sort jets by descending pt
//
// we make a new container first to preserve the indices
std::vector<const xAOD::Jet*> sorted_jets(jets->begin(), jets->end());
std::sort(sorted_jets.begin(), sorted_jets.end(),
[](const auto* j1, const auto* j2) {
return j1->pt() > j2->pt();
});
// second loop: select calibrated jets and write out to HDF5
//
// we'll need truth particles and event info
//
const xAOD::EventInfo *event_info = nullptr;
check_rc( event.retrieve(event_info, "EventInfo") );
std::vector<const xAOD::TruthParticle*> truth_leptons;
try {
const xAOD::TruthParticleContainer *tpc = nullptr;
check_rc( event.retrieve(tpc, "TruthParticles") );
std::vector<const xAOD::TruthParticle*> truth_particles(
tpc->begin(), tpc->end());
truth_leptons = truth::getLeptonsFromWZ(truth_particles);
tools.truth_counts.n_successful_truth_reads++;
} catch (truth::TruthRecordError& err) {
std::cerr << getEventInfoString(err, *event_info) << std::endl;
tools.truth_counts.n_failed_truth_reads++;
return;
}
// save some primary vertex information on eventinfo
const xAOD::VertexContainer* primary_vertices = nullptr;
check_rc( event.retrieve(primary_vertices, jobcfg.vertex_collection));
tools.dec.n_primary_vertices(*event_info) = primary_vertices->size();
const xAOD::TruthEventContainer* truthEventContainer = nullptr;
const xAOD::TruthVertex* truth_PV = nullptr;
if ( jobcfg.decorate.track_truth_info ) {
check_rc( event.retrieve(truthEventContainer, "TruthEvents"));
// truthEventContainer always has size == 1?
truth_PV = truthEventContainer->at(0)->truthVertex(0);
}
tools.dec.primary_vertex_detector_z(*event_info) =
primary(*primary_vertices)->z();
std::vector<const xAOD::Jet*> jets_to_write;
unsigned int rank = 0;
for (const xAOD::Jet* calib_jet: sorted_jets) {
if (truth::is_overlaping_lepton(*calib_jet, truth_leptons, 0.3)) {
continue;
}
tools.dec.jet_rank(*calib_jet) = rank++;
if (jobcfg.do_calibration){
// if we're calibrating jets we need to check the JVT again
float updated_jvt_value= tools.jvttool.updateJvt(*calib_jet);
tools.dec.jvt(*calib_jet) = updated_jvt_value;
bool fail_jvt = (
calib_jet->pt() > 20_GeV &&
calib_jet->pt() < 60_GeV &&
std::abs(calib_jet->eta()) < 2.4 &&
updated_jvt_value < jobcfg.jvt_cut );
if (fail_jvt) continue;
if ( ! tools.cleaning_tool.keep(*calib_jet)) continue;
}
else {
tools.dec.jvt(*calib_jet) = NAN;
}
if (calib_jet->pt() < jobcfg.pt_cut || std::abs(calib_jet->eta()) > 2.5) {
continue;
}
jets_to_write.push_back(calib_jet);
for (auto& tracktool: tools.tracks ) {
const xAOD::Jet* uncalib_jet = raw_jets->at(calib_jet->index());
auto tracks = tracktool.selector.get_tracks(*uncalib_jet);
for (const auto& track: tracks) {
tracktool.track_accessor.augment_with_ip(*track, *uncalib_jet);
if ( jobcfg.decorate.track_sv_info ) {
tools.trkVertexDecorator.decorate(*track, *uncalib_jet);
}
}
if ( jobcfg.decorate.track_truth_info ) {
tools.trkTruthDecorator.decorateAll(tracks, truth_PV);
}
sort(tracks.begin(), tracks.end(), tracktool.sort);
tracktool.writer.write(tracks, *uncalib_jet);
}
}
if (!jets_to_write.empty()) {
tools.jet_writer.write(jets_to_write, event_info);
}
}
// Concrete versions of the templated function above. These are the
// ones that are exposed to be used in other files.
//
void processSingleBTagEvent(xAOD::TEvent& e,
const SingleBTagConfig& c,
SingleBTagTools& t) {
processSingleBTagEventImpl(e, c, t);
}
// StoreGateSvc isn't defined in AnalysisBase...
#ifndef XAOD_STANDALONE
void processSingleBTagEvent(StoreGateSvc& e,
const SingleBTagConfig& c,
SingleBTagTools& t) {
processSingleBTagEventImpl(e, c, t);
}
#else
// ... but SgTEvent is
void processSingleBTagEvent(asg::SgTEvent& e,
const SingleBTagConfig& c,
SingleBTagTools& t) {
processSingleBTagEventImpl(e, c, t);
}
#endif // XAOD_STANDALONE
#ifndef PROCESS_SINGLE_BTAG_EVENT_HH
#define PROCESS_SINGLE_BTAG_EVENT_HH
class SingleBTagConfig;
class SingleBTagTools;
namespace xAOD {
class TEvent;
}
// the joy of "duel use": there doesn't appear to be a common
// baseclass that implements the one method we need in TEvent and
// StoreGate (i.e. `retrieve`). So we use overloads here, templates
// are defined in processSingleBTagEvent.tcc
void processSingleBTagEvent(xAOD::TEvent&,
const SingleBTagConfig&,
SingleBTagTools&);
// we only build this guy if we're using a Gaudi build
#ifndef XAOD_STANDALONE
class StoreGateSvc;
void processSingleBTagEvent(StoreGateSvc&,
const SingleBTagConfig&,
SingleBTagTools&);
#else
// for the non-Gaudi build, there's SgTEvent, which is some kind of
// wrapper for TEvent that looks (kind of) like StoreGate
namespace asg {
class SgTEvent;
}
void processSingleBTagEvent(asg::SgTEvent&,
const SingleBTagConfig&,
SingleBTagTools&);
#endif // XAOD_STANDALONE
#endif // PROCESS_SINGLE_BTAG_EVENT_HH
#include <cstddef>
#include <memory>
#include <cmath>
#include "src/SingleBTagConfig.hh"
#include "SingleBTagConfig.hh"
#include "SingleBTagTools.hh"
#include "SingleBTagOptions.hh"
#include "DecoratorExample.hh"
#include "TrackTruthDecorator.hh"
#include "TrackVertexDecorator.hh"
#include "TruthTools.hh"
#include "SafeShallowCopy.hh"
#include "TruthCorruptionCounter.hh"
#include "processSingleBTagEvent.hh"
#include "xAODRootAccess/Init.h"
#include "xAODRootAccess/TEvent.h"
#include "xAODCore/ShallowCopy.h"
#include "xAODTruth/TruthParticleContainer.h"
#include "xAODJet/JetContainer.h"
#include "xAODTracking/TrackParticleContainer.h"
#include "xAODBTagging/BTaggingUtilities.h"
#include "PathResolver/PathResolver.h"
#include "H5Cpp.h"
......@@ -37,9 +17,6 @@ namespace {
void check_rc(StatusCode code) {
if (!code.isSuccess()) throw std::runtime_error("bad return code");
}
constexpr float operator "" _GeV(unsigned long long d) { return d*1e3; }
constexpr float operator "" _GeV(long double d) { return d*1e3; }
}
int main (int argc, char *argv[]) {
......@@ -58,13 +35,6 @@ int main (int argc, char *argv[]) {
H5::H5File output(opts.out, H5F_ACC_TRUNC);
SingleBTagTools tools(jobcfg, output);
// this is just an example augmenter, it doesn't do anything important
DecoratorExample example_decorator;
TrackTruthDecorator trkTruthDecorator;
TrackVertexDecorator trkVertexDecorator;
TruthCorruptionCounter truth_counts;
// Loop over the specified files:
for (const std::string& file: opts.in) {
// Open the file:
......@@ -95,122 +65,14 @@ int main (int argc, char *argv[]) {
Info( APP_NAME, "Processing entry %lld / %lld", entry, entries );
}
const xAOD::JetContainer *raw_jets = nullptr;
check_rc( event.retrieve(raw_jets, jobcfg.jet_collection) );
const xAOD::TruthParticleContainer *tpc = nullptr;
check_rc( event.retrieve(tpc, "TruthParticles") );
std::vector<const xAOD::TruthParticle*> truth_particles(
tpc->begin(), tpc->end());
const xAOD::EventInfo *event_info = nullptr;
check_rc( event.retrieve(event_info, "EventInfo") );
std::vector<const xAOD::TruthParticle*> truth_leptons;
try {
truth_leptons = truth::getLeptonsFromWZ(truth_particles);
truth_counts.n_successful_truth_reads++;
} catch (truth::TruthRecordError& err) {
truth_counts.n_failed_truth_reads++;
std::cerr << getEventInfoString(err, *event_info) << std::endl;
continue;
}
const xAOD::VertexContainer* primary_vertices = nullptr;
check_rc( event.retrieve(primary_vertices, jobcfg.vertex_collection));
tools.dec.n_primary_vertices(*event_info) = primary_vertices->size();
// most of the real work happens in this function
processSingleBTagEvent(event, jobcfg, tools);
const xAOD::TruthEventContainer* truthEventContainer = nullptr;
const xAOD::TruthVertex* truth_PV = nullptr;
if ( jobcfg.decorate.track_truth_info ) {
check_rc( event.retrieve(truthEventContainer, "TruthEvents"));
truth_PV = truthEventContainer->at(0)->truthVertex(0); // truthEventContainer always has size == 1?
}
} // end event loop
// first loop: add decorations to jets. These have to be done
// before calibration to be consistent with reconstruction
auto [jets, jetaux] = safeShallowCopyContainer(*raw_jets);
for (const xAOD::Jet *const uncalib_jet : *jets) {
// this is just applying example decorations, for no reason
// other than to show that they work.
example_decorator.decorate(*uncalib_jet);
// this is more important stuff
const xAOD::BTagging* btag = xAOD::BTaggingUtilities::getBTagging(
*uncalib_jet);
if (!btag) throw