From 2a570659497e947d79c81ef488e05b9fd7fd2d64 Mon Sep 17 00:00:00 2001 From: doloresgarcia Date: Wed, 5 Jun 2024 17:39:25 +0200 Subject: [PATCH] add collection of hits output --- Tracking/include/GenFitter.h | 2 ++ Tracking/src/GenFitter.cpp | 23 +++++++++++++++-------- Tracking/test/test_example.py | 2 +- runIDEAtracker.py | 1 + 4 files changed, 19 insertions(+), 9 deletions(-) diff --git a/Tracking/include/GenFitter.h b/Tracking/include/GenFitter.h index baaf3a1..50f9c97 100644 --- a/Tracking/include/GenFitter.h +++ b/Tracking/include/GenFitter.h @@ -20,6 +20,7 @@ namespace edm4hep { #include "extension/TrackCollection.h" #include "extension/TrackCollection.h" +#include "extension/TrackerHit3DCollection.h" // EDM4HEP extension #include "extension/DriftChamberDigiCollection.h" #include "extension/DriftChamberDigiLocalCollection.h" @@ -82,4 +83,5 @@ class GenFitter : public Gaudi::Algorithm { // Output track collection name mutable DataHandle m_output_tracks{"outputTracks", Gaudi::DataHandle::Writer, this}; + mutable DataHandle m_output_hits{"outputHits", Gaudi::DataHandle::Writer, this}; }; diff --git a/Tracking/src/GenFitter.cpp b/Tracking/src/GenFitter.cpp index 9baf50d..dc82339 100644 --- a/Tracking/src/GenFitter.cpp +++ b/Tracking/src/GenFitter.cpp @@ -71,6 +71,7 @@ GenFitter::GenFitter(const std::string& aName, ISvcLocator* aSvcLoc) : Gaudi::Al declareProperty("inputHits_VTXD", m_input_hits_VTXD, "Input VTXD tracker hit collection name"); declareProperty("inputHits_VTXOB", m_input_hits_VTXOB, "Input VTXOB tracker hit collection name"); declareProperty("outputTracks", m_output_tracks, "Output track collection name"); + declareProperty("outputHits", m_output_hits, "Output hits collection name"); } GenFitter::~GenFitter() {} @@ -200,6 +201,7 @@ StatusCode GenFitter::execute(const EventContext&) const { // std::cout << "indices: " << clustering.index({indices}) << std::endl; extension::TrackCollection* output_tracks = m_output_tracks.createAndPut(); + extension::TrackerHit3DCollection* output_hits= m_output_hits.createAndPut(); int64_t number_of_tracks = unique_tensor.numel(); std::cout << "number_of_tracks: " << number_of_tracks << std::endl; @@ -222,9 +224,9 @@ StatusCode GenFitter::execute(const EventContext&) const { if ((torch::sum(mask_VTXD)>0).item()){ // The hit belong to vtxd auto hit = inputHits_VTXD->at(index_id.item()); - extension::MutableTrackerHit3D hit_extension; + auto hit_extension = output_hits->create(); hit_extension.setCellID(hit.getCellID()); - hit_extension.setType(hit.getType()); + hit_extension.setType(1); hit_extension.setEDep(hit.getEDep()); hit_extension.setPosition(hit.getPosition()); // output_track.addToTrackerHits(hit_extension); @@ -234,9 +236,9 @@ StatusCode GenFitter::execute(const EventContext&) const { index_id = index_id-it_0; // std::cout << "taking hit from inputHits_VTXIB" << std::endl; auto hit = inputHits_VTXIB->at(index_id.item()); - extension::MutableTrackerHit3D hit_extension; + auto hit_extension = output_hits->create(); hit_extension.setCellID(hit.getCellID()); - hit_extension.setType(hit.getType()); + hit_extension.setType(1); hit_extension.setEDep(hit.getEDep()); hit_extension.setPosition(hit.getPosition()); // output_track.addToTrackerHits(hit_extension); @@ -244,17 +246,22 @@ StatusCode GenFitter::execute(const EventContext&) const { index_id = index_id-(it_1+it_0); auto hit = inputHits_VTXOB->at(index_id.item()); // std::cout << "taking hit from VTOB" << std::endl; - extension::MutableTrackerHit3D hit_extension; + auto hit_extension = output_hits->create(); hit_extension.setCellID(hit.getCellID()); - hit_extension.setType(hit.getType()); + hit_extension.setType(1); hit_extension.setEDep(hit.getEDep()); hit_extension.setPosition(hit.getPosition()); - // output_track.addToTrackerHits(hit_extension); + output_track.addToTrackerHits(hit_extension); } else if ((torch::sum(mask_CDC)>0).item()){ index_id = index_id-(it_1+it_2 +it_0); // std::cout << "taking hit from CDC" << std::endl; auto hit = input_hits_CDC->at(index_id.item()); - output_track.addToTrackerHits(hit); + auto hit_extension = output_hits->create(); + hit_extension.setCellID(hit.getCellID()); + hit_extension.setType(0); + hit_extension.setEDep(hit.getEDep()); + hit_extension.setPosition(hit.getLeftPosition()); + output_track.addToTrackerHits(hit_extension); } } } diff --git a/Tracking/test/test_example.py b/Tracking/test/test_example.py index 6f8f18e..8f22e44 100644 --- a/Tracking/test/test_example.py +++ b/Tracking/test/test_example.py @@ -20,7 +20,7 @@ hits_in_track = track.getTrackerHits() coordinates = [] for hit in hits_in_track: - position = hit.getLeftPosition() + position = hit.getPosition() x = position.x y = position.y z = position.z diff --git a/runIDEAtracker.py b/runIDEAtracker.py index cbac463..a4993a4 100644 --- a/runIDEAtracker.py +++ b/runIDEAtracker.py @@ -21,6 +21,7 @@ inputHits_VTXIB="VTXIBDigis", inputHits_VTXOB="VTXOBDigis", outputTracks="CDCHTracks", + outputHits="Hits", OutputLevel=INFO, )