3232#include " EMCALBase/Geometry.h"
3333#include " EMCALBase/ClusterFactory.h"
3434#include " EMCALReconstruction/Clusterizer.h"
35+ #include " PWGJE/Core/JetUtilities.h"
3536#include " TVector2.h"
3637
3738using namespace o2 ;
@@ -41,11 +42,17 @@ using namespace o2::framework::expressions;
4142struct EmcalCorrectionTask {
4243 Produces<o2::aod::EMCALClusters> clusters;
4344 Produces<o2::aod::EMCALAmbiguousClusters> clustersAmbiguous;
45+ Produces<o2::aod::EMCALClusterCells> clustercells; // cells belonging to given cluster
46+ Produces<o2::aod::EMCALAmbiguousClusterCells> clustercellsambiguous;
47+ Produces<o2::aod::EMCALMatchedTracks> matchedTracks;
4448
49+ Preslice<aod::Tracks> perCollision = aod::track::collisionId;
4550 // Options for the clusterization
4651 // 1 corresponds to EMCAL cells based on the Run2 definition.
4752 Configurable<int > selectedCellType{" selectedCellType" , 1 , " EMCAL Cell type" };
4853 Configurable<std::string> clusterDefinitions{" clusterDefinition" , " kV3Default" , " cluster definition to be selected, e.g. V3Default. Multiple definitions can be specified separated by comma" };
54+ Configurable<float > maxMatchingDistance{" maxMatchingDistance" , 0 .4f , " Max matching distance track-cluster" };
55+
4956 // CDB service (for geometry)
5057 Service<o2::ccdb::BasicCCDBManager> mCcdbManager ;
5158
@@ -56,6 +63,8 @@ struct EmcalCorrectionTask {
5663 std::vector<std::unique_ptr<o2::emcal::ClusterFactory<o2::emcal::Cell>>> mClusterFactories ;
5764 // Cells and clusters
5865 std::vector<o2::emcal::Cell> mEmcalCells ;
66+ // map of cellId (local in BC) to global cell index in cell table in AO2D
67+ std::map<int , int64_t > mCellIdToCellGlobalIndex ;
5968 std::vector<o2::emcal::AnalysisCluster> mAnalysisClusters ;
6069
6170 std::vector<o2::aod::EMCALClusterDefinition> mClusterDefinitions ;
@@ -100,14 +109,15 @@ struct EmcalCorrectionTask {
100109 }
101110 }
102111 for (auto & clusterDefinition : mClusterDefinitions ) {
103- mClusterizers .emplace_back (std::make_unique<o2::emcal::Clusterizer<o2::emcal::Cell>>(1 , clusterDefinition.timeMin , clusterDefinition.timeMax , clusterDefinition.gradientCut , true , clusterDefinition.seedEnergy , clusterDefinition.minCellEnergy ));
112+ mClusterizers .emplace_back (std::make_unique<o2::emcal::Clusterizer<o2::emcal::Cell>>(1E9 , clusterDefinition.timeMin , clusterDefinition.timeMax , clusterDefinition.gradientCut , clusterDefinition. doGradientCut , clusterDefinition.seedEnergy , clusterDefinition.minCellEnergy ));
104113 mClusterFactories .emplace_back (std::make_unique<o2::emcal::ClusterFactory<o2::emcal::Cell>>());
105114 LOG (info) << " Cluster definition initialized: " << clusterDefinition.toString ();
106115 LOG (info) << " timeMin: " << clusterDefinition.timeMin ;
107116 LOG (info) << " timeMax: " << clusterDefinition.timeMax ;
108117 LOG (info) << " gradientCut: " << clusterDefinition.gradientCut ;
109118 LOG (info) << " seedEnergy: " << clusterDefinition.seedEnergy ;
110119 LOG (info) << " minCellEnergy: " << clusterDefinition.minCellEnergy ;
120+ LOG (info) << " storageID" << clusterDefinition.storageID ;
111121 }
112122 for (auto & clusterizer : mClusterizers ) {
113123 clusterizer->setGeometry (geometry);
@@ -134,15 +144,17 @@ struct EmcalCorrectionTask {
134144 // void process(aod::BCs const& bcs, aod::Collision const& collision, aod::Calos const& cells)
135145
136146 // Appears to need the BC to be accessed to be available in the collision table...
137- void process (aod::BC const & bc, aod::Collisions const & collisions, aod::Calos const & cells)
147+ void process (aod::BC const & bc, aod::Collisions const & collisions, aod::Tracks const & tracks, aod:: Calos const & cells)
138148 {
139149 LOG (debug) << " Starting process." ;
140150 // Convert aod::Calo to o2::emcal::Cell which can be used with the clusterizer.
141151 // In particular, we need to filter only EMCAL cells.
142152 mEmcalCells .clear ();
153+ mCellIdToCellGlobalIndex .clear ();
154+ int c = 0 ;
143155 for (auto & cell : cells) {
144156 if (cell.caloType () != selectedCellType) {
145- // LOG(debug) << "Rejected";
157+ LOG (debug) << " Rejected" ;
146158 continue ;
147159 }
148160 // LOG(debug) << "Cell E: " << cell.getEnergy();
@@ -152,7 +164,11 @@ struct EmcalCorrectionTask {
152164 cell.amplitude (),
153165 cell.time (),
154166 o2::emcal::intToChannelType (cell.cellType ())));
167+ mCellIdToCellGlobalIndex .insert (std::make_pair (c, cell.globalIndex ()));
168+ LOG (debug) << " Creating map " << c << " -> " << cell.globalIndex ();
169+ c++;
155170 }
171+ LOG (debug) << " Number of cells (CF): " << mEmcalCells .size ();
156172
157173 // Cell QA
158174 // For convenience, use the clusterizer stored geometry to get the eta-phi
@@ -197,6 +213,7 @@ struct EmcalCorrectionTask {
197213 for (int icl = 0 ; icl < mClusterFactories .at (i)->getNumberOfClusters (); icl++) {
198214 auto analysisCluster = mClusterFactories .at (i)->buildCluster (icl);
199215 mAnalysisClusters .emplace_back (analysisCluster);
216+ LOG (debug) << " Cluster " << icl << " : E: " << analysisCluster.E () << " , NCells " << analysisCluster.getNCells ();
200217 }
201218 LOG (debug) << " Converted to analysis clusters." ;
202219
@@ -212,8 +229,40 @@ struct EmcalCorrectionTask {
212229 vz = col.posZ ();
213230 hasCollision = true ;
214231
232+ // store positions of all tracks of collision
233+ auto groupedTracks = tracks.sliceBy (perCollision, col.globalIndex ());
234+
235+ std::vector<double > trackPhi;
236+ std::vector<double > trackEta;
237+ std::vector<int64_t > trackGlobalIndex;
238+ for (auto & track : groupedTracks) {
239+ // TODO this actually needs to use the eta phi
240+ // of track propagated to EMC surface! Will be provided centrally according to Ruben
241+ // TODO only consider tracks in current emcal/dcal acceptanc
242+ trackPhi.emplace_back (TVector2::Phi_0_2pi (track.phi ()));
243+ trackEta.emplace_back (track.eta ());
244+ trackGlobalIndex.emplace_back (track.globalIndex ());
245+ }
246+
247+ std::vector<double > clusterPhi;
248+ std::vector<double > clusterEta;
249+
250+ // TODO one loop that could in principle be combined with the other loop to improve performance
251+ for (const auto & cluster : mAnalysisClusters ) {
252+ // Determine the cluster eta, phi, correcting for the vertex position.
253+ auto pos = cluster.getGlobalPosition ();
254+ pos = pos - math_utils::Point3D<float >{vx, vy, vz};
255+ // Normalize the vector and rescale by energy.
256+ pos *= (cluster.E () / std::sqrt (pos.Mag2 ()));
257+ clusterPhi.emplace_back (TVector2::Phi_0_2pi (pos.Phi ()));
258+ clusterEta.emplace_back (pos.Eta ());
259+ }
260+ auto && [clusterToTrackIndexMap, trackToClusterIndexMap] = JetUtilities::MatchClustersAndTracks (clusterPhi, clusterEta, trackPhi, trackEta, maxMatchingDistance, 5 );
215261 // we found a collision, put the clusters into the none ambiguous table
216262 clusters.reserve (mAnalysisClusters .size ());
263+ int cellindex = -1 ;
264+
265+ unsigned int k = 0 ;
217266 for (const auto & cluster : mAnalysisClusters ) {
218267
219268 // Determine the cluster eta, phi, correcting for the vertex position.
@@ -228,9 +277,23 @@ struct EmcalCorrectionTask {
228277 cluster.getM02 (), cluster.getM20 (), cluster.getNCells (), cluster.getClusterTime (),
229278 cluster.getIsExotic (), cluster.getDistanceToBadChannel (), cluster.getNExMax (), static_cast <int >(mClusterDefinitions .at (i)));
230279
280+ clustercells.reserve (cluster.getNCells ());
281+ // loop over cells in cluster and save to table
282+ for (int ncell = 0 ; ncell < cluster.getNCells (); ncell++) {
283+ cellindex = cluster.getCellIndex (ncell);
284+ LOG (debug) << " trying to find cell index " << cellindex << " in map" ;
285+ clustercells (clusters.lastIndex (), mCellIdToCellGlobalIndex .at (cellindex));
286+ }
231287 // fill histograms
232288 hClusterE->Fill (cluster.E ());
233289 hClusterEtaPhi->Fill (pos.Eta (), TVector2::Phi_0_2pi (pos.Phi ()));
290+ for (unsigned int iTrack = 0 ; iTrack < clusterToTrackIndexMap[k].size (); iTrack++) {
291+ if (clusterToTrackIndexMap[k][iTrack] >= 0 ) {
292+ LOG (debug) << " Found track " << trackGlobalIndex[clusterToTrackIndexMap[k][iTrack]] << " in cluster " << cluster.getID ();
293+ matchedTracks (clusters.lastIndex (), trackGlobalIndex[clusterToTrackIndexMap[k][iTrack]]);
294+ }
295+ }
296+ k++;
234297 } // end of cluster loop
235298 } // end of collision loop
236299 }
@@ -241,6 +304,7 @@ struct EmcalCorrectionTask {
241304 // Store the clusters in the table where a mathcing collision could
242305 // be identified.
243306 if (!hasCollision) { // ambiguous
307+ int cellindex = -1 ;
244308 clustersAmbiguous.reserve (mAnalysisClusters .size ());
245309 for (const auto & cluster : mAnalysisClusters ) {
246310 auto pos = cluster.getGlobalPosition ();
@@ -254,6 +318,11 @@ struct EmcalCorrectionTask {
254318 clustersAmbiguous (bc, cluster.getID (), cluster.E (), cluster.getCoreEnergy (), pos.Eta (), TVector2::Phi_0_2pi (pos.Phi ()),
255319 cluster.getM02 (), cluster.getM20 (), cluster.getNCells (), cluster.getClusterTime (),
256320 cluster.getIsExotic (), cluster.getDistanceToBadChannel (), cluster.getNExMax (), static_cast <int >(mClusterDefinitions .at (i)));
321+ clustercellsambiguous.reserve (cluster.getNCells ());
322+ for (int ncell = 0 ; ncell < cluster.getNCells (); ncell++) {
323+ cellindex = cluster.getCellIndex (ncell);
324+ clustercellsambiguous (clustersAmbiguous.lastIndex (), mCellIdToCellGlobalIndex .at (cellindex));
325+ }
257326 }
258327 }
259328 LOG (debug) << " Cluster loop done for clusterizer " << i;
0 commit comments