[EMCAL-734][EMCAL-689] implemented track matching and connection between cell and cluster

PWGJE/Core/JetUtilities.h

@@ -278,6 +278,111 @@ std::tuple<std::vector<int>, std::vector<int>> MatchJetsGeometrically(
   return std::make_tuple(baseToTagMap, tagToBaseMap);
 }
 
+/**
+ * Match clusters and tracks.
+ *
+ * Match cluster with tracks, where maxNumberMatches are considered in dR=maxMatchingDistance.
+ * If no unique match was found for a jet, an index of -1 is stored.
+ * The same map is created for clusters matched to tracks e.g. for electron analyses.
+ *
+ * @param clusterPhi cluster collection phi.
+ * @param clusterEta cluster collection eta.
+ * @param trackPhi track collection phi.
+ * @param trackEta track collection eta.
+ * @param maxMatchingDistance Maximum matching distance.
+ * @param maxNumberMatches Maximum number of matches (e.g. 5 closest).
+ *
+ * @returns (cluster to track index map, track to cluster index map)
+ */
+template <typename T>
+std::tuple<std::vector<std::vector<int>>, std::vector<std::vector<int>>> MatchClustersAndTracks(
+  std::vector<T>& clusterPhi,
+  std::vector<T>& clusterEta,
+  std::vector<T>& trackPhi,
+  std::vector<T>& trackEta,
+  double maxMatchingDistance,
+  int maxNumberMatches)
+{
+  // test
+  // Validation
+  const std::size_t nClusters = clusterEta.size();
+  const std::size_t nTracks = trackEta.size();
+  if (!(nClusters && nTracks)) {
+    // There are no jets, so nothing to be done.
+    return std::make_tuple(std::vector<std::vector<int>>(nClusters, std::vector<int>(maxNumberMatches, -1)), std::vector<std::vector<int>>(nTracks, std::vector<int>(maxNumberMatches, -1)));
+  }
+  // Input sizes must match
+  if (clusterPhi.size() != clusterEta.size()) {
+    throw std::invalid_argument("cluster collection eta and phi sizes don't match. Check the inputs.");
+  }
+  if (trackPhi.size() != trackEta.size()) {
+    throw std::invalid_argument("track collection eta and phi sizes don't match. Check the inputs.");
+  }
+
+  // for (std::size_t iTrack = 0; iTrack < nTracks; iTrack++) {
+  //   if (trackEta[iTrack] == 0)
+  //     LOG(warning) << "Track eta is 0!";
+  // }
+
+  // Build the KD-trees using vectors
+  // We build two trees:
+  // treeBase, which contains the base collection.
+  // treeTag, which contains the tag collection.
+  // The trees are built to match in two dimensions (eta, phi)
+  TKDTree<int, T> treeCluster(clusterEta.size(), 2, 1), treeTrack(trackEta.size(), 2, 1);
+  // By utilizing SetData, we can avoid having to copy the data again.
+  treeCluster.SetData(0, clusterEta.data());
+  treeCluster.SetData(1, clusterPhi.data());
+  treeCluster.Build();
+  treeTrack.SetData(0, trackEta.data());
+  treeTrack.SetData(1, trackPhi.data());
+  treeTrack.Build();
+
+  // Storage for the cluster matching indices.
+  std::vector<std::vector<int>> matchIndexTrack(nClusters, std::vector<int>(maxNumberMatches, -1));
+  std::vector<std::vector<int>> matchIndexCluster(nTracks, std::vector<int>(maxNumberMatches, -1));
+
+  // Find the track closest to each cluster.
+  for (std::size_t iCluster = 0; iCluster < nClusters; iCluster++) {
+    T point[2] = {clusterEta[iCluster], clusterPhi[iCluster]};
+    int index[50];  // size 50 for safety
+    T distance[50]; // size 50 for safery
+    std::fill_n(index, 50, -1);
+    std::fill_n(distance, 50, std::numeric_limits<T>::max());
+    treeTrack.FindNearestNeighbors(point, maxNumberMatches, index, distance);
+    // test whether indices are matching:
+    matchIndexTrack[iCluster] = std::vector<int>(maxNumberMatches);
+    for (int m = 0; m < maxNumberMatches; m++) {
+      if (index[m] >= 0 && distance[m] < maxMatchingDistance) {
+        matchIndexTrack[iCluster][m] = index[m];
+      } else {
+        // no match or no more matches found, fill -1
+        matchIndexTrack[iCluster][m] = -1;
+      }
+    }
+  }
+
+  // Find the base jet closest to each tag jet
+  for (std::size_t iTrack = 0; iTrack < nTracks; iTrack++) {
+    T point[2] = {trackEta[iTrack], trackPhi[iTrack]};
+    int index[50];  // size 50 for safety
+    T distance[50]; // size 50 for safery
+    std::fill_n(index, 50, -1);
+    std::fill_n(distance, 50, std::numeric_limits<T>::max());
+    treeCluster.FindNearestNeighbors(point, maxNumberMatches, index, distance);
+    matchIndexCluster[iTrack] = std::vector<int>(maxNumberMatches);
+    // loop over maxNumberMatches closest matches
+    for (int m = 0; m < maxNumberMatches; m++) {
+      if (index[m] >= 0 && distance[m] < maxMatchingDistance) {
+        matchIndexCluster[iTrack][m] = index[m];
+      } else {
+        // no match jet or no more matches found, fill -1
+        matchIndexCluster[iTrack][m] = -1;
+      }
+    }
+  }
+  return std::make_tuple(matchIndexTrack, matchIndexCluster);
+}
 }; // namespace JetUtilities
 
 #endif

PWGJE/DataModel/EMCALClusterDefinition.h

@@ -39,12 +39,13 @@ struct EMCALClusterDefinition {
   double minCellEnergy = 0.05;     // minimum cell energy (GeV)
   double timeMin = -10000;         // minimum time (ns)
   double timeMax = 10000;          // maximum time (ns)
-  double gradientCut = 0.03;       // gradient cut
+  bool doGradientCut = true;       // apply gradient cut if true
+  double gradientCut = -1;         // gradient cut
 
   // default constructor
   EMCALClusterDefinition() = default;
   // constructor
-  EMCALClusterDefinition(ClusterAlgorithm_t pAlgorithm, int pStorageID, int pSelectedCellType, std::string pName, double pSeedEnergy, double pMinCellEnergy, double pTimeMin, double pTimeMax, double pGradientCut)
+  EMCALClusterDefinition(ClusterAlgorithm_t pAlgorithm, int pStorageID, int pSelectedCellType, std::string pName, double pSeedEnergy, double pMinCellEnergy, double pTimeMin, double pTimeMax, bool pDoGradientCut, double pGradientCut)
   {
     algorithm = pAlgorithm;
     storageID = pStorageID;
@@ -54,13 +55,14 @@ struct EMCALClusterDefinition {
     minCellEnergy = pMinCellEnergy;
     timeMin = pTimeMin;
     timeMax = pTimeMax;
+    doGradientCut = pDoGradientCut;
     gradientCut = pGradientCut;
   }
 
   // implement comparison operators for int std::string and ClusterAlgorithm_t
   bool operator==(const EMCALClusterDefinition& rhs) const
   {
-    return (algorithm == rhs.algorithm && storageID == rhs.storageID && name == rhs.name && seedEnergy == rhs.seedEnergy && minCellEnergy == rhs.minCellEnergy && timeMin == rhs.timeMin && timeMax == rhs.timeMax && gradientCut == rhs.gradientCut);
+    return (algorithm == rhs.algorithm && storageID == rhs.storageID && name == rhs.name && seedEnergy == rhs.seedEnergy && minCellEnergy == rhs.minCellEnergy && timeMin == rhs.timeMin && timeMax == rhs.timeMax && gradientCut == rhs.gradientCut && doGradientCut == rhs.doGradientCut);
   }
   bool operator!=(const EMCALClusterDefinition& rhs) const
   {

PWGJE/DataModel/EMCALClusters.h

@@ -27,12 +27,12 @@ namespace emcalcluster
 // define global cluster definitions
 // the V1 algorithm is not yet implemented, but the V3 algorithm is
 // New definitions should be added here!
-const EMCALClusterDefinition kV1Default(ClusterAlgorithm_t::kV1, 0, 1, "kV1Default", 0.5, 0.1, -10000, 10000, 0.03);
-const EMCALClusterDefinition kV1Variation1(ClusterAlgorithm_t::kV3, 1, 1, "kV1Variation1", 0.3, 0.1, -10000, 10000, 0.03);
-const EMCALClusterDefinition kV1Variation2(ClusterAlgorithm_t::kV3, 2, 1, "kV1Variation2", 0.2, 0.1, -10000, 10000, 0.03);
-const EMCALClusterDefinition kV3Default(ClusterAlgorithm_t::kV3, 10, 1, "kV3Default", 0.5, 0.1, -10000, 10000, 0.03);
-const EMCALClusterDefinition kV3Variation1(ClusterAlgorithm_t::kV3, 11, 1, "kV3Variation1", 0.3, 0.1, -10000, 10000, 0.03);
-const EMCALClusterDefinition kV3Variation2(ClusterAlgorithm_t::kV3, 12, 1, "kV3Variation2", 0.2, 0.1, -10000, 10000, 0.03);
+const EMCALClusterDefinition kV1Default(ClusterAlgorithm_t::kV1, 0, 1, "kV1Default", 0.5, 0.1, -10000, 10000, true, 0.03);       //dummy
+const EMCALClusterDefinition kV1Variation1(ClusterAlgorithm_t::kV3, 1, 1, "kV1Variation1", 0.3, 0.1, -10000, 10000, true, 0.03); //dummy
+const EMCALClusterDefinition kV1Variation2(ClusterAlgorithm_t::kV3, 2, 1, "kV1Variation2", 0.2, 0.1, -10000, 10000, true, 0.03); //dummy
+const EMCALClusterDefinition kV3Default(ClusterAlgorithm_t::kV3, 10, 1, "kV3Default", 0.5, 0.1, -10000, 10000, true, 0.03);
+const EMCALClusterDefinition kV3Variation1(ClusterAlgorithm_t::kV3, 11, 1, "kV3Variation1", 0.5, 0.1, -10000, 10000, true, 0.);
+const EMCALClusterDefinition kV3Variation2(ClusterAlgorithm_t::kV3, 12, 1, "kV3Variation2", 0.5, 0.1, -10000, 10000, false, 0.);
 
 /// \brief function returns EMCALClusterDefinition for the given name
 /// \param name name of the cluster definition
@@ -89,6 +89,27 @@ DECLARE_SOA_TABLE(EMCALAmbiguousClusters, "AOD", "EMCALAMBCLUS", //!
 using EMCALCluster = EMCALClusters::iterator;
 using EMCALAmbiguousCluster = EMCALAmbiguousClusters::iterator;
 
-} // namespace o2::aod
+namespace emcalclustercell
+{
+// declare index column pointing to cluster table
+DECLARE_SOA_INDEX_COLUMN(EMCALCluster, emcalcluster); //! linked to EMCalClusters table
+DECLARE_SOA_INDEX_COLUMN(Calo, calo);                 //! linked to calo cells
 
+// declare index column pointing to ambiguous cluster table
+DECLARE_SOA_INDEX_COLUMN(EMCALAmbiguousCluster, emcalambiguouscluster); //! linked to EMCalAmbiguousClusters table
+} // namespace emcalclustercell
+DECLARE_SOA_TABLE(EMCALClusterCells, "AOD", "EMCCLUSCELLS",                                               //!
+                  o2::soa::Index<>, emcalclustercell::EMCALClusterId, emcalclustercell::CaloId);          //!
+DECLARE_SOA_TABLE(EMCALAmbiguousClusterCells, "AOD", "EMCAMBBCLUSCLS",                                    //!
+                  o2::soa::Index<>, emcalclustercell::EMCALAmbiguousClusterId, emcalclustercell::CaloId); //!
+using EMCALClusterCell = EMCALClusterCells::iterator;
+using EMCALAmbiguousClusterCell = EMCALAmbiguousClusterCells::iterator;
+namespace emcalmatchedtrack
+{
+DECLARE_SOA_INDEX_COLUMN(Track, track); //! linked to Track table only for tracks that were matched
+} // namespace emcalmatchedtrack
+DECLARE_SOA_TABLE(EMCALMatchedTracks, "AOD", "EMCMATCHTRACKS",                                     //!
+                  o2::soa::Index<>, emcalclustercell::EMCALClusterId, emcalmatchedtrack::TrackId); //!
+using EMCALMatchedTrack = EMCALMatchedTracks::iterator;
+} // namespace o2::aod
 #endif

PWGJE/TableProducer/emcalCorrectionTask.cxx

@@ -32,6 +32,7 @@
 #include "EMCALBase/Geometry.h"
 #include "EMCALBase/ClusterFactory.h"
 #include "EMCALReconstruction/Clusterizer.h"
+#include "PWGJE/Core/JetUtilities.h"
 #include "TVector2.h"
 
 using namespace o2;
@@ -41,11 +42,17 @@ using namespace o2::framework::expressions;
 struct EmcalCorrectionTask {
   Produces<o2::aod::EMCALClusters> clusters;
   Produces<o2::aod::EMCALAmbiguousClusters> clustersAmbiguous;
+  Produces<o2::aod::EMCALClusterCells> clustercells; // cells belonging to given cluster
+  Produces<o2::aod::EMCALAmbiguousClusterCells> clustercellsambiguous;
+  Produces<o2::aod::EMCALMatchedTracks> matchedTracks;
 
+  Preslice<aod::Tracks> perCollision = aod::track::collisionId;
   // Options for the clusterization
   // 1 corresponds to EMCAL cells based on the Run2 definition.
   Configurable<int> selectedCellType{"selectedCellType", 1, "EMCAL Cell type"};
   Configurable<std::string> clusterDefinitions{"clusterDefinition", "kV3Default", "cluster definition to be selected, e.g. V3Default. Multiple definitions can be specified separated by comma"};
+  Configurable<float> maxMatchingDistance{"maxMatchingDistance", 0.4f, "Max matching distance track-cluster"};
+
   // CDB service (for geometry)
   Service<o2::ccdb::BasicCCDBManager> mCcdbManager;
 
@@ -56,6 +63,8 @@ struct EmcalCorrectionTask {
   std::vector<std::unique_ptr<o2::emcal::ClusterFactory<o2::emcal::Cell>>> mClusterFactories;
   // Cells and clusters
   std::vector<o2::emcal::Cell> mEmcalCells;
+  // map of cellId (local in BC) to global cell index in cell table in AO2D
+  std::map<int, int64_t> mCellIdToCellGlobalIndex;
   std::vector<o2::emcal::AnalysisCluster> mAnalysisClusters;
 
   std::vector<o2::aod::EMCALClusterDefinition> mClusterDefinitions;
@@ -100,14 +109,15 @@ struct EmcalCorrectionTask {
       }
     }
     for (auto& clusterDefinition : mClusterDefinitions) {
-      mClusterizers.emplace_back(std::make_unique<o2::emcal::Clusterizer<o2::emcal::Cell>>(1, clusterDefinition.timeMin, clusterDefinition.timeMax, clusterDefinition.gradientCut, true, clusterDefinition.seedEnergy, clusterDefinition.minCellEnergy));
+      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));
       mClusterFactories.emplace_back(std::make_unique<o2::emcal::ClusterFactory<o2::emcal::Cell>>());
       LOG(info) << "Cluster definition initialized: " << clusterDefinition.toString();
       LOG(info) << "timeMin: " << clusterDefinition.timeMin;
       LOG(info) << "timeMax: " << clusterDefinition.timeMax;
       LOG(info) << "gradientCut: " << clusterDefinition.gradientCut;
       LOG(info) << "seedEnergy: " << clusterDefinition.seedEnergy;
       LOG(info) << "minCellEnergy: " << clusterDefinition.minCellEnergy;
+      LOG(info) << "storageID" << clusterDefinition.storageID;
     }
     for (auto& clusterizer : mClusterizers) {
       clusterizer->setGeometry(geometry);
@@ -134,15 +144,17 @@ struct EmcalCorrectionTask {
   // void process(aod::BCs const& bcs, aod::Collision const& collision, aod::Calos const& cells)
 
   //  Appears to need the BC to be accessed to be available in the collision table...
-  void process(aod::BC const& bc, aod::Collisions const& collisions, aod::Calos const& cells)
+  void process(aod::BC const& bc, aod::Collisions const& collisions, aod::Tracks const& tracks, aod::Calos const& cells)
   {
     LOG(debug) << "Starting process.";
     // Convert aod::Calo to o2::emcal::Cell which can be used with the clusterizer.
     // In particular, we need to filter only EMCAL cells.
     mEmcalCells.clear();
+    mCellIdToCellGlobalIndex.clear();
+    int c = 0;
     for (auto& cell : cells) {
       if (cell.caloType() != selectedCellType) {
-        // LOG(debug) << "Rejected";
+        LOG(debug) << "Rejected";
         continue;
       }
       // LOG(debug) << "Cell E: " << cell.getEnergy();
@@ -152,7 +164,11 @@ struct EmcalCorrectionTask {
         cell.amplitude(),
         cell.time(),
         o2::emcal::intToChannelType(cell.cellType())));
+      mCellIdToCellGlobalIndex.insert(std::make_pair(c, cell.globalIndex()));
+      LOG(debug) << "Creating map " << c << " -> " << cell.globalIndex();
+      c++;
     }
+    LOG(debug) << "Number of cells (CF): " << mEmcalCells.size();
 
     // Cell QA
     // For convenience, use the clusterizer stored geometry to get the eta-phi
@@ -197,6 +213,7 @@ struct EmcalCorrectionTask {
       for (int icl = 0; icl < mClusterFactories.at(i)->getNumberOfClusters(); icl++) {
         auto analysisCluster = mClusterFactories.at(i)->buildCluster(icl);
         mAnalysisClusters.emplace_back(analysisCluster);
+        LOG(debug) << "Cluster " << icl << ": E: " << analysisCluster.E() << ", NCells " << analysisCluster.getNCells();
       }
       LOG(debug) << "Converted to analysis clusters.";
 
@@ -212,8 +229,40 @@ struct EmcalCorrectionTask {
           vz = col.posZ();
           hasCollision = true;
 
+          // store positions of all tracks of collision
+          auto groupedTracks = tracks.sliceBy(perCollision, col.globalIndex());
+
+          std::vector<double> trackPhi;
+          std::vector<double> trackEta;
+          std::vector<int64_t> trackGlobalIndex;
+          for (auto& track : groupedTracks) {
+            // TODO this actually needs to use the eta phi
+            // of track propagated to EMC surface! Will be provided centrally according to Ruben
+            // TODO only consider tracks in current emcal/dcal acceptanc
+            trackPhi.emplace_back(TVector2::Phi_0_2pi(track.phi()));
+            trackEta.emplace_back(track.eta());
+            trackGlobalIndex.emplace_back(track.globalIndex());
+          }
+
+          std::vector<double> clusterPhi;
+          std::vector<double> clusterEta;
+
+          // TODO one loop that could in principle be combined with the other loop to improve performance
+          for (const auto& cluster : mAnalysisClusters) {
+            // Determine the cluster eta, phi, correcting for the vertex position.
+            auto pos = cluster.getGlobalPosition();
+            pos = pos - math_utils::Point3D<float>{vx, vy, vz};
+            // Normalize the vector and rescale by energy.
+            pos *= (cluster.E() / std::sqrt(pos.Mag2()));
+            clusterPhi.emplace_back(TVector2::Phi_0_2pi(pos.Phi()));
+            clusterEta.emplace_back(pos.Eta());
+          }
+          auto&& [clusterToTrackIndexMap, trackToClusterIndexMap] = JetUtilities::MatchClustersAndTracks(clusterPhi, clusterEta, trackPhi, trackEta, maxMatchingDistance, 5);
           // we found a collision, put the clusters into the none ambiguous table
           clusters.reserve(mAnalysisClusters.size());
+          int cellindex = -1;
+
+          unsigned int k = 0;
           for (const auto& cluster : mAnalysisClusters) {
 
             // Determine the cluster eta, phi, correcting for the vertex position.
@@ -228,9 +277,23 @@ struct EmcalCorrectionTask {
                      cluster.getM02(), cluster.getM20(), cluster.getNCells(), cluster.getClusterTime(),
                      cluster.getIsExotic(), cluster.getDistanceToBadChannel(), cluster.getNExMax(), static_cast<int>(mClusterDefinitions.at(i)));
 
+            clustercells.reserve(cluster.getNCells());
+            // loop over cells in cluster and save to table
+            for (int ncell = 0; ncell < cluster.getNCells(); ncell++) {
+              cellindex = cluster.getCellIndex(ncell);
+              LOG(debug) << "trying to find cell index " << cellindex << " in map";
+              clustercells(clusters.lastIndex(), mCellIdToCellGlobalIndex.at(cellindex));
+            }
             // fill histograms
             hClusterE->Fill(cluster.E());
             hClusterEtaPhi->Fill(pos.Eta(), TVector2::Phi_0_2pi(pos.Phi()));
+            for (unsigned int iTrack = 0; iTrack < clusterToTrackIndexMap[k].size(); iTrack++) {
+              if (clusterToTrackIndexMap[k][iTrack] >= 0) {
+                LOG(debug) << "Found track " << trackGlobalIndex[clusterToTrackIndexMap[k][iTrack]] << " in cluster " << cluster.getID();
+                matchedTracks(clusters.lastIndex(), trackGlobalIndex[clusterToTrackIndexMap[k][iTrack]]);
+              }
+            }
+            k++;
           } // end of cluster loop
         }   // end of collision loop
       }
@@ -241,6 +304,7 @@ struct EmcalCorrectionTask {
       // Store the clusters in the table where a mathcing collision could
       // be identified.
       if (!hasCollision) { // ambiguous
+        int cellindex = -1;
         clustersAmbiguous.reserve(mAnalysisClusters.size());
         for (const auto& cluster : mAnalysisClusters) {
           auto pos = cluster.getGlobalPosition();
@@ -254,6 +318,11 @@ struct EmcalCorrectionTask {
           clustersAmbiguous(bc, cluster.getID(), cluster.E(), cluster.getCoreEnergy(), pos.Eta(), TVector2::Phi_0_2pi(pos.Phi()),
                             cluster.getM02(), cluster.getM20(), cluster.getNCells(), cluster.getClusterTime(),
                             cluster.getIsExotic(), cluster.getDistanceToBadChannel(), cluster.getNExMax(), static_cast<int>(mClusterDefinitions.at(i)));
+          clustercellsambiguous.reserve(cluster.getNCells());
+          for (int ncell = 0; ncell < cluster.getNCells(); ncell++) {
+            cellindex = cluster.getCellIndex(ncell);
+            clustercellsambiguous(clustersAmbiguous.lastIndex(), mCellIdToCellGlobalIndex.at(cellindex));
+          }
         }
       }
       LOG(debug) << "Cluster loop done for clusterizer " << i;