diff --git a/DataFormats/HGCalReco/interface/TICLGraph.h b/DataFormats/HGCalReco/interface/TICLGraph.h new file mode 100644 index 0000000000000..6150199bf1339 --- /dev/null +++ b/DataFormats/HGCalReco/interface/TICLGraph.h @@ -0,0 +1,40 @@ +#ifndef DataFormats_HGCalReco_TICLGraph_h +#define DataFormats_HGCalReco_TICLGraph_h + +#include "DataFormats/HGCalReco/interface/Trackster.h" +#include "DataFormats/TrackReco/interface/Track.h" + +class Node { +public: + Node() = default; + Node(unsigned index, bool isTrackster = true) : index_(index), isTrackster_(isTrackster){}; + + void addInner(unsigned int trackster_id) { innerNodes_.push_back(trackster_id); } + void addOuter(unsigned int trackster_id) { outerNodes_.push_back(trackster_id); } + + const unsigned int getId() const { return index_; } + std::vector getInner() const { return innerNodes_; } + std::vector getOuter() const { return outerNodes_; } + + ~Node() = default; + +private: + unsigned index_; + bool isTrackster_; + std::vector innerNodes_; + std::vector outerNodes_; +}; + +class TICLGraph { +public: + TICLGraph() = default; + TICLGraph(std::vector &n) { nodes_ = n; }; + const std::vector &getNodes() const { return nodes_; } + const Node &getNode(int i) const { return nodes_[i]; } + ~TICLGraph() = default; + +private: + std::vector nodes_; +}; + +#endif \ No newline at end of file diff --git a/DataFormats/HGCalReco/src/classes.h b/DataFormats/HGCalReco/src/classes.h index d871bfb485a71..87f8c7bb7c8b0 100644 --- a/DataFormats/HGCalReco/src/classes.h +++ b/DataFormats/HGCalReco/src/classes.h @@ -5,3 +5,4 @@ #include "DataFormats/HGCalReco/interface/TICLSeedingRegion.h" #include "DataFormats/HGCalReco/interface/TICLCandidate.h" #include "DataFormats/Common/interface/Wrapper.h" +#include "DataFormats/HGCalReco/interface/TICLGraph.h" diff --git a/DataFormats/HGCalReco/src/classes_def.xml b/DataFormats/HGCalReco/src/classes_def.xml index 9ac5b29b4f2e2..e8eacdafcc557 100644 --- a/DataFormats/HGCalReco/src/classes_def.xml +++ b/DataFormats/HGCalReco/src/classes_def.xml @@ -54,4 +54,16 @@ + + + + + + + + + + + + diff --git a/RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py b/RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py index 9c5ea39db66ed..8fc7241af03ac 100644 --- a/RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py +++ b/RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py @@ -19,7 +19,8 @@ 'keep *_ticlTrackstersHFNoseMIP_*_*', 'keep *_ticlTrackstersHFNoseHAD_*_*', 'keep *_ticlTrackstersHFNoseMerge_*_*',] + - ['keep *_pfTICL_*_*'] + ['keep *_pfTICL_*_*'] + + ['keep *_ticlGraph_*_*'] ) ) TICL_RECO.outputCommands.extend(TICL_AOD.outputCommands) diff --git a/RecoHGCal/TICL/plugins/LinkingAlgoByGNN.cc b/RecoHGCal/TICL/plugins/LinkingAlgoByGNN.cc index 568d192e21ada..e244aff93bc64 100644 --- a/RecoHGCal/TICL/plugins/LinkingAlgoByGNN.cc +++ b/RecoHGCal/TICL/plugins/LinkingAlgoByGNN.cc @@ -40,14 +40,14 @@ void Graph::addEdge(int v, int w) { void Graph::DFSUtil(int v) { // Mark the current node as visited and print it visited[v] = true; - std::cout << v << " "; - connected_components.back().push_back(v); + std::cout << v << std::endl; // Recur for all the vertices adjacent to this vertex list::iterator i; for (auto i = adj[v].begin(); i != adj[v].end(); ++i) if (!visited[*i]) { - connected_components.emplace_back(*i); + connected_components.back().push_back(*i); + std::cout << "Pushed back " << *i << std::endl; DFSUtil(*i); } } @@ -59,7 +59,9 @@ void Graph::DFS() { // traversal starting from all vertices one by one for (auto i : adj) if (visited[i.first] == false) { - //connected_components.emplace_back(i.first); + std::cout << "Emplaced back: " << i.first << std::endl; + connected_components.emplace_back(1, i.first); // {i.first} + std::cout << "Starting DFS from node: " << i.first << std::endl; DFSUtil(i.first); } } @@ -174,33 +176,36 @@ void LinkingAlgoByGNN::linkTracksters(const edm::Handle } input_shapes.push_back({1, N, shapeFeatures}); - input_shapes.push_back({1, 2, 3 * N}); - data.emplace_back(features); // Creating Edges: uncomment when have a Graph as an input - /* - std::vector edges_src; - std::vector edges_dst; - for (int i = 0; i < N; i++){ - for (auto & i_neighbour : graph.node_linked_inners[i]){ - // Create an edge between the tracksters - edges_src.push_back(i_neighbour); - edges_dst.push_back(i); - } - } - */ + + //std::vector edges_src; + //std::vector edges_dst; + //for (int i = 0; i < N; i++){ + // for (auto & i_neighbour : graph.node_linked_inners[i]){ + // // Create an edge between the tracksters + // edges_src.push_back(i_neighbour); + // edges_dst.push_back(i); + // } + //} // Create fully connected graph for testing std::vector edges_src; std::vector edges_dst; + for (int i = 0; i < N; i++) { - for (int j = i; i < N; j++) { - edges_src.push_back(i); - edges_dst.push_back(j); + std::cout << "i: " << i << std::endl; + for (int j = i; j < N; j++) { + std::cout << "j: " << j << std::endl; + edges_src.push_back(static_cast(i)); + edges_dst.push_back(static_cast(j)); } } + long unsigned int numEdges = edges_src.size(); + input_shapes.push_back({1, 2, static_cast(numEdges)}); + std::cout << "Num edges: " << numEdges << std::endl; data.emplace_back(edges_src); for (auto &dst : edges_dst) { @@ -208,6 +213,14 @@ void LinkingAlgoByGNN::linkTracksters(const edm::Handle } std::vector edge_predictions = cache->run(input_names, data, input_shapes)[0]; + + std::cout << "Network output shape is " << edge_predictions.size() << std::endl; + + for (long unsigned int i = 0; i < edge_predictions.size(); i++) { + std::cout << "Network output for edge " << data[1][i] << "-" << data[1][numEdges + i] + << " is: " << edge_predictions[i] << std::endl; + } + // Create a graph Graph g; const auto classification_threshold = 0.7; @@ -223,8 +236,10 @@ void LinkingAlgoByGNN::linkTracksters(const edm::Handle } } - std::cout << "Following is Depth First Traversal\n"; - std::cout << "Connected components are:\n"; + std::cout << "HERE 8" << std::endl; + + std::cout << "Following is Depth First Traversal" << std::endl; + std::cout << "Connected components are: " << std::endl; g.DFS(); int i = 0; @@ -235,18 +250,11 @@ void LinkingAlgoByGNN::linkTracksters(const edm::Handle for (auto &trackster_id : component) { std::cout << "Component " << i << ": trackster id " << trackster_id << std::endl; tracksterCandidate.addTrackster(edm::Ptr(tsH, trackster_id)); - i++; } + i++; connectedCandidates.push_back(tracksterCandidate); } - std::cout << "Network output shape is " << edge_predictions.size() << std::endl; - - for (long unsigned int i = 0; i < edge_predictions.size(); i++) { - std::cout << "Network output for edge " << data[1][i] << "-" << data[1][numEdges + i] - << " is: " << edge_predictions[i] << std::endl; - } - // The final candidates are passed to `resultLinked` resultLinked.insert(std::end(resultLinked), std::begin(connectedCandidates), std::end(connectedCandidates)); diff --git a/RecoHGCal/TICL/plugins/TICLGraphProducer.cc b/RecoHGCal/TICL/plugins/TICLGraphProducer.cc new file mode 100644 index 0000000000000..7c63af1f2c7fc --- /dev/null +++ b/RecoHGCal/TICL/plugins/TICLGraphProducer.cc @@ -0,0 +1,189 @@ +#include + +#include "FWCore/Framework/interface/stream/EDProducer.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/Utilities/interface/ESGetToken.h" +#include "FWCore/Framework/interface/ESHandle.h" +#include "FWCore/Framework/interface/Frameworkfwd.h" +#include "FWCore/Framework/interface/MakerMacros.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" + +#include "DataFormats/HGCalReco/interface/TICLGraph.h" +#include "DataFormats/HGCalReco/interface/Trackster.h" +#include "DataFormats/HGCalReco/interface/TICLLayerTile.h" + +#include "DataFormats/TrackReco/interface/Track.h" + +#include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h" +#include "TrackingTools/GeomPropagators/interface/Propagator.h" +#include "TrackingTools/Records/interface/TrackingComponentsRecord.h" + +#include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h" +#include "CommonTools/Utils/interface/StringCutObjectSelector.h" + +#include "MagneticField/Engine/interface/MagneticField.h" +#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h" + +#include "Geometry/HGCalCommonData/interface/HGCalDDDConstants.h" +#include "Geometry/Records/interface/IdealGeometryRecord.h" +#include "Geometry/CaloGeometry/interface/CaloGeometry.h" +#include "Geometry/Records/interface/CaloGeometryRecord.h" +#include "Geometry/CommonDetUnit/interface/GeomDet.h" + +using namespace ticl; + +class TICLGraphProducer : public edm::stream::EDProducer<> { +public: + explicit TICLGraphProducer(const edm::ParameterSet &ps); + ~TICLGraphProducer() override{}; + void produce(edm::Event &, const edm::EventSetup &) override; + static void fillDescriptions(edm::ConfigurationDescriptions &descriptions); + + void beginJob(); + void endJob(); + + void beginRun(edm::Run const &iEvent, edm::EventSetup const &es) override; + +private: + typedef math::XYZVector Vector; + const edm::EDGetTokenT> tracksters_clue3d_token_; + const edm::EDGetTokenT> tracks_token_; + const StringCutObjectSelector cutTk_; + const edm::ESGetToken geometry_token_; + const std::string detector_; + const std::string propName_; + const edm::ESGetToken bfield_token_; + const edm::ESGetToken propagator_token_; + + const HGCalDDDConstants *hgcons_; + hgcal::RecHitTools rhtools_; + edm::ESGetToken hdc_token_; +}; + +TICLGraphProducer::TICLGraphProducer(const edm::ParameterSet &ps) + : tracksters_clue3d_token_(consumes>(ps.getParameter("trackstersclue3d"))), + tracks_token_(consumes>(ps.getParameter("tracks"))), + cutTk_(ps.getParameter("cutTk")), + geometry_token_(esConsumes()), + detector_(ps.getParameter("detector")), + propName_(ps.getParameter("propagator")), + bfield_token_(esConsumes()), + propagator_token_( + esConsumes(edm::ESInputTag("", propName_))) { + produces(); + std::string detectorName_ = (detector_ == "HFNose") ? "HGCalHFNoseSensitive" : "HGCalEESensitive"; + hdc_token_ = + esConsumes(edm::ESInputTag("", detectorName_)); +} + +void TICLGraphProducer::beginJob() {} + +void TICLGraphProducer::endJob(){}; + +void TICLGraphProducer::beginRun(edm::Run const &iEvent, edm::EventSetup const &es) { + edm::ESHandle hdc = es.getHandle(hdc_token_); + hgcons_ = hdc.product(); + + edm::ESHandle geom = es.getHandle(geometry_token_); + rhtools_.setGeometry(*geom); + + edm::ESHandle bfield = es.getHandle(bfield_token_); + edm::ESHandle propagator = es.getHandle(propagator_token_); +}; + +void TICLGraphProducer::produce(edm::Event &evt, const edm::EventSetup &es) { + edm::Handle> trackstersclue3d_h; + evt.getByToken(tracksters_clue3d_token_, trackstersclue3d_h); + auto trackstersclue3d = *trackstersclue3d_h; + + //std::vector trackstersclue3d_sorted(trackstersclue3d); + //std::sort(trackstersclue3d_sorted.begin(), trackstersclue3d_sorted.end(), [](Trackster& t1, Trackster& t2){return t1.barycenter().z() < t2.barycenter().z();}); + + TICLLayerTile tracksterTilePos; + TICLLayerTile tracksterTileNeg; + + for (size_t id_t = 0; id_t < trackstersclue3d.size(); ++id_t) { + auto t = trackstersclue3d[id_t]; + if (t.barycenter().eta() > 0.) { + tracksterTilePos.fill(t.barycenter().eta(), t.barycenter().phi(), id_t); + } else if (t.barycenter().eta() < 0.) { + tracksterTileNeg.fill(t.barycenter().eta(), t.barycenter().phi(), id_t); + } + } + + std::vector allNodes; + + for (size_t id_t = 0; id_t < trackstersclue3d.size(); ++id_t) { + auto t = trackstersclue3d[id_t]; + + Node tNode(id_t); + + auto bary = t.barycenter(); + double del = 0.1; + + double eta_min = std::max(abs(bary.eta()) - del, (double)TileConstants::minEta); + double eta_max = std::min(abs(bary.eta()) + del, (double)TileConstants::maxEta); + + if (bary.eta() > 0.) { + std::array search_box = + tracksterTilePos.searchBoxEtaPhi(eta_min, eta_max, bary.phi() - del, bary.phi() + del); + if (search_box[2] > search_box[3]) { + search_box[3] += TileConstants::nPhiBins; + } + + for (int eta_i = search_box[0]; eta_i <= search_box[1]; ++eta_i) { + for (int phi_i = search_box[2]; phi_i <= search_box[3]; ++phi_i) { + auto &neighbours = tracksterTilePos[tracksterTilePos.globalBin(eta_i, (phi_i % TileConstants::nPhiBins))]; + for (auto n : neighbours) { + if (trackstersclue3d[n].barycenter().z() < bary.z()) { + tNode.addInner(n); + } else if (trackstersclue3d[n].barycenter().z() > bary.z()) { + tNode.addOuter(n); + } + } + } + } + } + + else if (bary.eta() < 0.) { + std::array search_box = + tracksterTileNeg.searchBoxEtaPhi(eta_min, eta_max, bary.phi() - del, bary.phi() + del); + if (search_box[2] > search_box[3]) { + search_box[3] += TileConstants::nPhiBins; + } + + for (int eta_i = search_box[0]; eta_i <= search_box[1]; ++eta_i) { + for (int phi_i = search_box[2]; phi_i <= search_box[3]; ++phi_i) { + auto &neighbours = tracksterTileNeg[tracksterTileNeg.globalBin(eta_i, (phi_i % TileConstants::nPhiBins))]; + for (auto n : neighbours) { + if (abs(trackstersclue3d[n].barycenter().z()) < abs(bary.z())) { + tNode.addInner(n); + } else if (abs(trackstersclue3d[n].barycenter().z()) > abs(bary.z())) { + tNode.addOuter(n); + } + } + } + } + } + allNodes.push_back(tNode); + } + auto resultGraph = std::make_unique(allNodes); + + evt.put(std::move(resultGraph)); +} + +void TICLGraphProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) { + edm::ParameterSetDescription desc; + + desc.add("trackstersclue3d", edm::InputTag("ticlTrackstersCLUE3DHigh")); + desc.add("tracks", edm::InputTag("generalTracks")); + desc.add("muons", edm::InputTag("muons1stStep")); + desc.add("detector", "HGCAL"); + desc.add("propagator", "PropagatorWithMaterial"); + desc.add("cutTk", + "1.48 < abs(eta) < 3.0 && pt > 1. && quality(\"highPurity\") && " + "hitPattern().numberOfLostHits(\"MISSING_OUTER_HITS\") < 5"); + descriptions.add("ticlGraphProducer", desc); +} + +DEFINE_FWK_MODULE(TICLGraphProducer); \ No newline at end of file diff --git a/RecoHGCal/TICL/plugins/TrackstersMergeProducer.cc b/RecoHGCal/TICL/plugins/TrackstersMergeProducer.cc index 02f8f8f3356d4..faea3d0dc0c6d 100644 --- a/RecoHGCal/TICL/plugins/TrackstersMergeProducer.cc +++ b/RecoHGCal/TICL/plugins/TrackstersMergeProducer.cc @@ -15,6 +15,8 @@ #include "DataFormats/HGCalReco/interface/Common.h" #include "DataFormats/HGCalReco/interface/TICLLayerTile.h" #include "DataFormats/HGCalReco/interface/Trackster.h" +#include "DataFormats/HGCalReco/interface/TICLSeedingRegion.h" +#include "DataFormats/HGCalReco/interface/TICLGraph.h" #include "DataFormats/TrackReco/interface/Track.h" #include "DataFormats/MuonReco/interface/Muon.h" #include "DataFormats/GeometrySurface/interface/BoundDisk.h" @@ -82,6 +84,8 @@ class TrackstersMergeProducer : public edm::stream::EDProducer linkingAlgo_; const edm::EDGetTokenT> tracksters_clue3d_token_; + const edm::EDGetTokenT ticlGraph_token_; + const edm::EDGetTokenT> seedingTrk_token_; const edm::EDGetTokenT> clusters_token_; const edm::EDGetTokenT>> clustersTime_token_; const edm::EDGetTokenT> tracks_token_; @@ -138,6 +142,7 @@ class TrackstersMergeProducer : public edm::stream::EDProducer>(ps.getParameter("trackstersclue3d"))), + ticlGraph_token_(consumes(ps.getParameter("ticlgraph"))), clusters_token_(consumes>(ps.getParameter("layer_clusters"))), clustersTime_token_( consumes>>(ps.getParameter("layer_clustersTime"))), @@ -326,6 +331,10 @@ void TrackstersMergeProducer::produce(edm::Event &evt, const edm::EventSetup &es std::back_inserter(outTrackster.vertex_multiplicity())); } + edm::Handle ticlGraph_h; + evt.getByToken(ticlGraph_token_, ticlGraph_h); + const auto &ticlGraph = *ticlGraph_h; + LogDebug("TrackstersMergeProducer") << std::endl; // Find duplicate LCs @@ -355,6 +364,22 @@ void TrackstersMergeProducer::produce(edm::Event &evt, const edm::EventSetup &es if (!outTrackster.vertices().empty()) { resultTrackstersMerged->push_back(outTrackster); } + // Compute timing + assignTimeToCandidates(*resultCandidates); + + // if (debug_) { + // // print info from graph + // std::cout << "From graph:" << std::endl; + // const auto nodes = ticlGraph.getNodes(); + // for (const auto &n : nodes) { + // std::cout << "Trackster : " << n.getId() << std::endl; + // std::cout << "inners : "; + // for (auto &inner : n.getInner()) std::cout << (int)inner << " "; + // std::cout << std::endl << "outers : "; + // for (auto &outer : n.getOuter()) std::cout << (int)outer << " "; + // std::cout << std::endl; + // } + // } } assignPCAtoTracksters(*resultTrackstersMerged, @@ -601,6 +626,7 @@ void TrackstersMergeProducer::fillDescriptions(edm::ConfigurationDescriptions &d desc.add("linkingPSet", linkingDesc); desc.add("trackstersclue3d", edm::InputTag("ticlTrackstersCLUE3DHigh")); + desc.add("ticlgraph", edm::InputTag("ticlGraph")); desc.add("layer_clusters", edm::InputTag("hgcalLayerClusters")); desc.add("layer_clustersTime", edm::InputTag("hgcalLayerClusters", "timeLayerCluster")); desc.add("tracks", edm::InputTag("generalTracks")); diff --git a/RecoHGCal/TICL/python/iterativeTICL_cff.py b/RecoHGCal/TICL/python/iterativeTICL_cff.py index 69c9989ca8955..a1b20e4aed579 100644 --- a/RecoHGCal/TICL/python/iterativeTICL_cff.py +++ b/RecoHGCal/TICL/python/iterativeTICL_cff.py @@ -13,12 +13,14 @@ from RecoHGCal.TICL.pfTICLProducer_cfi import pfTICLProducer as _pfTICLProducer from RecoHGCal.TICL.trackstersMergeProducer_cfi import trackstersMergeProducer as _trackstersMergeProducer from RecoHGCal.TICL.trackstersMergeProducerV3_cfi import trackstersMergeProducerV3 as _trackstersMergeProducerV3 +from RecoHGCal.TICL.ticlGraphProducer_cfi import ticlGraphProducer as _ticlGraphProducer from RecoHGCal.TICL.tracksterSelectionTf_cfi import * ticlLayerTileTask = cms.Task(ticlLayerTileProducer) ticlTrackstersMerge = _trackstersMergeProducer.clone() ticlTrackstersMergeV3 = _trackstersMergeProducerV3.clone() +ticlGraph = _ticlGraphProducer.clone() pfTICL = _pfTICLProducer.clone() ticlPFTask = cms.Task(pfTICL) @@ -42,12 +44,15 @@ ticlTracksterMergeTask = cms.Task(ticlTrackstersMerge) ticlTracksterMergeTaskV3 = cms.Task(ticlTrackstersMergeV3) +ticlGraphTask = cms.Task(ticlGraph) ticl_v3.toModify(pfTICL, ticlCandidateSrc = "ticlTrackstersMergeV3") mergeTICLTask = cms.Task(ticlLayerTileTask ,ticlIterationsTask + ,ticlGraphTask ,ticlTracksterMergeTask + ,ticlPFTask ) ticl_v3.toModify(mergeTICLTask, func=lambda x : x.add(ticlTracksterMergeTaskV3))