Commit b2b2a9f9 authored by Alexander Leopold's avatar Alexander Leopold Committed by Tadej Novak
Browse files

Hgtd extension, the algorithm

parent 0fb91904
################################################################################
# Package: HGTD_TrackTimeExtension
################################################################################
# Declare the package name:
atlas_subdir( HGTD_TrackTimeExtension )
# Declare the package's dependencies:
atlas_depends_on_subdirs( PUBLIC
Control/AthenaBaseComps
GaudiKernel)
# Component(s) in the package:
atlas_add_component( TrackTimeExtensionAlgorithm
src/*.cxx
src/components/*.cxx
LINK_LIBRARIES AthenaBaseComps GaudiKernel
xAODTracking HGTD_TrackTimeExtensionTools HGTD_RIO_OnTrack HGTD_TruthTools GeneratorObjects
xAODTruth InDetSimData)
# Install files from the package:
atlas_install_headers( HGTD_TrackTimeExtension )
atlas_install_python_modules( python/*.py)
atlas_install_joboptions( share/*.py )
/**
* Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration.
*
* @file HGTD_TrackTimeExtension/TrackTimeExtensionAlg.h
* @author Alexander Leopold <alexander.leopold@cern.ch>
* @author
* @date September, 2021
*
* @brief Calls the track extension tool and adds the output in addition to
* truth information as decorations to the track particles.
* In future iterations, the output of the extension can be a dedicated object
* written to file, with the TrackParticle linking to its exension (if it has
* one).
*
* TODO:
* - add all relevant decorations
* - access the truth information to classify extensions
*/
#ifndef HGTD_RECALGS_TRACKTIMEEXTENSIONALG_H
#define HGTD_RECALGS_TRACKTIMEEXTENSIONALG_H
#include "AthenaBaseComps/AthAlgorithm.h"
#include "GaudiKernel/ToolHandle.h"
#include "GeneratorObjects/McEventCollection.h"
#include "HGTD_PrepRawData/HGTD_ClusterContainer.h"
#include "HGTD_RecToolInterfaces/IHGTD_ClusterTruthTool.h"
#include "HGTD_RecToolInterfaces/IHGTD_TrackTimeExtensionTool.h"
#include "StoreGate/ReadHandleKey.h"
#include "xAODTracking/TrackParticle.h"
#include "xAODTracking/TrackParticleContainer.h"
#include "InDetSimData/InDetSimDataCollection.h"
#include <string>
namespace HGTD {
using HGTDExtension_t =
std::array<std::unique_ptr<const Trk::TrackStateOnSurface>, 4>;
class TrackTimeExtensionAlg : public AthAlgorithm {
public:
TrackTimeExtensionAlg(const std::string& name, ISvcLocator* pSvcLocator);
virtual ~TrackTimeExtensionAlg() {}
virtual StatusCode initialize() override final;
virtual StatusCode execute() override final;
private:
StatusCode decorateTrackParticle(const xAOD::TrackParticle* track_ptkl,
const HGTDExtension_t& extension,
const InDetSimDataCollection* sdo_collection,
const HepMC::GenEvent* hs_event,
bool skip_deco = false);
ToolHandle<HGTD::IHGTD_TrackTimeExtensionTool> m_extension_tool;
SG::ReadHandleKey<HGTD::HGTD_ClusterContainer> m_clustercont_rh_key;
SG::ReadHandleKey<InDetSimDataCollection> m_sdo_coll_rh_key;
SG::ReadHandleKey<McEventCollection> m_mc_coll_rh_key;
SG::ReadHandleKey<xAOD::TrackParticleContainer> m_trk_ptkl_rh_key;
std::string m_deco_prefix;
float m_eta_cut;
ToolHandle<HGTD::IHGTD_ClusterTruthTool> m_truth_tool;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<bool>>>
m_dec_layer_has_ext;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<float>>>
m_dec_layer_ext_chi2;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<float>>>
m_dec_layer_cluster_raw_time;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<float>>>
m_dec_layer_cluster_time;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<int>>>
m_dec_layer_cluster_truth_class;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<bool>>>
m_dec_layer_cluster_cluster_shadowed;
std::unique_ptr<SG::AuxElement::Decorator<std::vector<bool>>>
m_dec_layer_cluster_cluster_merged;
// to look for primary hits that have been missed, additional info is needed.
// this will go in in a next iteration
// std::unique_ptr<SG::AuxElement::Decorator<std::vector<bool>>>
// m_dec_layer_cluster_cluster_expected;
};
} // namespace HGTD
#endif // HGTD_RECALGS_TRACKTIMEEXTENSIONALG_H
extrapolator = ToolSvc.AtlasExtrapolator
kalman_updator = ToolSvc.InDetUpdator
from HGTD_TimeCalibrationTools.HGTD_TimeCalibrationToolsConf import HGTD__StraightLineTOFcorrectionTool
hgtd_tof_corr_tool = HGTD__StraightLineTOFcorrectionTool(name="StraightLineTOFcorrectionTool")
hgtd_tof_corr_tool.OutputLevel = VERBOSE
ToolSvc += hgtd_tof_corr_tool
from HGTD_TrackTimeExtensionTools.HGTD_TrackTimeExtensionToolsConf import HGTD__HGTD_IterativeExtensionTool
htgd_ext_tool = HGTD__HGTD_IterativeExtensionTool(name="IterativeExtensionTool")
htgd_ext_tool.OutputLevel = VERBOSE
htgd_ext_tool.ExtrapolatorTool = extrapolator
htgd_ext_tool.UpdatorTool = kalman_updator
htgd_ext_tool.TOFCorrTool = hgtd_tof_corr_tool
htgd_ext_tool.Chi2Cut = 5.0
ToolSvc += htgd_ext_tool
from HGTD_TruthTools.HGTD_TruthToolsConf import HGTD__ClusterTruthTool
hgtd_cluster_truth_tool = HGTD__ClusterTruthTool(name="ClusterTruthTool")
ToolSvc += hgtd_cluster_truth_tool
from HGTD_TrackTimeExtension.TrackTimeExtensionAlgConf import HGTD__TrackTimeExtensionAlg
hgtd_alg = HGTD__TrackTimeExtensionAlg(name="TrackTimeExtensionAlg")
hgtd_alg.OutputLevel = VERBOSE
hgtd_alg.TimeExtensionTool = htgd_ext_tool
hgtd_alg.TruthTool = hgtd_cluster_truth_tool
topSequence += hgtd_alg
# print "========= NOW SETTING UP FOR OUTPUT ============"
# # --- load setup
from AthenaPoolCnvSvc.WriteAthenaPool import AthenaPoolOutputStream
# --- check dictionary
ServiceMgr.AthenaSealSvc.CheckDictionary = True
# --- commit interval (test)
ServiceMgr.AthenaPoolCnvSvc.CommitInterval = 10
# MC truth information
truthList = [ 'xAOD::TruthEventContainer#TruthEvents',
'xAOD::TruthEventAuxContainer#TruthEventsAux.',
'xAOD::TruthEventContainer#TruthEvents_PU',
'xAOD::TruthEventAuxContainer#TruthEvents_PUAux.',
'xAOD::TruthEventContainer#SpclMCPU',
'xAOD::TruthEventAuxContainer#SpclMCPUAux.',
'xAOD::TruthPileupEventContainer#TruthPileupEvents',
'xAOD::TruthPileupEventAuxContainer#TruthPileupEventsAux.',
'xAOD::TruthParticleContainer#TruthParticles',
'xAOD::TruthParticleAuxContainer#TruthParticlesAux.',
'xAOD::TruthVertexContainer#TruthVertices',
'xAOD::TruthVertexAuxContainer#TruthVerticesAux.',
'PileUpEventInfo#OverlayEvent' ]
# --- create stream
StreamESD = AthenaPoolOutputStream ( "StreamESD2", "ESD.HGTD.pool.root",asAlg=True)
# --- save MC collections if truth turned on
StreamESD.ItemList += truthList
# ---- load list of objects
include ( "InDetRecExample/WriteInDetESD.py" )
StreamESD.ItemList += InDetESDList
StreamESD.ItemList += [ 'xAOD::EventInfo#EventInfo' , 'xAOD::EventAuxInfo#EventInfoAux.' ]
StreamESD.ForceRead = True # otherwise unread stuff is not copied
# StreamESD.ItemList += InDetESDList;
#topSequence.StreamESD.ItemList+=[\"xAOD::EventInfo#EventInfo\", \"xAOD::EventAuxInfo#EventInfoAux.\"];topSequence.StreamESD.ItemList += [ \"xAOD::TrackParticleAuxContainer#InDetTrackParticlesAuxDyn.HGTD_*\"];topSequence.StreamESD.ForceRead = True;
/**
* Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration.
*
* @file HGTD_TrackTimeExtension/src/TrackTimeExtensionAlg.cxx
* @author Alexander Leopold <alexander.leopold@cern.ch>
* @author
* @date July, 2021
* @brief
*/
#include "HGTD_TrackTimeExtension/TrackTimeExtensionAlg.h"
#include "HGTD_RIO_OnTrack/HGTD_ClusterOnTrack.h"
#include "StoreGate/WriteHandle.h"
#include "xAODTruth/TruthParticleContainer.h"
namespace {
static constexpr unsigned short n_hgtd_layers = 4;
}
namespace HGTD {
TrackTimeExtensionAlg::TrackTimeExtensionAlg(const std::string& name,
ISvcLocator* pSvcLocator)
: AthAlgorithm(name, pSvcLocator),
m_extension_tool(
"HGTD_IterativeExtensionTool/HGTD_IterativeExtensionTool"),
m_deco_prefix("HGTD"),
m_eta_cut(2.38),
m_truth_tool("ClusterTruthTool/ClusterTruthTool") {
declareProperty("TimeExtensionTool", m_extension_tool);
declareProperty("HGTDClusterContainerName",
m_clustercont_rh_key = "HGTD_Cluster");
declareProperty("HGTD_SDOCollectionName", m_sdo_coll_rh_key = "HGTD_SDO_Map");
declareProperty("MC_CollectionName", m_mc_coll_rh_key = "TruthEvent");
declareProperty("TrackParticleContainerName",
m_trk_ptkl_rh_key = "InDetTrackParticles");
declareProperty("DecorationPrefix", m_deco_prefix);
declareProperty("MinEtaAcceptance", m_eta_cut);
declareProperty("TruthTool", m_truth_tool);
}
StatusCode TrackTimeExtensionAlg::initialize() {
ATH_CHECK(m_extension_tool.retrieve());
ATH_CHECK(m_clustercont_rh_key.initialize());
ATH_CHECK(m_sdo_coll_rh_key.initialize());
ATH_CHECK(m_mc_coll_rh_key.initialize());
ATH_CHECK(m_trk_ptkl_rh_key.initialize());
m_dec_layer_has_ext =
std::make_unique<SG::AuxElement::Decorator<std::vector<bool>>>(
m_deco_prefix + "_has_extension");
m_dec_layer_ext_chi2 =
std::make_unique<SG::AuxElement::Decorator<std::vector<float>>>(
m_deco_prefix + "_extension_chi2");
m_dec_layer_cluster_raw_time =
std::make_unique<SG::AuxElement::Decorator<std::vector<float>>>(
m_deco_prefix + "_cluster_raw_time");
m_dec_layer_cluster_time =
std::make_unique<SG::AuxElement::Decorator<std::vector<float>>>(
m_deco_prefix + "_cluster_time");
m_dec_layer_cluster_truth_class =
std::make_unique<SG::AuxElement::Decorator<std::vector<int>>>(
m_deco_prefix + "_cluster_truth_class");
m_dec_layer_cluster_cluster_shadowed =
std::make_unique<SG::AuxElement::Decorator<std::vector<bool>>>(
m_deco_prefix + "_cluster_shadowed");
m_dec_layer_cluster_cluster_merged =
std::make_unique<SG::AuxElement::Decorator<std::vector<bool>>>(
m_deco_prefix + "_cluster_merged");
// m_dec_layer_cluster_cluster_expected =
// std::make_unique<SG::AuxElement::Decorator<std::vector<bool>>>(
// m_deco_prefix + "_cluster_expected");
return StatusCode::SUCCESS;
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
StatusCode TrackTimeExtensionAlg::execute() {
ATH_MSG_DEBUG("Start event");
SG::ReadHandle<HGTD::HGTD_ClusterContainer> cluster_container_handle(
m_clustercont_rh_key);
const HGTD::HGTD_ClusterContainer* cluster_container =
cluster_container_handle.cptr();
if (not cluster_container) {
ATH_MSG_ERROR("[TrackTimeExtensionAlg] HGTD_ClusterContainer not found, "
"aborting execute!");
return StatusCode::FAILURE;
}
SG::ReadHandle<InDetSimDataCollection> sdo_collection_handle(
m_sdo_coll_rh_key);
const InDetSimDataCollection* sdo_collection = sdo_collection_handle.cptr();
if (not sdo_collection) {
ATH_MSG_ERROR(
"[TrackTimeExtensionAlg] SDO Collection not found, aborting execute!");
return StatusCode::FAILURE;
}
SG::ReadHandle<McEventCollection> mc_collection_handle(m_mc_coll_rh_key);
const McEventCollection* mc_coll = mc_collection_handle.cptr();
if (not mc_coll) {
ATH_MSG_ERROR("[TrackTimeExtensionAlg] McEventCollection not found, "
"aborting execute!");
return StatusCode::FAILURE;
}
const HepMC::GenEvent* hs_event = nullptr;
if (mc_coll->size() > 0) {
hs_event = mc_coll->at(0);
} else {
hs_event = nullptr;
}
// for each track, run the extension if the track is in HGTD acceptance
SG::ReadHandle<xAOD::TrackParticleContainer> trk_ptkl_container_handle(
m_trk_ptkl_rh_key);
const xAOD::TrackParticleContainer* track_particles =
trk_ptkl_container_handle.cptr();
if (not track_particles) {
ATH_MSG_ERROR("[TrackTimeExtensionAlg] TrackParticleContainer not found, "
"aborting execute!");
return StatusCode::FAILURE;
}
for (const auto* track_ptkl : *track_particles) {
ATH_MSG_DEBUG("Track eta: " << track_ptkl->eta()
<< " pt: " << track_ptkl->pt());
HGTDExtension_t extension{nullptr, nullptr, nullptr, nullptr};
if (std::abs(track_ptkl->eta()) < m_eta_cut) {
ATH_MSG_DEBUG("Track out of acceptance");
// decorate all track particle objects to avoid issues with the
// decorations
ATH_CHECK(decorateTrackParticle(track_ptkl, extension, sdo_collection,
hs_event, true));
continue;
}
// this should not happen?
if (track_ptkl->track() == nullptr) {
ATH_MSG_DEBUG("There is no Trk::Track");
ATH_CHECK(decorateTrackParticle(track_ptkl, extension, sdo_collection,
hs_event, true));
continue;
}
// return 4 track states on surface objects as a result of the extension
extension = m_extension_tool->extendTrackToHGTD(*(track_ptkl->track()),
cluster_container);
// TODO here:
// retrieve truth info for associated and non-associated HGTD hits, merging
// or shadowing info
// decorate the track
ATH_CHECK(decorateTrackParticle(track_ptkl, extension, sdo_collection,
hs_event, false));
} // END LOOP over tracks
return StatusCode::SUCCESS;
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
StatusCode TrackTimeExtensionAlg::decorateTrackParticle(
const xAOD::TrackParticle* track_ptkl, const HGTDExtension_t& extension,
const InDetSimDataCollection* sdo_collection,
const HepMC::GenEvent* hs_event, bool skip_deco) {
std::vector<bool> has_cluster_vec;
has_cluster_vec.reserve(n_hgtd_layers);
std::vector<float> chi2_vec;
chi2_vec.reserve(n_hgtd_layers);
std::vector<float> raw_time_vec;
raw_time_vec.reserve(n_hgtd_layers);
std::vector<float> time_vec;
time_vec.reserve(n_hgtd_layers);
std::vector<int> truth_vec;
truth_vec.reserve(n_hgtd_layers);
std::vector<bool> is_shadowed_vec;
is_shadowed_vec.reserve(n_hgtd_layers);
std::vector<bool> is_merged_vec;
is_merged_vec.reserve(n_hgtd_layers);
for (const std::unique_ptr<const Trk::TrackStateOnSurface>& trk_state :
extension) {
if (trk_state) {
has_cluster_vec.emplace_back(true);
chi2_vec.emplace_back(
trk_state->fitQualityOnSurface()->chiSquared() /
trk_state->fitQualityOnSurface()->doubleNumberDoF());
const HGTD::HGTD_ClusterOnTrack* cot =
dynamic_cast<const HGTD::HGTD_ClusterOnTrack*>(
trk_state->measurementOnTrack());
time_vec.emplace_back(cot->time());
// get the cluster
const HGTD::HGTD_Cluster* cluster = cot->prepRawData();
raw_time_vec.emplace_back(cluster->time());
// get the truth particle
static SG::AuxElement::Accessor<ElementLink<xAOD::TruthParticleContainer>>
acc_tpl("truthParticleLink");
const xAOD::TruthParticle* truth_particle = nullptr;
if (acc_tpl.isAvailable(*track_ptkl)) {
auto truth_match_link = acc_tpl(*track_ptkl);
if (truth_match_link.isValid()) {
truth_particle = *truth_match_link;
}
}
HGTD::ClusterTruthOrigin truth_origin = ClusterTruthOrigin::UNIDENTIFIED;
bool is_shadowed = false;
bool is_merged = false;
if (truth_particle) {
auto truth_info = m_truth_tool->classifyCluster(
cluster, truth_particle, sdo_collection, hs_event);
truth_origin = truth_info.origin;
is_shadowed = truth_info.is_shadowed;
is_merged = truth_info.is_merged;
}
truth_vec.emplace_back(truth_origin);
is_shadowed_vec.emplace_back(is_shadowed);
is_merged_vec.emplace_back(is_merged);
} else {
has_cluster_vec.emplace_back(false);
if (not skip_deco) {
chi2_vec.emplace_back(-1.);
raw_time_vec.emplace_back(-1.);
time_vec.emplace_back(-1.);
truth_vec.emplace_back(-1);
is_shadowed_vec.emplace_back(false);
is_merged_vec.emplace_back(false);
}
}
} // END LOOP over TrackStateOnSurface
m_dec_layer_has_ext->set(*track_ptkl, has_cluster_vec);
m_dec_layer_ext_chi2->set(*track_ptkl, chi2_vec);
m_dec_layer_cluster_raw_time->set(*track_ptkl, raw_time_vec);
m_dec_layer_cluster_time->set(*track_ptkl, time_vec);
m_dec_layer_cluster_truth_class->set(*track_ptkl, truth_vec);
m_dec_layer_cluster_cluster_shadowed->set(*track_ptkl, is_shadowed_vec);
m_dec_layer_cluster_cluster_merged->set(*track_ptkl, is_merged_vec);
return StatusCode::SUCCESS;
}
} // namespace HGTD
#include "GaudiKernel/DeclareFactoryEntries.h"
#include "HGTD_TrackTimeExtension/TrackTimeExtensionAlg.h"
DECLARE_NAMESPACE_ALGORITHM_FACTORY(HGTD, TrackTimeExtensionAlg)
DECLARE_FACTORY_ENTRIES(TrackTimeExtensionAlg) {
DECLARE_NAMESPACE_ALGORITHM(HGTD, TrackTimeExtensionAlg)
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment